package run_test

import (
	"context"
	"encoding/json"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/store/workflow/run"
)

func TestSaveTurnAndLoadRoundTripRunStateLargerThanSinglePayloadLimit(t *testing.T) {
	ctx := context.Background()
	db := startRollbackDB(t, ctx)
	fixture := seedRunFixture(t, ctx, db, "large_state")
	store := run.NewWithOptions(db, run.StoreOptions{WorkspaceGroupID: fixture.GroupID})
	claimed, err := store.ClaimRunnableRuns(ctx, run.ClaimRunnableParams{
		WorkerID: "graph", Now: fixture.Now, LeaseTTL: time.Minute, Limit: 1,
	})
	require.NoError(t, err)
	require.Len(t, claimed, 1)
	payload := strings.Repeat("x", 5<<20)
	stateJSON, err := json.Marshal(map[string]any{
		"status": "waiting", "input": map[string]any{}, "variables": map[string]any{},
		"activations": []any{}, "scopes": []any{}, "nextActivation": 0, "nextScope": 0,
		"payload": payload,
	})
	require.NoError(t, err)
	saved, err := store.SaveTurn(ctx, run.SaveTurnParams{
		Fence: claimed[0].Fence, ExpectedRevision: claimed[0].Run.Revision,
		Status: run.RunStatusWaiting, StateJSON: stateJSON,
		RunnableAt: fixture.Now.Add(time.Hour), StartedAt: claimed[0].Run.StartedAt,
		ReleaseLease: true, Now: fixture.Now.Add(time.Second),
	})
	require.NoError(t, err)
	require.Greater(t, len(saved.Run.StateJSON), 4<<20)
	loaded, err := store.LoadRun(ctx, fixture.WorkspaceID, fixture.RunID)
	require.NoError(t, err)
	require.JSONEq(t, string(stateJSON), string(loaded.StateJSON))
	canceled, _, err := store.CancelRun(ctx, run.CancelRunParams{
		WorkspaceID: fixture.WorkspaceID, RunID: fixture.RunID, Now: fixture.Now.Add(2 * time.Second),
	})
	require.NoError(t, err)
	require.Equal(t, run.RunStatusCanceled, canceled.Status)
	require.Greater(t, len(canceled.StateJSON), 4<<20)
}

func TestFailClaimedRunCancelsPendingChildrenBeforeTerminal(t *testing.T) {
	ctx := context.Background()
	db := startRollbackDB(t, ctx)
	fixture := seedRunFixture(t, ctx, db, "failure_pending")
	store, activationID, effectID := seedFailureChildren(t, ctx, db, fixture)

	claimed, err := store.ClaimRunnableRuns(ctx, run.ClaimRunnableParams{
		WorkerID: "graph-failure", Now: fixture.Now.Add(2 * time.Second), LeaseTTL: time.Minute, Limit: 1,
	})
	require.NoError(t, err)
	require.Len(t, claimed, 1)
	result, err := store.FailClaimedRun(ctx, run.FailClaimedRunParams{
		Fence: claimed[0].Fence, ExpectedRevision: claimed[0].Run.Revision,
		StateJSON: deterministicFailureState(), ErrorCode: "run.plan_version_invalid",
		ErrorJSON: json.RawMessage(`{"code":"run.plan_version_invalid","message":"invalid artifact"}`),
		Now:       fixture.Now.Add(2 * time.Second),
	})
	require.NoError(t, err)
	require.False(t, result.Draining)
	require.Equal(t, run.RunStatusFailed, result.Run.Status)
	require.NotNil(t, result.Run.CompletedAt)
	require.NotNil(t, result.Event)
	require.Equal(t, "run_failed", result.Event.Type)

	activation, err := store.LoadActivation(ctx, fixture.WorkspaceID, fixture.RunID, activationID)
	require.NoError(t, err)
	require.Equal(t, run.ActivationStatusCanceled, activation.Status)
	require.Equal(t, "run.plan_version_invalid", activation.ErrorCode)
	effect, err := store.LoadEffect(ctx, fixture.WorkspaceID, fixture.RunID, effectID)
	require.NoError(t, err)
	require.Equal(t, run.EffectStatusCanceled, effect.Status)
	require.Empty(t, effect.LeaseToken)
}

func TestFailClaimedRunDrainsCommittedDispatchBeforeTerminal(t *testing.T) {
	ctx := context.Background()
	db := startRollbackDB(t, ctx)
	fixture := seedRunFixture(t, ctx, db, "failure_dispatch")
	store, activationID, effectID := seedFailureChildren(t, ctx, db, fixture)

	candidates, err := store.ListPendingEffects(ctx, run.ListPendingEffectsParams{
		Now: fixture.Now.Add(time.Second), Limit: 10,
	})
	require.NoError(t, err)
	require.Len(t, candidates, 1)
	dispatching, ok, err := store.ClaimEffect(ctx, run.ClaimEffectParams{
		Candidate: candidates[0], WorkerID: "effect-worker", Now: fixture.Now.Add(time.Second), LeaseTTL: time.Minute,
	})
	require.NoError(t, err)
	require.True(t, ok)

	claimed, err := store.ClaimRunnableRuns(ctx, run.ClaimRunnableParams{
		WorkerID: "graph-failure", Now: fixture.Now.Add(2 * time.Second), LeaseTTL: time.Minute, Limit: 1,
	})
	require.NoError(t, err)
	require.Len(t, claimed, 1)
	first, err := store.FailClaimedRun(ctx, run.FailClaimedRunParams{
		Fence: claimed[0].Fence, ExpectedRevision: claimed[0].Run.Revision,
		StateJSON: deterministicFailureState(), ErrorCode: "run.plan_version_invalid",
		ErrorJSON: json.RawMessage(`{"code":"run.plan_version_invalid","message":"invalid artifact"}`),
		Now:       fixture.Now.Add(2 * time.Second),
	})
	require.NoError(t, err)
	require.True(t, first.Draining)
	require.Equal(t, run.RunStatusWaiting, first.Run.Status)
	require.Nil(t, first.Run.CompletedAt)
	require.Equal(t, *dispatching.LeaseExpiresAt, first.Run.RunnableAt)

	activation, err := store.LoadActivation(ctx, fixture.WorkspaceID, fixture.RunID, activationID)
	require.NoError(t, err)
	require.Equal(t, run.ActivationStatusCanceled, activation.Status)
	stillDispatching, err := store.LoadEffect(ctx, fixture.WorkspaceID, fixture.RunID, effectID)
	require.NoError(t, err)
	require.Equal(t, run.EffectStatusDispatching, stillDispatching.Status)
	require.Equal(t, dispatching.LeaseToken, stillDispatching.LeaseToken)

	// This models an external writer that already succeeded before the graph
	// discovered its deterministic failure. The in-flight receipt must remain
	// writable while the aggregate is draining.
	_, _, err = store.CompleteEffect(
		ctx, fixture.WorkspaceID, fixture.RunID, effectID, dispatching.LeaseToken,
		json.RawMessage(`{"port":"sent","output":{}}`), nil, fixture.Now.Add(3*time.Second),
	)
	require.NoError(t, err)

	claimed, err = store.ClaimRunnableRuns(ctx, run.ClaimRunnableParams{
		WorkerID: "graph-failure", Now: fixture.Now.Add(4 * time.Second), LeaseTTL: time.Minute, Limit: 1,
	})
	require.NoError(t, err)
	require.Len(t, claimed, 1)
	second, err := store.FailClaimedRun(ctx, run.FailClaimedRunParams{
		Fence: claimed[0].Fence, ExpectedRevision: claimed[0].Run.Revision,
		// A repeated failure cannot replace the first authoritative cause.
		StateJSON: deterministicFailureState(), ErrorCode: "run.different_failure",
		ErrorJSON: json.RawMessage(`{"code":"run.different_failure","message":"ignored"}`),
		Now:       fixture.Now.Add(4 * time.Second),
	})
	require.NoError(t, err)
	require.False(t, second.Draining)
	require.Equal(t, run.RunStatusFailed, second.Run.Status)
	require.Equal(t, "run.plan_version_invalid", second.Run.ErrorCode)
	require.NotNil(t, second.Run.CompletedAt)
	require.Nil(t, second.Event, "failure drain finalization must not duplicate run_failed")

	completed, err := store.LoadEffect(ctx, fixture.WorkspaceID, fixture.RunID, effectID)
	require.NoError(t, err)
	require.Equal(t, run.EffectStatusCompleted, completed.Status)
	events, err := store.ListEvents(ctx, run.ListEventsParams{WorkspaceID: fixture.WorkspaceID, RunID: fixture.RunID, Limit: 20})
	require.NoError(t, err)
	require.Equal(t, []string{"run_accepted", "run_failed"}, []string{events[0].Type, events[1].Type})
}

func TestFailClaimedRunPreservesCrossedDispatchRetryUntilDefinitiveReceipt(t *testing.T) {
	ctx := context.Background()
	db := startRollbackDB(t, ctx)
	fixture := seedRunFixture(t, ctx, db, "failure_dispatch_retry")
	store, _, effectID := seedFailureChildren(t, ctx, db, fixture)

	candidates, err := store.ListPendingEffects(ctx, run.ListPendingEffectsParams{
		Now: fixture.Now.Add(time.Second), Limit: 10,
	})
	require.NoError(t, err)
	require.Len(t, candidates, 1)
	dispatching, ok, err := store.ClaimEffect(ctx, run.ClaimEffectParams{
		Candidate: candidates[0], WorkerID: "effect-worker", Now: fixture.Now.Add(time.Second), LeaseTTL: time.Minute,
	})
	require.NoError(t, err)
	require.True(t, ok)
	require.EqualValues(t, 1, dispatching.AttemptCount)

	claimed, err := store.ClaimRunnableRuns(ctx, run.ClaimRunnableParams{
		WorkerID: "graph-failure", Now: fixture.Now.Add(2 * time.Second), LeaseTTL: time.Minute, Limit: 1,
	})
	require.NoError(t, err)
	require.Len(t, claimed, 1)
	first, err := store.FailClaimedRun(ctx, run.FailClaimedRunParams{
		Fence: claimed[0].Fence, ExpectedRevision: claimed[0].Run.Revision,
		StateJSON: deterministicFailureState(), ErrorCode: "run.plan_version_invalid",
		ErrorJSON: json.RawMessage(`{"code":"run.plan_version_invalid","message":"invalid artifact"}`),
		Now:       fixture.Now.Add(2 * time.Second),
	})
	require.NoError(t, err)
	require.True(t, first.Draining)
	require.Equal(t, run.RunStatusWaiting, first.Run.Status)

	// The writer may have committed externally and then lost its response. A
	// retryable transport error therefore returns the Effect to pending with a
	// stable idempotency key; deterministic Run failure must not cancel it.
	retryAt := fixture.Now.Add(5 * time.Second)
	resolvedPending, _, err := store.ResolveEffect(ctx, run.ResolveEffectParams{
		WorkspaceID: fixture.WorkspaceID, RunID: fixture.RunID, EffectID: effectID,
		LeaseToken: dispatching.LeaseToken, Status: run.EffectStatusPending,
		ErrorCode: "effect.transport_error", ErrorJSON: json.RawMessage(`{"message":"response lost"}`),
		RetryAt: &retryAt, RunnableAt: retryAt, Now: fixture.Now.Add(3 * time.Second),
	})
	require.NoError(t, err)
	require.Equal(t, run.EffectStatusPending, resolvedPending.Status)
	require.EqualValues(t, 1, resolvedPending.AttemptCount)
	require.Equal(t, retryAt, *resolvedPending.RetryAt)
	runWaitingForRetry, err := store.LoadRun(ctx, fixture.WorkspaceID, fixture.RunID)
	require.NoError(t, err)
	require.Equal(t, retryAt, runWaitingForRetry.RunnableAt, "pending retries retain their retry schedule while failure drains")

	claimed, err = store.ClaimRunnableRuns(ctx, run.ClaimRunnableParams{
		WorkerID: "graph-failure", Now: retryAt, LeaseTTL: time.Minute, Limit: 1,
	})
	require.NoError(t, err)
	require.Len(t, claimed, 1)
	second, err := store.FailClaimedRun(ctx, run.FailClaimedRunParams{
		Fence: claimed[0].Fence, ExpectedRevision: claimed[0].Run.Revision,
		StateJSON: deterministicFailureState(), ErrorCode: "run.different_failure",
		ErrorJSON: json.RawMessage(`{"code":"run.different_failure","message":"ignored"}`),
		Now:       retryAt,
	})
	require.NoError(t, err)
	require.True(t, second.Draining)
	require.Equal(t, run.RunStatusWaiting, second.Run.Status)
	require.Equal(t, retryAt, second.Run.RunnableAt)
	stillPending, err := store.LoadEffect(ctx, fixture.WorkspaceID, fixture.RunID, effectID)
	require.NoError(t, err)
	require.Equal(t, run.EffectStatusPending, stillPending.Status)
	require.EqualValues(t, 1, stillPending.AttemptCount)

	candidates, err = store.ListPendingEffects(ctx, run.ListPendingEffectsParams{Now: retryAt, Limit: 10})
	require.NoError(t, err)
	require.Len(t, candidates, 1)
	replayed, ok, err := store.ClaimEffect(ctx, run.ClaimEffectParams{
		Candidate: candidates[0], WorkerID: "effect-worker", Now: retryAt, LeaseTTL: time.Minute,
	})
	require.NoError(t, err)
	require.True(t, ok)
	require.Equal(t, dispatching.IdempotencyKey, replayed.IdempotencyKey)
	require.EqualValues(t, 2, replayed.AttemptCount)
	definitiveAt := retryAt.Add(time.Second)
	_, _, err = store.CompleteEffect(
		ctx, fixture.WorkspaceID, fixture.RunID, effectID, replayed.LeaseToken,
		json.RawMessage(`{"port":"sent","output":{"messageId":"msg_replayed"}}`), nil, definitiveAt,
	)
	require.NoError(t, err)

	claimed, err = store.ClaimRunnableRuns(ctx, run.ClaimRunnableParams{
		WorkerID: "graph-failure", Now: definitiveAt, LeaseTTL: time.Minute, Limit: 1,
	})
	require.NoError(t, err)
	require.Len(t, claimed, 1)
	finalized, err := store.FailClaimedRun(ctx, run.FailClaimedRunParams{
		Fence: claimed[0].Fence, ExpectedRevision: claimed[0].Run.Revision,
		StateJSON: deterministicFailureState(), ErrorCode: "run.different_failure",
		ErrorJSON: json.RawMessage(`{"code":"run.different_failure","message":"ignored"}`),
		Now:       definitiveAt,
	})
	require.NoError(t, err)
	require.False(t, finalized.Draining)
	require.Equal(t, run.RunStatusFailed, finalized.Run.Status)
	require.Equal(t, "run.plan_version_invalid", finalized.Run.ErrorCode)
	require.Nil(t, finalized.Event, "re-entering failure drain must not duplicate run_failed")
	completed, err := store.LoadEffect(ctx, fixture.WorkspaceID, fixture.RunID, effectID)
	require.NoError(t, err)
	require.Equal(t, run.EffectStatusCompleted, completed.Status)
	require.EqualValues(t, 2, completed.AttemptCount)
	events, err := store.ListEvents(ctx, run.ListEventsParams{WorkspaceID: fixture.WorkspaceID, RunID: fixture.RunID, Limit: 20})
	require.NoError(t, err)
	require.Equal(t, []string{"run_accepted", "run_failed"}, []string{events[0].Type, events[1].Type})
}

func TestResolveEffectWakesDeterministicallyFailingRunImmediatelyWhenAwaiting(t *testing.T) {
	ctx := context.Background()
	db := startRollbackDB(t, ctx)
	fixture := seedRunFixture(t, ctx, db, "failure_awaiting_wake")
	store, _, effectID := seedFailureChildren(t, ctx, db, fixture)

	candidates, err := store.ListPendingEffects(ctx, run.ListPendingEffectsParams{
		Now: fixture.Now.Add(time.Second), Limit: 10,
	})
	require.NoError(t, err)
	require.Len(t, candidates, 1)
	dispatching, ok, err := store.ClaimEffect(ctx, run.ClaimEffectParams{
		Candidate: candidates[0], WorkerID: "approval-writer", Now: fixture.Now.Add(time.Second), LeaseTTL: time.Minute,
	})
	require.NoError(t, err)
	require.True(t, ok)

	claimed, err := store.ClaimRunnableRuns(ctx, run.ClaimRunnableParams{
		WorkerID: "graph-failure", Now: fixture.Now.Add(2 * time.Second), LeaseTTL: time.Minute, Limit: 1,
	})
	require.NoError(t, err)
	require.Len(t, claimed, 1)
	first, err := store.FailClaimedRun(ctx, run.FailClaimedRunParams{
		Fence: claimed[0].Fence, ExpectedRevision: claimed[0].Run.Revision,
		StateJSON: deterministicFailureState(), ErrorCode: "run.plan_version_invalid",
		ErrorJSON: json.RawMessage(`{"code":"run.plan_version_invalid","message":"invalid artifact"}`),
		Now:       fixture.Now.Add(2 * time.Second),
	})
	require.NoError(t, err)
	require.True(t, first.Draining)

	resolvedAt := fixture.Now.Add(3 * time.Second)
	authoredDeadline := fixture.Now.Add(24 * time.Hour)
	awaiting, _, err := store.ResolveEffect(ctx, run.ResolveEffectParams{
		WorkspaceID: fixture.WorkspaceID, RunID: fixture.RunID, EffectID: effectID,
		LeaseToken: dispatching.LeaseToken, Status: run.EffectStatusAwaiting,
		ResultJSON: json.RawMessage(`{"decision":"approved"}`), DeliveredAt: &resolvedAt,
		RunnableAt: authoredDeadline, Now: resolvedAt,
	})
	require.NoError(t, err)
	require.Equal(t, run.EffectStatusAwaiting, awaiting.Status)
	woken, err := store.LoadRun(ctx, fixture.WorkspaceID, fixture.RunID)
	require.NoError(t, err)
	require.Equal(t, resolvedAt, woken.RunnableAt, "failure drain must not wait for the authored callback deadline")

	claimed, err = store.ClaimRunnableRuns(ctx, run.ClaimRunnableParams{
		WorkerID: "graph-failure", Now: resolvedAt, LeaseTTL: time.Minute, Limit: 1,
	})
	require.NoError(t, err)
	require.Len(t, claimed, 1)
	finalized, err := store.FailClaimedRun(ctx, run.FailClaimedRunParams{
		Fence: claimed[0].Fence, ExpectedRevision: claimed[0].Run.Revision,
		StateJSON: deterministicFailureState(), ErrorCode: "run.different_failure",
		ErrorJSON: json.RawMessage(`{"code":"run.different_failure","message":"ignored"}`),
		Now:       resolvedAt,
	})
	require.NoError(t, err)
	require.False(t, finalized.Draining)
	require.Equal(t, run.RunStatusFailed, finalized.Run.Status)
	canceled, err := store.LoadEffect(ctx, fixture.WorkspaceID, fixture.RunID, effectID)
	require.NoError(t, err)
	require.Equal(t, run.EffectStatusCanceled, canceled.Status)
}

func seedFailureChildren(t *testing.T, ctx context.Context, db coredb.DB, fixture runFixture) (*run.Store, string, string) {
	t.Helper()
	store := run.NewWithOptions(db, run.StoreOptions{WorkspaceGroupID: fixture.GroupID})
	claimed, err := store.ClaimRunnableRuns(ctx, run.ClaimRunnableParams{
		WorkerID: "graph", Now: fixture.Now, LeaseTTL: time.Minute, Limit: 1,
	})
	require.NoError(t, err)
	require.Len(t, claimed, 1)
	activationID := "wra_" + fixture.RunID
	effectID := "wfx_" + fixture.RunID
	retryAt := fixture.Now.Add(time.Second)
	_, err = store.SaveTurn(ctx, run.SaveTurnParams{
		Fence: claimed[0].Fence, ExpectedRevision: claimed[0].Run.Revision,
		Status: run.RunStatusWaiting, StateJSON: json.RawMessage(`{"status":"waiting"}`),
		RunnableAt: fixture.Now.Add(2 * time.Second), StartedAt: claimed[0].Run.StartedAt,
		ReleaseLease: true, Now: fixture.Now,
		Activations: []run.ActivationMutation{{Operation: run.MutationInsert, Value: run.Activation{
			ID: activationID, LogicalKey: "approval/1", NodeID: "approval", Opcode: run.EffectKindApproval,
			Status: run.ActivationStatusWaiting, InputJSON: json.RawMessage(`{}`), ScopeContextJSON: json.RawMessage(`{}`),
		}}},
		Effects: []run.EffectMutation{{Operation: run.MutationInsert, Value: run.Effect{
			ID: effectID, ActivationID: activationID, Kind: run.EffectKindApproval, IdempotencyKey: "approval/1",
			Status: run.EffectStatusPending, RequestJSON: json.RawMessage(`{"summary":"work?"}`), RetryAt: &retryAt,
		}}},
	})
	require.NoError(t, err)
	return store, activationID, effectID
}

func deterministicFailureState() json.RawMessage {
	return json.RawMessage(`{
		"status":"failed",
		"input":{"seed":"start"},
		"variables":{},
		"failure":{"code":"run.plan_version_invalid","message":"invalid artifact"},
		"activations":[{"id":"a1","sequence":1,"nodeId":"approval","inputPort":"in","input":{},"selectedPorts":[],"status":"canceled","lineage":[],"ownedScopeIds":[],"failure":{"code":"run.plan_version_invalid"}}],
		"scopes":[],"nextActivation":1,"nextScope":0
	}`)
}