package automation

import (
	"context"
	"encoding/base64"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/internal/automation/blob"
	"github.com/tidbcloud/AgentTeam/internal/service/automation/runtime"
	automationstore "github.com/tidbcloud/AgentTeam/internal/store/automation"
)

func TestAutomationRunnerAcquireLongPollRechecksRevokedSession(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	svc, store, now := newAutomationRunnerTestService(t)

	registration, err := svc.RegisterAutomationRunner(ctx, "ws-1", "user-1", AutomationRunnerRegisterMutation{RunnerID: "runner-1", Name: "runner-1"})
	require.NoError(t, err)
	session, err := svc.CreateAutomationRunnerSession(ctx, "ws-1", registration.RegistrationToken, AutomationRunnerSessionMutation{RunnerID: "runner-1", SessionID: "session-1"})
	require.NoError(t, err)

	secondPoll := make(chan struct{})
	allowSecondPoll := make(chan struct{})
	var releaseSecondPoll sync.Once
	release := func() {
		releaseSecondPoll.Do(func() { close(allowSecondPoll) })
	}
	defer release()
	var nowMu sync.Mutex
	nowCalls := 0
	svc.now = func() time.Time {
		nowMu.Lock()
		nowCalls++
		call := nowCalls
		nowMu.Unlock()
		if call == 2 {
			close(secondPoll)
			<-allowSecondPoll
		}
		return now()
	}

	type acquireResult struct {
		jobs []AutomationRunnerJob
		err  error
	}
	results := make(chan acquireResult, 1)
	go func() {
		jobs, err := svc.AcquireAutomationRunnerJobs(ctx, "ws-1", session.SessionToken, AutomationRunnerAcquireMutation{RunnerID: "runner-1", SessionID: "session-1", Limit: 1, WaitSeconds: 1})
		results <- acquireResult{jobs: jobs, err: err}
	}()

	select {
	case <-secondPoll:
	case <-time.After(2 * time.Second):
		t.Fatal("AcquireAutomationRunnerJobs did not enter long poll")
	}
	_, err = svc.RegisterAutomationRunner(ctx, "ws-1", "user-1", AutomationRunnerRegisterMutation{RunnerID: "runner-1", Name: "runner-1"})
	require.NoError(t, err)
	job := createAutomationRunnerReadyJob(t, ctx, store, now(), "run-revoked-acquire", "job-revoked-acquire")
	release()

	select {
	case result := <-results:
		require.ErrorIs(t, result.err, ErrUnauthorized)
		require.Empty(t, result.jobs)
	case <-time.After(2 * time.Second):
		t.Fatal("AcquireAutomationRunnerJobs did not return after session revocation")
	}

	stored, err := store.GetJob(ctx, job.ID)
	require.NoError(t, err)
	require.Equal(t, automationstore.AutomationJobStatusReady, stored.Status)
	require.Empty(t, stored.LeaseOwner)
}

func TestAutomationRunnerAcquireRejectsReservedRunnerGroupSession(t *testing.T) {
	ctx := context.Background()
	svc, store, now := newAutomationRunnerTestService(t)
	current := now()
	registrationToken := "reserved-registration-token"
	require.NoError(t, store.RegisterRunnerWithToken(ctx, automationstore.AutomationRunner{
		WorkspaceID:  "ws-1",
		RunnerID:     "runner-reserved",
		Name:         "reserved",
		RunnerGroup:  runtime.ServerManagedRunnerGroup,
		LabelsJSON:   automationRunnerJSON([]string{"ubuntu-latest"}),
		Status:       automationstore.AutomationRunnerStatusActive,
		RegisteredAt: current,
		CreatedAt:    current,
		UpdatedAt:    current,
	}, automationstore.AutomationRunnerToken{
		TokenHash: automationRunnerTokenHash(registrationToken),
		TokenKind: automationstore.AutomationRunnerTokenKindRegistration,
		ExpiresAt: current.Add(time.Hour),
		CreatedAt: current,
		UpdatedAt: current,
	}, current))
	session, err := svc.CreateAutomationRunnerSession(ctx, "ws-1", registrationToken, AutomationRunnerSessionMutation{RunnerID: "runner-reserved", SessionID: "session-reserved"})
	require.NoError(t, err)
	_, err = store.CreateRun(ctx, automationstore.AutomationRun{ID: "run-reserved", WorkspaceID: "ws-1", Name: "Reserved", Status: automationstore.AutomationRunStatusRunning, CreatedAt: current, UpdatedAt: current})
	require.NoError(t, err)
	require.NoError(t, store.CreateJobRuns(ctx, []automationstore.AutomationJobRun{{ID: "job-reserved", WorkspaceID: "ws-1", RunID: "run-reserved", JobID: "build", JobType: "github_actions_job", Status: automationstore.AutomationJobRunStatusQueued, ScheduledAt: current, CreatedAt: current, UpdatedAt: current}}))
	job, err := store.EnqueueJob(ctx, automationstore.AutomationJob{WorkspaceID: "ws-1", RunID: "run-reserved", JobRunID: "job-reserved", QueueName: runtime.DefaultQueueName, Status: automationstore.AutomationJobStatusReady, RunnerGroup: runtime.ServerManagedRunnerGroup, LabelsJSON: automationRunnerJSON([]string{"ubuntu-latest"}), ScheduledAt: current, MaxAttempts: 3, IdempotencyKey: "run-reserved:build", ShardBucket: 0, CreatedAt: current, UpdatedAt: current})
	require.NoError(t, err)

	jobs, err := svc.AcquireAutomationRunnerJobs(ctx, "ws-1", session.SessionToken, AutomationRunnerAcquireMutation{RunnerID: "runner-reserved", SessionID: "session-reserved", Limit: 1})

	require.ErrorIs(t, err, ErrInvalidArgument)
	require.Contains(t, err.Error(), "group is reserved")
	require.Empty(t, jobs)
	stored, err := store.GetJob(ctx, job.ID)
	require.NoError(t, err)
	require.Equal(t, automationstore.AutomationJobStatusReady, stored.Status)
	require.Empty(t, stored.LeaseOwner)
}

func TestAutomationRunnerAcquireLongPollWakesOnWorkspaceSignal(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	svc, store, now := newAutomationRunnerTestService(t)
	current := now()
	registration, err := svc.RegisterAutomationRunner(ctx, "ws-1", "user-1", AutomationRunnerRegisterMutation{RunnerID: "runner-1", Name: "runner-1"})
	require.NoError(t, err)
	session, err := svc.CreateAutomationRunnerSession(ctx, "ws-1", registration.RegistrationToken, AutomationRunnerSessionMutation{RunnerID: "runner-1", SessionID: "session-1"})
	require.NoError(t, err)

	firstClaim := make(chan struct{})
	var calls atomic.Int32
	svc.now = func() time.Time {
		if calls.Add(1) == 1 {
			close(firstClaim)
		}
		return current
	}

	type acquireResult struct {
		jobs []AutomationRunnerJob
		err  error
	}
	results := make(chan acquireResult, 1)
	go func() {
		jobs, err := svc.AcquireAutomationRunnerJobs(ctx, "ws-1", session.SessionToken, AutomationRunnerAcquireMutation{RunnerID: "runner-1", SessionID: "session-1", Limit: 1, WaitSeconds: 5})
		results <- acquireResult{jobs: jobs, err: err}
	}()

	select {
	case <-firstClaim:
	case <-time.After(time.Second):
		t.Fatal("AcquireAutomationRunnerJobs did not attempt initial claim")
	}
	job := createAutomationRunnerReadyJob(t, ctx, store, current, "run-signal-acquire", "job-signal-acquire")
	svc.signalRunnerWorkspace("ws-1")

	select {
	case result := <-results:
		require.NoError(t, result.err)
		require.Len(t, result.jobs, 1)
		require.Equal(t, job.ID, result.jobs[0].QueueJobID)
	case <-time.After(time.Second):
		t.Fatal("AcquireAutomationRunnerJobs did not wake on workspace signal")
	}
}

func TestAutomationRunnerWorkspaceWaitUnregistersOnCancel(t *testing.T) {
	svc, _, _ := newAutomationRunnerTestService(t)
	wake, cancel := svc.runnerWorkspaceWait("ws-1")
	require.NotNil(t, wake)

	cancel()

	svc.runnerWaitMu.Lock()
	defer svc.runnerWaitMu.Unlock()
	require.Empty(t, svc.runnerWaiters)
}

func TestAutomationRunnerHeartbeatSkipsWriteUntilRefreshWindow(t *testing.T) {
	ctx := context.Background()
	db, err := openServiceTestSQLiteDB(t)
	require.NoError(t, err)
	migrateAutomationServiceTestSchema(t, db, &automationstore.AutomationRunner{}, &automationstore.AutomationRunnerToken{}, &automationstore.AutomationRunnerSession{}, &automationstore.AutomationActionsRunnerCredential{}, &automationstore.AutomationLogChunk{}, &automationstore.AutomationArtifact{})
	store := automationstore.New(db)
	current := time.Date(2026, 6, 6, 12, 0, 0, 0, time.UTC)
	svc := NewAutomationsService(store, func() time.Time { return current }, nil)

	registration, err := svc.RegisterAutomationRunner(ctx, "ws-1", "user-1", AutomationRunnerRegisterMutation{RunnerID: "runner-1", Name: "runner-1"})
	require.NoError(t, err)
	session, err := svc.CreateAutomationRunnerSession(ctx, "ws-1", registration.RegistrationToken, AutomationRunnerSessionMutation{RunnerID: "runner-1", SessionID: "session-1"})
	require.NoError(t, err)
	tokenHash := automationRunnerTokenHash(session.SessionToken)

	heartbeat, err := svc.HeartbeatAutomationRunnerSessionByTokenHash(ctx, "ws-1", tokenHash, AutomationRunnerHeartbeatMutation{RunnerID: "runner-1", SessionID: "session-1"})
	require.NoError(t, err)
	require.Equal(t, session.HeartbeatAt, heartbeat.HeartbeatAt)
	require.Equal(t, session.ExpiresAt, heartbeat.ExpiresAt)
	stored, err := store.GetActiveRunnerSessionByToken(ctx, tokenHash, "ws-1", "runner-1", "session-1", current)
	require.NoError(t, err)
	require.Equal(t, session.HeartbeatAt, stored.HeartbeatAt)
	require.Equal(t, session.ExpiresAt, stored.ExpiresAt)

	current = session.ExpiresAt.Add(-automationRunnerHeartbeatRefresh).Add(time.Second)
	heartbeat, err = svc.HeartbeatAutomationRunnerSessionByTokenHash(ctx, "ws-1", tokenHash, AutomationRunnerHeartbeatMutation{RunnerID: "runner-1", SessionID: "session-1"})
	require.NoError(t, err)
	require.Equal(t, current, heartbeat.HeartbeatAt)
	require.Equal(t, current.Add(automationRunnerSessionTTL), heartbeat.ExpiresAt)
}

func TestAutomationRunnerAcquireIncludesTriggerContextJSON(t *testing.T) {
	ctx := context.Background()
	svc, store, now := newAutomationRunnerTestService(t)
	registration, err := svc.RegisterAutomationRunner(ctx, "ws-1", "user-1", AutomationRunnerRegisterMutation{RunnerID: "runner-1", Name: "runner-1"})
	require.NoError(t, err)
	session, err := svc.CreateAutomationRunnerSession(ctx, "ws-1", registration.RegistrationToken, AutomationRunnerSessionMutation{RunnerID: "runner-1", SessionID: "session-1"})
	require.NoError(t, err)

	runNow := now()
	triggerContext := buildAutomationRunMetadata(&AutomationTriggerContext{
		TriggerKind:          "schedule",
		ScheduleDefinitionID: "sched_1",
		AutomationRunID:      "arun_1",
	})
	_, err = store.CreateRun(ctx, automationstore.AutomationRun{
		ID:          "run-trigger-context",
		WorkspaceID: "ws-1",
		Name:        "Acquire",
		Status:      automationstore.AutomationRunStatusRunning,
		Metadata:    triggerContext,
		CreatedAt:   runNow,
		UpdatedAt:   runNow,
	})
	require.NoError(t, err)
	require.NoError(t, store.CreateJobRuns(ctx, []automationstore.AutomationJobRun{{
		ID:          "job-trigger-context",
		WorkspaceID: "ws-1",
		RunID:       "run-trigger-context",
		JobID:       "build",
		JobType:     "github_actions_job",
		Status:      automationstore.AutomationJobRunStatusQueued,
		ScheduledAt: runNow,
		CreatedAt:   runNow,
		UpdatedAt:   runNow,
	}}))
	_, err = store.EnqueueJob(ctx, automationstore.AutomationJob{
		WorkspaceID:    "ws-1",
		RunID:          "run-trigger-context",
		JobRunID:       "job-trigger-context",
		QueueName:      runtime.DefaultQueueName,
		Status:         automationstore.AutomationJobStatusReady,
		ScheduledAt:    runNow,
		MaxAttempts:    3,
		IdempotencyKey: "run-trigger-context:build",
		ShardBucket:    0,
		CreatedAt:      runNow,
		UpdatedAt:      runNow,
	})
	require.NoError(t, err)

	jobs, err := svc.AcquireAutomationRunnerJobs(ctx, "ws-1", session.SessionToken, AutomationRunnerAcquireMutation{RunnerID: "runner-1", SessionID: "session-1", Limit: 1, WaitSeconds: 1})
	require.NoError(t, err)
	require.Len(t, jobs, 1)
	require.JSONEq(t, `{"trigger_kind":"schedule","schedule_definition_id":"sched_1","automation_run_id":"arun_1"}`, jobs[0].TriggerContextJSON)
}

func TestAutomationRunnerCallbacksScrubReferencedSecretsFromPersistedPayloads(t *testing.T) {
	ctx := context.Background()
	svc, store, now := newAutomationRunnerTestService(t)
	blobStore, err := blob.NewFilesystem(t.TempDir())
	require.NoError(t, err)
	svc.SetBlobStore(blobStore)
	svc.SetAutomationCredentialsKey([]byte("0123456789abcdef0123456789abcdef"))
	const secretValue = "plain-secret-token"
	_, err = svc.UpsertAutomationSecret(ctx, "ws-1", "user-1", AutomationSecretMutation{
		RepositoryOwner: "pingcap",
		RepositoryName:  "agentteam",
		Name:            "deploy_token",
		Value:           secretValue,
	})
	require.NoError(t, err)
	registration, err := svc.RegisterAutomationRunner(ctx, "ws-1", "user-1", AutomationRunnerRegisterMutation{RunnerID: "runner-1", Name: "runner-1"})
	require.NoError(t, err)
	session, err := svc.CreateAutomationRunnerSession(ctx, "ws-1", registration.RegistrationToken, AutomationRunnerSessionMutation{RunnerID: "runner-1", SessionID: "session-1"})
	require.NoError(t, err)

	completeJob := createAutomationRunnerSecretLeasedJob(t, ctx, store, now(), session.LeaseOwner, "run-secret-complete", "job-secret-complete")
	completed, err := svc.CompleteAutomationRunnerJob(ctx, "ws-1", session.SessionToken, AutomationRunnerJobMutation{
		RunnerID:   "runner-1",
		SessionID:  "session-1",
		QueueJobID: completeJob.ID,
		OutputJSON: `{"token":"` + secretValue + `","nested":["` + secretValue + `"]}`,
	})
	require.NoError(t, err)
	require.NotContains(t, completed.OutputJSON, secretValue)
	require.Contains(t, completed.OutputJSON, "***")

	failJob := createAutomationRunnerSecretLeasedJob(t, ctx, store, now(), session.LeaseOwner, "run-secret-fail", "job-secret-fail")
	failed, err := svc.FailAutomationRunnerJob(ctx, "ws-1", session.SessionToken, AutomationRunnerJobMutation{
		RunnerID:        "runner-1",
		SessionID:       "session-1",
		QueueJobID:      failJob.ID,
		Message:         "failed with " + secretValue,
		OutputJSON:      `{"stderr":"` + secretValue + `"}`,
		TerminalFailure: true,
	})
	require.NoError(t, err)
	require.NotContains(t, failed.ErrorJSON, secretValue)
	require.NotContains(t, failed.OutputJSON, secretValue)
	require.Contains(t, failed.ErrorJSON, "***")
	failedRun, err := store.GetRun(ctx, "ws-1", "run-secret-fail")
	require.NoError(t, err)
	require.NotContains(t, failedRun.ErrorJSON, secretValue)

	logJob := createAutomationRunnerSecretLeasedJob(t, ctx, store, now(), session.LeaseOwner, "run-secret-log", "job-secret-log")
	logPayload := "before " + secretValue + " after"
	_, err = svc.AppendAutomationRunnerJobLog(ctx, "ws-1", session.SessionToken, AutomationRunnerJobMutation{
		RunnerID:    "runner-1",
		SessionID:   "session-1",
		QueueJobID:  logJob.ID,
		LogStream:   "stdout",
		LogContent:  logPayload,
		LogSequence: 1,
		LogSHA256:   blob.SHA256Hex([]byte(logPayload)),
	})
	require.NoError(t, err)
	logContent, err := svc.ReadAutomationLogChunkContent(ctx, "ws-1", "user-1", "run-secret-log", "job-secret-log", "stdout", 1)
	require.NoError(t, err)
	require.NotContains(t, logContent.Content, secretValue)
	require.Equal(t, "before *** after", logContent.Content)

	artifactJob := createAutomationRunnerSecretLeasedJob(t, ctx, store, now(), session.LeaseOwner, "run-secret-artifact", "job-secret-artifact")
	artifactPayload := []byte("artifact")
	artifactEvent, err := svc.AppendAutomationRunnerJobArtifact(ctx, "ws-1", session.SessionToken, AutomationRunnerJobMutation{
		RunnerID:              "runner-1",
		SessionID:             "session-1",
		QueueJobID:            artifactJob.ID,
		ArtifactName:          "artifact-" + secretValue + ".txt",
		ArtifactSize:          int64(len(artifactPayload)),
		ArtifactContentBase64: base64.StdEncoding.EncodeToString(artifactPayload),
		ArtifactSHA256:        blob.SHA256Hex(artifactPayload),
		ArtifactContentType:   "text/plain",
		MetadataJSON:          `{"label":"` + secretValue + `"}`,
	})
	require.NoError(t, err)
	require.NotContains(t, artifactEvent.MetadataJSON, secretValue)
	require.Contains(t, artifactEvent.MetadataJSON, "***")
	artifacts, err := store.ListArtifacts(ctx, "ws-1", "run-secret-artifact", "job-secret-artifact", 1, 10)
	require.NoError(t, err)
	require.Len(t, artifacts, 1)
	require.NotContains(t, artifacts[0].Name, secretValue)
	require.NotContains(t, artifacts[0].ObjectKey, secretValue)
}

func newAutomationRunnerTestService(t *testing.T) (*AutomationsService, *automationstore.Store, func() time.Time) {
	t.Helper()
	db, err := openServiceTestConcurrentSQLiteDB(t)
	require.NoError(t, err)
	migrateAutomationServiceTestSchema(t, db, &automationstore.AutomationRunner{}, &automationstore.AutomationRunnerToken{}, &automationstore.AutomationRunnerSession{}, &automationstore.AutomationActionsRunnerCredential{}, &automationstore.AutomationLogChunk{}, &automationstore.AutomationArtifact{}, &automationstore.AutomationSecret{})
	store := automationstore.New(db)
	current := time.Date(2026, 6, 6, 11, 0, 0, 0, time.UTC)
	return NewAutomationsService(store, func() time.Time { return current }, nil), store, func() time.Time { return current }
}

func createAutomationRunnerReadyJob(t *testing.T, ctx context.Context, store *automationstore.Store, now time.Time, runID string, jobRunID string) automationstore.AutomationJob {
	t.Helper()
	_, err := store.CreateRun(ctx, automationstore.AutomationRun{ID: runID, WorkspaceID: "ws-1", Name: "Acquire", Status: automationstore.AutomationRunStatusRunning, CreatedAt: now, UpdatedAt: now})
	require.NoError(t, err)
	require.NoError(t, store.CreateJobRuns(ctx, []automationstore.AutomationJobRun{{ID: jobRunID, WorkspaceID: "ws-1", RunID: runID, JobID: "build", JobType: "github_actions_job", Status: automationstore.AutomationJobRunStatusQueued, ScheduledAt: now, CreatedAt: now, UpdatedAt: now}}))
	job, err := store.EnqueueJob(ctx, automationstore.AutomationJob{WorkspaceID: "ws-1", RunID: runID, JobRunID: jobRunID, QueueName: runtime.DefaultQueueName, Status: automationstore.AutomationJobStatusReady, ScheduledAt: now, MaxAttempts: 3, IdempotencyKey: runID + ":build", ShardBucket: 0, CreatedAt: now, UpdatedAt: now})
	require.NoError(t, err)
	return job
}

func createAutomationRunnerSecretLeasedJob(t *testing.T, ctx context.Context, store *automationstore.Store, now time.Time, leaseOwner string, runID string, jobRunID string) automationstore.AutomationJob {
	t.Helper()
	_, err := store.CreateRun(ctx, automationstore.AutomationRun{ID: runID, WorkspaceID: "ws-1", Name: "Secret callback", Status: automationstore.AutomationRunStatusRunning, RepositoryOwner: "pingcap", RepositoryName: "agentteam", CreatedAt: now, UpdatedAt: now})
	require.NoError(t, err)
	require.NoError(t, store.CreateJobRuns(ctx, []automationstore.AutomationJobRun{{
		ID:          jobRunID,
		WorkspaceID: "ws-1",
		RunID:       runID,
		JobID:       "build",
		JobType:     "github_actions_job",
		Status:      automationstore.AutomationJobRunStatusRunning,
		InputJSON:   `{"env":{"TOKEN":"${{ secrets.deploy_token }}"},"steps":[{"run":"echo hi"}]}`,
		ScheduledAt: now,
		CreatedAt:   now,
		UpdatedAt:   now,
	}}))
	leaseUntil := now.Add(time.Hour)
	job, err := store.EnqueueJob(ctx, automationstore.AutomationJob{WorkspaceID: "ws-1", RunID: runID, JobRunID: jobRunID, QueueName: runtime.DefaultQueueName, Status: automationstore.AutomationJobStatusLeased, LeaseOwner: leaseOwner, LeaseUntil: &leaseUntil, ScheduledAt: now, MaxAttempts: 1, IdempotencyKey: runID + ":build", ShardBucket: 0, CreatedAt: now, UpdatedAt: now})
	require.NoError(t, err)
	return job
}