package server

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/config"
)

func TestAutomationRunnerDispatchAPI(t *testing.T) {
	automationYAML := "name: Runner flow\non:\n  manual: {}\njobs:\n  build:\n    steps:\n      - id: build\n        run: make build\n  test:\n    needs:\n      - build\n    steps:\n      - id: test\n        run: go test ./...\n"
	commitSHA := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
	srv := newAutomationRunnerTestServer(t, automationYAML, commitSHA)
	humanAuth := "Bearer " + jwt(map[string]any{"sub": "runner-user"})
	workspaceID := "ws-runner-dispatch"
	run := startAutomationRunnerTestRun(t, srv, humanAuth, workspaceID, ".agentteam/runner.yml")

	registerReq := httptest.NewRequest(http.MethodPost, "/api/v1/automations/runners/register", bytes.NewBufferString(`{"runnerId":"runner-1","name":"Runner 1","labels":["ubuntu-latest","self-hosted"]}`))
	setAutomationRunnerHeaders(registerReq, humanAuth, workspaceID)
	registerRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(registerRec, registerReq)
	require.Equal(t, http.StatusCreated, registerRec.Code, registerRec.Body.String())
	require.Contains(t, registerRec.Body.String(), `"runnerId":"runner-1"`)
	var registration automationRunnerRegistration
	require.NoError(t, json.Unmarshal(registerRec.Body.Bytes(), &registration))
	require.NotEmpty(t, registration.RegistrationToken)

	sessionReq := httptest.NewRequest(http.MethodPost, "/api/v1/automations/runners/runner-1/sessions", bytes.NewBufferString(`{"sessionId":"session-1"}`))
	setAutomationRunnerHeaders(sessionReq, "Bearer "+registration.RegistrationToken, workspaceID)
	sessionRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(sessionRec, sessionReq)
	require.Equal(t, http.StatusCreated, sessionRec.Code, sessionRec.Body.String())
	var session automationRunnerSession
	require.NoError(t, json.Unmarshal(sessionRec.Body.Bytes(), &session))
	require.NotEmpty(t, session.SessionToken)
	require.NotEmpty(t, session.LeaseOwner)
	require.NotEqual(t, "session-1", session.LeaseOwner)
	runnerAuth := "Bearer " + session.SessionToken

	reuseRegistrationReq := httptest.NewRequest(http.MethodPost, "/api/v1/automations/runners/runner-1/sessions", bytes.NewBufferString(`{"sessionId":"session-reuse"}`))
	setAutomationRunnerHeaders(reuseRegistrationReq, "Bearer "+registration.RegistrationToken, workspaceID)
	reuseRegistrationRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(reuseRegistrationRec, reuseRegistrationReq)
	require.Equal(t, http.StatusUnauthorized, reuseRegistrationRec.Code, reuseRegistrationRec.Body.String())

	firstJob := acquireAutomationRunnerJob(t, srv, runnerAuth, workspaceID, "runner-1", "session-1")
	require.Equal(t, "build", firstJob.JobID)
	require.Equal(t, session.LeaseOwner, firstJob.LeaseOwner)
	require.Equal(t, "pingcap/agentteam", firstJob.RepositoryFullName)
	require.Equal(t, commitSHA, firstJob.CommitSHA)
	require.Contains(t, firstJob.InputJSON, "make build")
	require.JSONEq(t, `{"trigger_kind":"manual"}`, firstJob.TriggerContextJSON)

	wrongAcceptReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-1/jobs/%d/accepted", firstJob.QueueJobID), bytes.NewBufferString(`{"sessionId":"other-session"}`))
	setAutomationRunnerHeaders(wrongAcceptReq, runnerAuth, workspaceID)
	wrongAcceptRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(wrongAcceptRec, wrongAcceptReq)
	require.Equal(t, http.StatusUnauthorized, wrongAcceptRec.Code, wrongAcceptRec.Body.String())

	acceptReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-1/jobs/%d/accepted", firstJob.QueueJobID), bytes.NewBufferString(`{"sessionId":"session-1"}`))
	setAutomationRunnerHeaders(acceptReq, runnerAuth, workspaceID)
	acceptRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(acceptRec, acceptReq)
	require.Equal(t, http.StatusOK, acceptRec.Code, acceptRec.Body.String())
	require.Contains(t, acceptRec.Body.String(), `"status":"running"`)

	renewReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-1/jobs/%d/renew", firstJob.QueueJobID), bytes.NewBufferString(`{"sessionId":"session-1","leaseSeconds":600}`))
	setAutomationRunnerHeaders(renewReq, runnerAuth, workspaceID)
	renewRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(renewRec, renewReq)
	require.Equal(t, http.StatusOK, renewRec.Code, renewRec.Body.String())
	require.Contains(t, renewRec.Body.String(), `"queueJobId":`)

	logReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-1/jobs/%d/logs", firstJob.QueueJobID), bytes.NewBufferString(`{"sessionId":"session-1","stream":"stdout","content":"building","sequence":1}`))
	setAutomationRunnerHeaders(logReq, runnerAuth, workspaceID)
	logRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(logRec, logReq)
	require.Equal(t, http.StatusCreated, logRec.Code, logRec.Body.String())
	require.Contains(t, logRec.Body.String(), `"eventType":"job_log_appended"`)
	require.NotContains(t, logRec.Body.String(), "building")

	logsReq := httptest.NewRequest(http.MethodGet, "/api/v1/automations/runs/"+run.ID+"/logs?jobRunId="+firstJob.JobRunID+"&stream=stdout", nil)
	setAutomationRunnerHeaders(logsReq, humanAuth, workspaceID)
	logsRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(logsRec, logsReq)
	require.Equal(t, http.StatusOK, logsRec.Code, logsRec.Body.String())
	var logChunks []automationLogChunk
	require.NoError(t, json.Unmarshal(logsRec.Body.Bytes(), &logChunks))
	require.Len(t, logChunks, 1)
	require.Equal(t, firstJob.JobRunID, logChunks[0].JobRunID)
	require.Equal(t, "stdout", logChunks[0].Stream)
	require.EqualValues(t, 1, logChunks[0].Sequence)
	require.EqualValues(t, len("building"), logChunks[0].SizeBytes)
	require.NotContains(t, logsRec.Body.String(), "building")

	logContentReq := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/api/v1/automations/runs/%s/jobs/%s/logs/stdout/1", run.ID, firstJob.JobRunID), nil)
	setAutomationRunnerHeaders(logContentReq, humanAuth, workspaceID)
	logContentRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(logContentRec, logContentReq)
	require.Equal(t, http.StatusOK, logContentRec.Code, logContentRec.Body.String())
	var logContent automationLogChunkContent
	require.NoError(t, json.Unmarshal(logContentRec.Body.Bytes(), &logContent))
	require.Equal(t, "building", logContent.Content)

	artifactPayload := []byte("build log artifact")
	artifactBody, err := json.Marshal(map[string]any{"sessionId": "session-1", "name": "build-log", "contentBase64": base64.StdEncoding.EncodeToString(artifactPayload), "size": len(artifactPayload), "sha256": testSHA256Hex(artifactPayload), "contentType": "text/plain"})
	require.NoError(t, err)
	artifactReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-1/jobs/%d/artifacts", firstJob.QueueJobID), bytes.NewBuffer(artifactBody))
	setAutomationRunnerHeaders(artifactReq, runnerAuth, workspaceID)
	artifactRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(artifactRec, artifactReq)
	require.Equal(t, http.StatusCreated, artifactRec.Code, artifactRec.Body.String())
	require.Contains(t, artifactRec.Body.String(), `"eventType":"job_artifact_recorded"`)
	require.NotContains(t, artifactRec.Body.String(), base64.StdEncoding.EncodeToString(artifactPayload))

	artifactsReq := httptest.NewRequest(http.MethodGet, "/api/v1/automations/runs/"+run.ID+"/artifacts?jobRunId="+firstJob.JobRunID, nil)
	setAutomationRunnerHeaders(artifactsReq, humanAuth, workspaceID)
	artifactsRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(artifactsRec, artifactsReq)
	require.Equal(t, http.StatusOK, artifactsRec.Code, artifactsRec.Body.String())
	var artifacts []automationArtifact
	require.NoError(t, json.Unmarshal(artifactsRec.Body.Bytes(), &artifacts))
	require.Len(t, artifacts, 1)
	require.Equal(t, firstJob.JobRunID, artifacts[0].JobRunID)
	require.Equal(t, "build-log", artifacts[0].Name)
	require.EqualValues(t, len(artifactPayload), artifacts[0].SizeBytes)
	require.NotContains(t, artifactsRec.Body.String(), base64.StdEncoding.EncodeToString(artifactPayload))

	artifactContentReq := httptest.NewRequest(http.MethodGet, "/api/v1/automations/runs/"+run.ID+"/artifacts/"+artifacts[0].ID, nil)
	setAutomationRunnerHeaders(artifactContentReq, humanAuth, workspaceID)
	artifactContentRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(artifactContentRec, artifactContentReq)
	require.Equal(t, http.StatusOK, artifactContentRec.Code, artifactContentRec.Body.String())
	var artifactContent automationArtifactContent
	require.NoError(t, json.Unmarshal(artifactContentRec.Body.Bytes(), &artifactContent))
	require.Equal(t, artifacts[0].ID, artifactContent.ID)
	require.Equal(t, base64.StdEncoding.EncodeToString(artifactPayload), artifactContent.ContentBase64)

	crossWorkspaceLogReq := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/api/v1/automations/runs/%s/jobs/%s/logs/stdout/1", run.ID, firstJob.JobRunID), nil)
	setAutomationRunnerHeaders(crossWorkspaceLogReq, humanAuth, workspaceID+"-other")
	crossWorkspaceLogRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(crossWorkspaceLogRec, crossWorkspaceLogReq)
	require.Equal(t, http.StatusNotFound, crossWorkspaceLogRec.Code, crossWorkspaceLogRec.Body.String())

	unauthorizedArtifactReq := httptest.NewRequest(http.MethodGet, "/api/v1/automations/runs/"+run.ID+"/artifacts/"+artifacts[0].ID, nil)
	unauthorizedArtifactReq.Header.Set("X-Workspace-Id", workspaceID)
	unauthorizedArtifactRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(unauthorizedArtifactRec, unauthorizedArtifactReq)
	require.Equal(t, http.StatusUnauthorized, unauthorizedArtifactRec.Code, unauthorizedArtifactRec.Body.String())

	wrongCompleteReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-1/jobs/%d/completed", firstJob.QueueJobID), bytes.NewBufferString(`{"sessionId":"other-session","outputJson":"{}"}`))
	setAutomationRunnerHeaders(wrongCompleteReq, runnerAuth, workspaceID)
	wrongCompleteRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(wrongCompleteRec, wrongCompleteReq)
	require.Equal(t, http.StatusUnauthorized, wrongCompleteRec.Code, wrongCompleteRec.Body.String())

	completeAutomationRunnerJob(t, srv, runnerAuth, workspaceID, "runner-1", "session-1", firstJob.QueueJobID, `{"ok":true}`)
	secondJob := acquireAutomationRunnerJob(t, srv, runnerAuth, workspaceID, "runner-1", "session-1")
	require.Equal(t, "test", secondJob.JobID)
	completeAutomationRunnerJob(t, srv, runnerAuth, workspaceID, "runner-1", "session-1", secondJob.QueueJobID, `{"ok":true}`)

	getRunReq := httptest.NewRequest(http.MethodGet, "/api/v1/automations/runs/"+run.ID, nil)
	setAutomationRunnerHeaders(getRunReq, humanAuth, workspaceID)
	getRunRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(getRunRec, getRunReq)
	require.Equal(t, http.StatusOK, getRunRec.Code, getRunRec.Body.String())
	require.Contains(t, getRunRec.Body.String(), `"status":"succeeded"`)

	eventsReq := httptest.NewRequest(http.MethodGet, "/api/v1/automations/runs/"+run.ID+"/events", nil)
	setAutomationRunnerHeaders(eventsReq, humanAuth, workspaceID)
	eventsRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(eventsRec, eventsReq)
	require.Equal(t, http.StatusOK, eventsRec.Code, eventsRec.Body.String())
	require.Contains(t, eventsRec.Body.String(), "job_log_appended")
	require.Contains(t, eventsRec.Body.String(), "job_artifact_recorded")
}

func TestAutomationRunnerFailureAndCancelAPI(t *testing.T) {
	automationYAML := "name: Failure flow\non:\n  manual: {}\njobs:\n  build:\n    steps:\n      - id: build\n        run: make build\n"
	srv := newAutomationRunnerTestServer(t, automationYAML, "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb")
	humanAuth := "Bearer " + jwt(map[string]any{"sub": "runner-user"})
	workspaceID := "ws-runner-failure"
	run := startAutomationRunnerTestRun(t, srv, humanAuth, workspaceID, ".agentteam/runner.yml")
	runnerAuth := registerAutomationRunnerSession(t, srv, humanAuth, workspaceID, "runner-1", "runner-1", nil)
	job := acquireAutomationRunnerJob(t, srv, runnerAuth, workspaceID, "runner-1", "runner-1")

	failReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-1/jobs/%d/failed", job.QueueJobID), bytes.NewBufferString(`{"sessionId":"runner-1","message":"runner failed"}`))
	setAutomationRunnerHeaders(failReq, runnerAuth, workspaceID)
	failRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(failRec, failReq)
	require.Equal(t, http.StatusOK, failRec.Code, failRec.Body.String())
	require.Contains(t, failRec.Body.String(), `"status":"queued"`)

	cancelReq := httptest.NewRequest(http.MethodPost, "/api/v1/automations/runs/"+run.ID+"/cancel", bytes.NewBufferString(`{"reason":"stop after failure"}`))
	setAutomationRunnerHeaders(cancelReq, humanAuth, workspaceID)
	cancelRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(cancelRec, cancelReq)
	require.Equal(t, http.StatusOK, cancelRec.Code, cancelRec.Body.String())
	require.Contains(t, cancelRec.Body.String(), `"status":"canceled"`)

	retryReq := httptest.NewRequest(http.MethodPost, "/api/v1/automations/runs/"+run.ID+"/retry", nil)
	setAutomationRunnerHeaders(retryReq, humanAuth, workspaceID)
	retryRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(retryRec, retryReq)
	require.Equal(t, http.StatusCreated, retryRec.Code, retryRec.Body.String())
	require.Contains(t, retryRec.Body.String(), `"status":"running"`)
	require.Contains(t, retryRec.Body.String(), `"triggerKind":"manual_retry"`)
	require.Contains(t, retryRec.Body.String(), `"triggerRef":"`+run.ID+`"`)
}

func TestAutomationRunnerRetryRejectsActiveRun(t *testing.T) {
	automationYAML := "name: Active retry flow\non:\n  manual: {}\njobs:\n  build:\n    steps:\n      - run: echo active\n"
	srv := newAutomationRunnerTestServer(t, automationYAML, "1212121212121212121212121212121212121212")
	auth := "Bearer " + jwt(map[string]any{"sub": "runner-active-retry-user"})
	workspaceID := "ws-runner-active-retry"
	run := startAutomationRunnerTestRun(t, srv, auth, workspaceID, ".agentteam/runner.yml")

	retryReq := httptest.NewRequest(http.MethodPost, "/api/v1/automations/runs/"+run.ID+"/retry", nil)
	setAutomationRunnerHeaders(retryReq, auth, workspaceID)
	retryRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(retryRec, retryReq)
	require.Equal(t, http.StatusBadRequest, retryRec.Code, retryRec.Body.String())
	require.Contains(t, retryRec.Body.String(), "terminal status")
}

func TestAutomationRunnerRetryUsesOriginalCommitForMutableBranchRef(t *testing.T) {
	automationYAML := "name: Mutable retry flow\non:\n  manual: {}\njobs:\n  build:\n    steps:\n      - run: echo mutable\n"
	oldCommit := "1111111111111111111111111111111111111111"
	newCommit := "2222222222222222222222222222222222222222"
	branchCommit := oldCommit
	srv := newAutomationRunnerMutableRefTestServer(t, automationYAML, &branchCommit, oldCommit, newCommit)
	auth := "Bearer " + jwt(map[string]any{"sub": "runner-mutable-retry-user"})
	workspaceID := "ws-runner-mutable-retry"
	run := startAutomationRunnerTestRun(t, srv, auth, workspaceID, ".agentteam/runner.yml")
	require.Equal(t, oldCommit, run.CommitSHA)

	branchCommit = newCommit
	cancelReq := httptest.NewRequest(http.MethodPost, "/api/v1/automations/runs/"+run.ID+"/cancel", bytes.NewBufferString(`{"reason":"make terminal before retry"}`))
	setAutomationRunnerHeaders(cancelReq, auth, workspaceID)
	cancelRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(cancelRec, cancelReq)
	require.Equal(t, http.StatusOK, cancelRec.Code, cancelRec.Body.String())

	retryReq := httptest.NewRequest(http.MethodPost, "/api/v1/automations/runs/"+run.ID+"/retry", nil)
	setAutomationRunnerHeaders(retryReq, auth, workspaceID)
	retryRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(retryRec, retryReq)
	require.Equal(t, http.StatusCreated, retryRec.Code, retryRec.Body.String())
	var retryRun automationRunnerTestRun
	require.NoError(t, json.Unmarshal(retryRec.Body.Bytes(), &retryRun))
	require.Equal(t, oldCommit, retryRun.CommitSHA)
	require.NotEqual(t, newCommit, retryRun.CommitSHA)
	require.Equal(t, "manual_retry", retryRun.TriggerKind)
	require.Equal(t, run.ID, retryRun.TriggerRef)
}

func TestAutomationRunnerAcquireHonorsRunsOnLabels(t *testing.T) {
	automationYAML := "name: Labels flow\non:\n  manual: {}\njobs:\n  linux:\n    runs-on: ubuntu-latest\n    steps:\n      - run: echo linux\n  macos:\n    runs-on: [self-hosted, macos]\n    steps:\n      - run: echo macos\n"
	srv := newAutomationRunnerTestServer(t, automationYAML, "cccccccccccccccccccccccccccccccccccccccc")
	humanAuth := "Bearer " + jwt(map[string]any{"sub": "runner-user"})
	workspaceID := "ws-runner-labels"
	macosAuth := registerAutomationRunnerSession(t, srv, humanAuth, workspaceID, "mac-runner", "mac-session", []string{"self-hosted", "macos"})
	_ = startAutomationRunnerTestRun(t, srv, humanAuth, workspaceID, ".agentteam/runner.yml")

	linuxAuth := registerAutomationRunnerSession(t, srv, humanAuth, workspaceID, "linux-runner", "linux-session", []string{"ubuntu-latest"})
	linux := acquireAutomationRunnerJobWithQuery(t, srv, linuxAuth, workspaceID, "linux-runner", "linux-session", "")
	require.Equal(t, "linux", linux.JobID)
	require.JSONEq(t, `"ubuntu-latest"`, linux.RunsOnJSON)
	require.Equal(t, []string{"ubuntu-latest"}, linux.Labels)

	macosMissingAuth := registerAutomationRunnerSession(t, srv, humanAuth, workspaceID, "mac-missing-runner", "mac-missing-session", []string{"self-hosted"})
	macosMissingLabel := acquireAutomationRunnerJobsWithQuery(t, srv, macosMissingAuth, workspaceID, "mac-missing-runner", "mac-missing-session", "")
	require.Empty(t, macosMissingLabel)
	macosQueryLabel := acquireAutomationRunnerJobsWithQuery(t, srv, macosMissingAuth, workspaceID, "mac-missing-runner", "mac-missing-session", "labels=self-hosted,macos")
	require.Empty(t, macosQueryLabel)

	macos := acquireAutomationRunnerJobWithQuery(t, srv, macosAuth, workspaceID, "mac-runner", "mac-session", "")
	require.Equal(t, "macos", macos.JobID)
	require.JSONEq(t, `["self-hosted","macos"]`, macos.RunsOnJSON)
	require.Equal(t, []string{"self-hosted", "macos"}, macos.Labels)
}

func TestAutomationRunnerAuthLeaseAndAcquireBoundaries(t *testing.T) {
	automationYAML := "name: Boundary flow\non:\n  manual: {}\njobs:\n  build:\n    steps:\n      - run: echo boundary\n"
	srv := newAutomationRunnerTestServer(t, automationYAML, "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee")
	humanAuth := "Bearer " + jwt(map[string]any{"sub": "runner-boundary-user"})
	workspaceID := "ws-runner-boundaries"
	_ = startAutomationRunnerTestRun(t, srv, humanAuth, workspaceID, ".agentteam/runner.yml")
	runnerAuth := registerAutomationRunnerSession(t, srv, humanAuth, workspaceID, "runner-1", "session-1", nil)

	unauthorizedReq := httptest.NewRequest(http.MethodGet, "/api/v1/automations/runners/runner-1/jobs:acquire?sessionId=session-1", nil)
	unauthorizedReq.Header.Set("X-Workspace-Id", workspaceID)
	unauthorizedRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(unauthorizedRec, unauthorizedReq)
	require.Equal(t, http.StatusUnauthorized, unauthorizedRec.Code, unauthorizedRec.Body.String())

	userJWTReq := httptest.NewRequest(http.MethodGet, "/api/v1/automations/runners/runner-1/jobs:acquire?sessionId=session-1", nil)
	setAutomationRunnerHeaders(userJWTReq, humanAuth, workspaceID)
	userJWTRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(userJWTRec, userJWTReq)
	require.Equal(t, http.StatusUnauthorized, userJWTRec.Code, userJWTRec.Body.String())

	revokedAuth := registerAutomationRunnerSession(t, srv, humanAuth, workspaceID, "revoked-runner", "revoked-session", nil)
	revokedResult := srv.meshxDB.Exec("UPDATE automation_runner_tokens SET revoked_at = ?, updated_at = ? WHERE workspace_id = ? AND runner_id = ? AND token_kind = ?", "2026-06-05 00:00:00", "2026-06-05 00:00:00", workspaceID, "revoked-runner", "session")
	require.NoError(t, revokedResult.Error)
	revokedReq := httptest.NewRequest(http.MethodGet, "/api/v1/automations/runners/revoked-runner/jobs:acquire?sessionId=revoked-session", nil)
	setAutomationRunnerHeaders(revokedReq, revokedAuth, workspaceID)
	revokedRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(revokedRec, revokedReq)
	require.Equal(t, http.StatusUnauthorized, revokedRec.Code, revokedRec.Body.String())

	missingSessionReq := httptest.NewRequest(http.MethodGet, "/api/v1/automations/runners/runner-1/jobs:acquire?limit=1", nil)
	setAutomationRunnerHeaders(missingSessionReq, runnerAuth, workspaceID)
	missingSessionRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(missingSessionRec, missingSessionReq)
	require.Equal(t, http.StatusBadRequest, missingSessionRec.Code, missingSessionRec.Body.String())

	invalidShardReq := httptest.NewRequest(http.MethodGet, "/api/v1/automations/runners/runner-1/jobs:acquire?sessionId=session-1&shardBucket=-1", nil)
	setAutomationRunnerHeaders(invalidShardReq, runnerAuth, workspaceID)
	invalidShardRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(invalidShardRec, invalidShardReq)
	require.Equal(t, http.StatusBadRequest, invalidShardRec.Code, invalidShardRec.Body.String())

	job := acquireAutomationRunnerJobWithQuery(t, srv, runnerAuth, workspaceID, "runner-1", "session-1", "leaseSeconds=1")

	missingCallbackSessionReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-1/jobs/%d/completed", job.QueueJobID), bytes.NewBufferString(`{"outputJson":"{}"}`))
	setAutomationRunnerHeaders(missingCallbackSessionReq, runnerAuth, workspaceID)
	missingCallbackSessionRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(missingCallbackSessionRec, missingCallbackSessionReq)
	require.Equal(t, http.StatusBadRequest, missingCallbackSessionRec.Code, missingCallbackSessionRec.Body.String())

	expireSessionResult := srv.meshxDB.Exec("UPDATE automation_runner_sessions SET expires_at = ? WHERE workspace_id = ? AND runner_id = ? AND session_id = ?", "2000-01-01 00:00:00", workspaceID, "runner-1", "session-1")
	require.NoError(t, expireSessionResult.Error)
	require.Equal(t, int64(1), expireSessionResult.RowsAffected)
	expiredSessionCompleteReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-1/jobs/%d/completed", job.QueueJobID), bytes.NewBufferString(`{"sessionId":"session-1","outputJson":"{}"}`))
	setAutomationRunnerHeaders(expiredSessionCompleteReq, runnerAuth, workspaceID)
	expiredSessionCompleteRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(expiredSessionCompleteRec, expiredSessionCompleteReq)
	require.Equal(t, http.StatusUnauthorized, expiredSessionCompleteRec.Code, expiredSessionCompleteRec.Body.String())
	restoreSessionResult := srv.meshxDB.Exec("UPDATE automation_runner_sessions SET expires_at = ? WHERE workspace_id = ? AND runner_id = ? AND session_id = ?", "2099-01-01 00:00:00", workspaceID, "runner-1", "session-1")
	require.NoError(t, restoreSessionResult.Error)
	require.Equal(t, int64(1), restoreSessionResult.RowsAffected)

	expireResult := srv.meshxDB.Exec("UPDATE automation_jobs SET lease_until = ? WHERE id = ?", "2000-01-01 00:00:00", job.QueueJobID)
	require.NoError(t, expireResult.Error)
	require.Equal(t, int64(1), expireResult.RowsAffected)
	expiredCompleteReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-1/jobs/%d/completed", job.QueueJobID), bytes.NewBufferString(`{"sessionId":"session-1","outputJson":"{}"}`))
	setAutomationRunnerHeaders(expiredCompleteReq, runnerAuth, workspaceID)
	expiredCompleteRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(expiredCompleteRec, expiredCompleteReq)
	require.Equal(t, http.StatusNotFound, expiredCompleteRec.Code, expiredCompleteRec.Body.String())
}

func TestAutomationRunnerLeaseOwnerIncludesRunnerIdentity(t *testing.T) {
	automationYAML := "name: Lease owner flow\non:\n  manual: {}\njobs:\n  build:\n    steps:\n      - run: echo build\n"
	srv := newAutomationRunnerTestServer(t, automationYAML, "8181818181818181818181818181818181818181")
	humanAuth := "Bearer " + jwt(map[string]any{"sub": "runner-lease-owner-user"})
	workspaceID := "ws-runner-lease-owner"
	_ = startAutomationRunnerTestRun(t, srv, humanAuth, workspaceID, ".agentteam/runner.yml")
	runnerOneAuth := registerAutomationRunnerSession(t, srv, humanAuth, workspaceID, "runner-1", "shared-session", []string{"ubuntu-latest"})
	runnerTwoAuth := registerAutomationRunnerSession(t, srv, humanAuth, workspaceID, "runner-2", "shared-session", []string{"ubuntu-latest"})

	job := acquireAutomationRunnerJob(t, srv, runnerOneAuth, workspaceID, "runner-1", "shared-session")
	crossAcceptReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-2/jobs/%d/accepted", job.QueueJobID), bytes.NewBufferString(`{"sessionId":"shared-session"}`))
	setAutomationRunnerHeaders(crossAcceptReq, runnerTwoAuth, workspaceID)
	crossAcceptRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(crossAcceptRec, crossAcceptReq)
	require.Equal(t, http.StatusNotFound, crossAcceptRec.Code, crossAcceptRec.Body.String())

	crossLogReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-2/jobs/%d/logs", job.QueueJobID), bytes.NewBufferString(`{"sessionId":"shared-session","stream":"stdout","content":"wrong runner","sequence":1}`))
	setAutomationRunnerHeaders(crossLogReq, runnerTwoAuth, workspaceID)
	crossLogRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(crossLogRec, crossLogReq)
	require.Equal(t, http.StatusNotFound, crossLogRec.Code, crossLogRec.Body.String())

	crossCompleteReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-2/jobs/%d/completed", job.QueueJobID), bytes.NewBufferString(`{"sessionId":"shared-session","outputJson":"{}"}`))
	setAutomationRunnerHeaders(crossCompleteReq, runnerTwoAuth, workspaceID)
	crossCompleteRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(crossCompleteRec, crossCompleteReq)
	require.Equal(t, http.StatusNotFound, crossCompleteRec.Code, crossCompleteRec.Body.String())
}

func TestAutomationRunnerArtifactUploadBodyLimit(t *testing.T) {
	automationYAML := "name: Artifact body limit flow\non:\n  manual: {}\njobs:\n  build:\n    steps:\n      - run: echo build\n"
	srv := newAutomationRunnerTestServerWithAutomationConfig(t, automationYAML, "8282828282828282828282828282828282828282", config.AutomationConfig{MaxArtifactBytes: 1})
	humanAuth := "Bearer " + jwt(map[string]any{"sub": "runner-artifact-limit-user"})
	workspaceID := "ws-runner-artifact-limit"
	runnerAuth := registerAutomationRunnerSession(t, srv, humanAuth, workspaceID, "runner-1", "session-1", []string{"ubuntu-latest"})

	body, err := json.Marshal(map[string]any{"sessionId": "session-1", "name": "huge.bin", "contentBase64": strings.Repeat("A", 200*1024), "size": 1, "sha256": strings.Repeat("0", 64)})
	require.NoError(t, err)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/automations/runners/runner-1/jobs/1/artifacts", bytes.NewBuffer(body))
	setAutomationRunnerHeaders(req, runnerAuth, workspaceID)
	rec := httptest.NewRecorder()
	srv.handler.ServeHTTP(rec, req)
	require.Equal(t, http.StatusRequestEntityTooLarge, rec.Code, rec.Body.String())
	require.Contains(t, rec.Body.String(), "Automation runner request body exceeds maximum size")

	trailingReq := httptest.NewRequest(http.MethodPost, "/api/v1/automations/runners/runner-1/jobs/1/artifacts", strings.NewReader(`{"sessionId":"session-1"}`+strings.Repeat(" ", 200*1024)))
	setAutomationRunnerHeaders(trailingReq, runnerAuth, workspaceID)
	trailingRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(trailingRec, trailingReq)
	require.Equal(t, http.StatusRequestEntityTooLarge, trailingRec.Code, trailingRec.Body.String())

	extraReq := httptest.NewRequest(http.MethodPost, "/api/v1/automations/runners/runner-1/jobs/1/artifacts", strings.NewReader(`{"sessionId":"session-1"} {}`))
	setAutomationRunnerHeaders(extraReq, runnerAuth, workspaceID)
	extraRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(extraRec, extraReq)
	require.Equal(t, http.StatusBadRequest, extraRec.Code, extraRec.Body.String())
	require.Contains(t, extraRec.Body.String(), "Invalid JSON body")
}

func TestAutomationRunnerArtifactAuthenticatesBeforeDecode(t *testing.T) {
	automationYAML := "name: Artifact auth flow\non:\n  manual: {}\njobs:\n  build:\n    steps:\n      - run: echo build\n"
	srv := newAutomationRunnerTestServer(t, automationYAML, "8383838383838383838383838383838383838383")
	req := httptest.NewRequest(http.MethodPost, "/api/v1/automations/runners/runner-1/jobs/1/artifacts", bytes.NewBufferString(`{"sessionId":"session-1","name":"artifact.bin","contentBase64":"!!!!","size":1,"sha256":"0000000000000000000000000000000000000000000000000000000000000000"}`))
	setAutomationRunnerHeaders(req, "Bearer invalid-runner-token", "ws-runner-artifact-auth")
	rec := httptest.NewRecorder()
	srv.handler.ServeHTTP(rec, req)
	require.Equal(t, http.StatusUnauthorized, rec.Code, rec.Body.String())
}

func TestAutomationRunnerHeartbeatReceivesCancelSignal(t *testing.T) {
	automationYAML := "name: Cancel signal flow\non:\n  manual: {}\njobs:\n  build:\n    steps:\n      - run: sleep 60\n"
	srv := newAutomationRunnerTestServer(t, automationYAML, "abababababababababababababababababababab")
	humanAuth := "Bearer " + jwt(map[string]any{"sub": "runner-cancel-user"})
	workspaceID := "ws-runner-cancel-signal"
	run := startAutomationRunnerTestRun(t, srv, humanAuth, workspaceID, ".agentteam/runner.yml")
	runnerAuth := registerAutomationRunnerSession(t, srv, humanAuth, workspaceID, "runner-1", "session-1", nil)
	job := acquireAutomationRunnerJob(t, srv, runnerAuth, workspaceID, "runner-1", "session-1")

	acceptReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-1/jobs/%d/accepted", job.QueueJobID), bytes.NewBufferString(`{"sessionId":"session-1"}`))
	setAutomationRunnerHeaders(acceptReq, runnerAuth, workspaceID)
	acceptRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(acceptRec, acceptReq)
	require.Equal(t, http.StatusOK, acceptRec.Code, acceptRec.Body.String())

	cancelReq := httptest.NewRequest(http.MethodPost, "/api/v1/automations/runs/"+run.ID+"/cancel", bytes.NewBufferString(`{"reason":"stop active job"}`))
	setAutomationRunnerHeaders(cancelReq, humanAuth, workspaceID)
	cancelRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(cancelRec, cancelReq)
	require.Equal(t, http.StatusOK, cancelRec.Code, cancelRec.Body.String())

	heartbeatReq := httptest.NewRequest(http.MethodPost, "/api/v1/automations/runners/runner-1/sessions/session-1/heartbeat", bytes.NewBufferString(fmt.Sprintf(`{"queueJobId":%d}`, job.QueueJobID)))
	setAutomationRunnerHeaders(heartbeatReq, runnerAuth, workspaceID)
	heartbeatRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(heartbeatRec, heartbeatReq)
	require.Equal(t, http.StatusOK, heartbeatRec.Code, heartbeatRec.Body.String())
	require.Contains(t, heartbeatRec.Body.String(), `"cancelRequested":true`)

	completeReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/automations/runners/runner-1/jobs/%d/completed", job.QueueJobID), bytes.NewBufferString(`{"sessionId":"session-1","outputJson":"{}"}`))
	setAutomationRunnerHeaders(completeReq, runnerAuth, workspaceID)
	completeRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(completeRec, completeReq)
	require.Equal(t, http.StatusNotFound, completeRec.Code, completeRec.Body.String())
}

func TestAutomationRunnerAcquireLongPollClaimsJobCreatedDuringWait(t *testing.T) {
	automationYAML := "name: Long poll flow\non:\n  manual: {}\njobs:\n  build:\n    steps:\n      - run: echo long-poll\n"
	srv := newAutomationRunnerTestServer(t, automationYAML, "ffffffffffffffffffffffffffffffffffffffff")
	humanAuth := "Bearer " + jwt(map[string]any{"sub": "runner-long-poll-user"})
	workspaceID := "ws-runner-long-poll"
	runnerAuth := 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=1", nil)
		setAutomationRunnerHeaders(acquireReq, runnerAuth, 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 any job was queued: status=%d body=%s", result.code, string(result.body))
	case <-time.After(25 * time.Millisecond):
	}
	_ = startAutomationRunnerTestRun(t, srv, humanAuth, workspaceID, ".agentteam/runner.yml")

	select {
	case result := <-resultCh:
		require.Equal(t, http.StatusOK, result.code, string(result.body))
		var response struct {
			Jobs []automationRunnerAcquiredJob `json:"jobs"`
		}
		require.NoError(t, json.Unmarshal(result.body, &response))
		require.Len(t, response.Jobs, 1)
		require.Equal(t, "build", response.Jobs[0].JobID)
		require.NotEmpty(t, response.Jobs[0].LeaseOwner)
		require.NotEqual(t, "session-1", response.Jobs[0].LeaseOwner)
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for long-poll acquire response")
	}
}

func TestAutomationContractCompatibilityAliases(t *testing.T) {
	automationYAML := "name: Contract flow\non:\n  manual: {}\njobs:\n  build:\n    steps:\n      - run: echo contract\n"
	srv := newAutomationRunnerTestServer(t, automationYAML, "dddddddddddddddddddddddddddddddddddddddd")
	auth := "Bearer " + jwt(map[string]any{"sub": "contract-user"})
	workspaceID := "ws-automation-contract"
	body := bytes.NewBufferString(`{"ref":"refs/heads/main","definitionPath":".agentteam/runner.yml","inputs":"{\"env\":\"prod\"}"}`)
	startReq := httptest.NewRequest(http.MethodPost, "/api/v1/automations/repositories/pingcap/agentteam/runs", body)
	setAutomationRunnerHeaders(startReq, auth, workspaceID)
	startRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(startRec, startReq)
	require.Equal(t, http.StatusCreated, startRec.Code, startRec.Body.String())
	require.Contains(t, startRec.Body.String(), `"inputJson":"{\"env\":\"prod\"}"`)
	var run automationRunnerTestRun
	require.NoError(t, json.Unmarshal(startRec.Body.Bytes(), &run))
	require.NotEmpty(t, run.ID)

	jobsReq := httptest.NewRequest(http.MethodGet, "/api/v1/automations/runs/"+run.ID+"/jobs", nil)
	setAutomationRunnerHeaders(jobsReq, auth, workspaceID)
	jobsRec := httptest.NewRecorder()
	srv.handler.ServeHTTP(jobsRec, jobsReq)
	require.Equal(t, http.StatusOK, jobsRec.Code, jobsRec.Body.String())
	require.Contains(t, jobsRec.Body.String(), `"jobId":"build"`)
}

func TestAutomationRunStartReturnsStructuredScheduleRejectionCode(t *testing.T) {
	automationYAML := "name: Schedule flow\non: push\njobs:\n  build:\n    steps:\n      - uses: actions/checkout@v4\n"
	srv := newAutomationRunnerTestServer(t, automationYAML, "cdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcd")
	auth := "Bearer " + jwt(map[string]any{"sub": "schedule-user"})
	workspaceID := "ws-automation-schedule-rejection"
	body := bytes.NewBufferString(`{"ref":"cdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcd","definitionPath":".agentteam/runner.yml","triggerKind":"schedule","triggerRef":"arun_1","targetKind":"schedule","targetRef":"sched_1"}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/automations/repositories/pingcap/agentteam/runs", body)
	setAutomationRunnerHeaders(req, auth, workspaceID)
	rec := httptest.NewRecorder()
	srv.handler.ServeHTTP(rec, req)
	require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String())
	var payload map[string]any
	require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &payload))
	require.Equal(t, "capability_unsupported_syntax", payload["error"])
	require.NotEmpty(t, payload["message"])
}

func TestAutomationRunStartCreatesFailedPushRunWithTriggerContext(t *testing.T) {
	commitSHA := "edededededededededededededededededededed"
	automationYAML := "name: Push flow\non:\n  push:\n    branches:\n      - main\njobs:\n  build:\n    steps:\n      - uses: actions/checkout@v4\n"
	srv := newAutomationRunnerTestServer(t, automationYAML, commitSHA)
	auth := "Bearer " + jwt(map[string]any{"sub": "push-user"})
	workspaceID := "ws-automation-push-rejection"
	body := bytes.NewBufferString(`{"ref":"` + commitSHA + `","definitionPath":".agentteam/runner.yml","branch":"main","beforeSha":"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb","actorRef":"user:1","actorDisplay":"Alice","triggerKind":"push","triggerRef":"evt_1"}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/automations/repositories/pingcap/agentteam/runs", body)
	setAutomationRunnerHeaders(req, auth, workspaceID)
	rec := httptest.NewRecorder()
	srv.handler.ServeHTTP(rec, req)
	require.Equal(t, http.StatusCreated, rec.Code, rec.Body.String())
	var payload map[string]any
	require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &payload))
	require.Equal(t, "failed", payload["status"])
	errorPayload, ok := payload["error"].(map[string]any)
	require.True(t, ok)
	require.Equal(t, "capability_unsupported_syntax", errorPayload["code"])
	triggerContext, ok := payload["triggerContext"].(map[string]any)
	require.True(t, ok)
	require.Equal(t, "push", triggerContext["triggerKind"])
	require.Equal(t, "refs/heads/main", triggerContext["ref"])
	require.Equal(t, "main", triggerContext["branch"])
	require.Equal(t, commitSHA, triggerContext["afterSha"])
	require.Equal(t, "Alice", triggerContext["actorDisplay"])
}