package server

import (
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/store/automation"
)

func TestAutomationRunnerAcquireLongPollRechecksRevokedSession(t *testing.T) {
	automationYAML := "name: Revoked long poll flow\non:\n  manual: {}\njobs:\n  build:\n    steps:\n      - run: echo revoked-long-poll\n"
	srv := newAutomationRunnerTestServer(t, automationYAML, "9090909090909090909090909090909090909090")
	humanAuth := "Bearer " + jwt(map[string]any{"sub": "runner-revoked-long-poll-user"})
	workspaceID := "ws-runner-revoked-long-poll"
	oldRunnerAuth := registerAutomationRunnerSession(t, srv, humanAuth, workspaceID, "runner-1", "session-1", nil)

	type acquireResult struct {
		code int
		body []byte
	}
	resultCh := make(chan acquireResult, 1)
	go func() {
		acquireReq := httptest.NewRequest(http.MethodGet, "/api/v1/automations/runners/runner-1/jobs:acquire?sessionId=session-1&limit=1&waitSeconds=2", nil)
		setAutomationRunnerHeaders(acquireReq, oldRunnerAuth, workspaceID)
		acquireRec := httptest.NewRecorder()
		srv.handler.ServeHTTP(acquireRec, acquireReq)
		resultCh <- acquireResult{code: acquireRec.Code, body: acquireRec.Body.Bytes()}
	}()

	select {
	case result := <-resultCh:
		t.Fatalf("long-poll acquire returned before session revocation: status=%d body=%s", result.code, string(result.body))
	case <-time.After(250 * time.Millisecond):
	}

	newRunnerAuth := registerAutomationRunnerSession(t, srv, humanAuth, workspaceID, "runner-1", "session-2", nil)
	run := startAutomationRunnerTestRun(t, srv, humanAuth, workspaceID, ".agentteam/runner.yml")

	select {
	case result := <-resultCh:
		require.Equal(t, http.StatusUnauthorized, result.code, string(result.body))
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for revoked long-poll acquire response")
	}

	var stored automation.AutomationJob
	require.NoError(t, srv.meshxDB.Where("workspace_id = ? AND run_id = ?", workspaceID, run.ID).First(&stored).Error)
	require.Equal(t, automation.AutomationJobStatusReady, stored.Status)
	require.Empty(t, stored.LeaseOwner)

	claimed := acquireAutomationRunnerJob(t, srv, newRunnerAuth, workspaceID, "runner-1", "session-2")
	require.Equal(t, stored.ID, claimed.QueueJobID)
	require.Equal(t, "build", claimed.JobID)
}