package invocation_test

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

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/store/workflow/invocation"
	"github.com/tidbcloud/AgentTeam/internal/workflow/binding"
	invocationlang "github.com/tidbcloud/AgentTeam/internal/workflow/invocation"
	"github.com/tidbcloud/AgentTeam/internal/workflow/plan"
)

func TestScheduledAcceptanceAdvancesFromOccurrenceAndReplaysAfterLifecycleChanges(t *testing.T) {
	ctx := context.Background()
	db := startRollbackDB(t, ctx)
	fixture := seedInvocationFixture(t, ctx, db, "schedule_recurring")
	groupID := "group-invocation-recurring"
	occurrence := fixture.Now.Add(-5 * time.Minute)
	interval := int64(60)
	configureScheduleInvocationFixture(t, ctx, db, fixture, groupID, binding.ScheduleSource{
		Mode: binding.ScheduleModeInterval, AnchorAt: &occurrence, IntervalSeconds: &interval,
		Inputs: json.RawMessage(`{"title":"scheduled"}`),
	})

	store := invocation.NewWithOptions(db, invocation.StoreOptions{WorkspaceGroupID: groupID})
	candidates, err := store.ListDueScheduleCandidates(ctx, invocationlang.ListDueScheduleInput{Now: fixture.Now, Limit: 10})
	require.NoError(t, err)
	require.Len(t, candidates, 1)
	candidate := candidates[0]
	require.Equal(t, fixture.WorkspaceID, candidate.WorkspaceID())
	require.Equal(t, fixture.SourceID, candidate.SourceID())
	require.Equal(t, occurrence, candidate.OccurrenceAt())

	created, err := store.MaterializeScheduled(ctx, candidate, fixture.Now)
	require.NoError(t, err)
	require.Equal(t, invocationlang.ScheduleMaterializeCreated, created.Disposition)
	require.False(t, created.Acceptance.Replayed)
	require.Empty(t, created.Acceptance.Invocation.InvokedBy)
	require.JSONEq(t, `{"title":"scheduled"}`, string(created.Acceptance.Invocation.InputJSON))
	require.Equal(t, occurrence, created.Acceptance.Invocation.SourceOccurredAt)

	var next time.Time
	var status string
	var sourceRevision int64
	require.NoError(t, db.QueryRow(ctx, `SELECT status, revision, next_occurrence_at FROM workflow_binding_sources WHERE id = $1`, fixture.SourceID).Scan(&status, &sourceRevision, &next))
	require.Equal(t, "active", status)
	require.Equal(t, int64(2), sourceRevision)
	require.Equal(t, occurrence.Add(time.Minute), next.UTC(), "backlog cursor must advance from the accepted occurrence, not dispatcher now")

	// The original candidate is now stale with respect to the mutable cursor,
	// but occurrence identity wins first and returns the immutable receipt.
	replayed, err := store.MaterializeScheduled(ctx, candidate, fixture.Now.Add(time.Minute))
	require.NoError(t, err)
	require.Equal(t, invocationlang.ScheduleMaterializeReplayed, replayed.Disposition)
	require.True(t, replayed.Acceptance.Replayed)
	require.Equal(t, created.Acceptance.Invocation.ID, replayed.Acceptance.Invocation.ID)
	require.Equal(t, created.Acceptance.Run.ID, replayed.Acceptance.Run.ID)

	closedAt := fixture.Now.Add(2 * time.Minute)
	mustExec(t, ctx, db, `
UPDATE workflow_binding_sources
SET status='disabled', closed_at=$1, closed_reason='disabled', revision=revision+1, updated_at=$1
WHERE id=$2
`, closedAt, fixture.SourceID)
	mustExec(t, ctx, db, `
UPDATE workflow_bindings
SET state='disabled', active_revision_id=NULL, active_source_id=NULL, revision=revision+1, updated_at=$1
WHERE id=$2
`, closedAt, fixture.BindingID)
	mustExec(t, ctx, db, `UPDATE members SET status='inactive', updated_at=$1 WHERE id=$2`, closedAt, fixture.MemberID)

	replayedAfterLifecycle, err := store.AcceptScheduled(ctx, candidate, closedAt)
	require.NoError(t, err)
	require.Equal(t, invocationlang.ScheduleMaterializeReplayed, replayedAfterLifecycle.Disposition)
	require.True(t, replayedAfterLifecycle.Acceptance.Replayed)
	require.Equal(t, created.Acceptance.Invocation.ID, replayedAfterLifecycle.Acceptance.Invocation.ID)

	var invocationCount, runCount int
	require.NoError(t, db.QueryRow(ctx, `SELECT COUNT(*) FROM workflow_invocations WHERE workspace_id=$1 AND binding_source_id=$2`, fixture.WorkspaceID, fixture.SourceID).Scan(&invocationCount))
	require.NoError(t, db.QueryRow(ctx, `
SELECT COUNT(*)
FROM workflow_run_instances run
JOIN workflow_invocations invocation
  ON invocation.workspace_id=run.workspace_id AND invocation.id=run.invocation_id
WHERE invocation.workspace_id=$1 AND invocation.binding_id=$2
`, fixture.WorkspaceID, fixture.BindingID).Scan(&runCount))
	require.Equal(t, 1, invocationCount)
	require.Equal(t, 1, runCount)
}

func TestScheduledOneShotCompletesSourceInAcceptanceTransaction(t *testing.T) {
	ctx := context.Background()
	db := startRollbackDB(t, ctx)
	fixture := seedInvocationFixture(t, ctx, db, "schedule_one_shot")
	groupID := "group-invocation-one-shot"
	occurrence := fixture.Now.Add(-time.Second)
	configureScheduleInvocationFixture(t, ctx, db, fixture, groupID, binding.ScheduleSource{
		Mode: binding.ScheduleModeAbsolute, ScheduledAt: &occurrence,
		Inputs: json.RawMessage(`{"title":"once"}`),
	})
	// Plan archive is authoring lifecycle only. The activated Binding remains
	// pinned to this exact immutable PlanVersion and must keep running.
	mustExec(t, ctx, db, `
UPDATE workflow_plans
SET archived_by_member_id=$1, archived_by_actor_type='human', archived_by_actor_id=$2,
    archived_at=$3, revision=revision+1, updated_at=$3
WHERE workspace_id=$4 AND id=$5
`, fixture.MemberID, fixture.UserID, fixture.Now, fixture.WorkspaceID, fixture.PlanID)

	store := invocation.NewWithOptions(db, invocation.StoreOptions{WorkspaceGroupID: groupID})
	candidates, err := store.ListDueScheduleCandidates(ctx, invocationlang.ListDueScheduleInput{Now: fixture.Now, Limit: 10})
	require.NoError(t, err)
	require.Len(t, candidates, 1)
	result, err := store.AcceptScheduled(ctx, candidates[0], fixture.Now)
	require.NoError(t, err)
	require.Equal(t, invocationlang.ScheduleMaterializeCreated, result.Disposition)

	var status string
	var next *time.Time
	var closedAt *time.Time
	var reason string
	require.NoError(t, db.QueryRow(ctx, `
SELECT status, next_occurrence_at, closed_at, closed_reason
FROM workflow_binding_sources WHERE id=$1
`, fixture.SourceID).Scan(&status, &next, &closedAt, &reason))
	require.Equal(t, "completed", status)
	require.Nil(t, next)
	require.NotNil(t, closedAt)
	require.Equal(t, fixture.Now, closedAt.UTC())
	require.Equal(t, "schedule_completed", reason)

	candidates, err = store.ListDueScheduleCandidates(ctx, invocationlang.ListDueScheduleInput{Now: fixture.Now.Add(time.Hour), Limit: 10})
	require.NoError(t, err)
	require.Empty(t, candidates)
}

func TestScheduledCandidateIsFencedByPlacementGeneration(t *testing.T) {
	ctx := context.Background()
	db := startRollbackDB(t, ctx)
	fixture := seedInvocationFixture(t, ctx, db, "schedule_fenced")
	groupID := "group-invocation-fenced"
	occurrence := fixture.Now.Add(-time.Second)
	configureScheduleInvocationFixture(t, ctx, db, fixture, groupID, binding.ScheduleSource{
		Mode: binding.ScheduleModeAbsolute, ScheduledAt: &occurrence,
		Inputs: json.RawMessage(`{"title":"fenced"}`),
	})
	store := invocation.NewWithOptions(db, invocation.StoreOptions{WorkspaceGroupID: groupID})
	candidates, err := store.ListDueScheduleCandidates(ctx, invocationlang.ListDueScheduleInput{Now: fixture.Now, Limit: 10})
	require.NoError(t, err)
	require.Len(t, candidates, 1)

	mustExec(t, ctx, db, `UPDATE workspace_placements SET placement_epoch=placement_epoch+1, updated_at=$1 WHERE workspace_id=$2`, fixture.Now, fixture.WorkspaceID)
	result, err := store.MaterializeScheduled(ctx, candidates[0], fixture.Now)
	require.NoError(t, err)
	require.Equal(t, invocationlang.ScheduleMaterializeFenced, result.Disposition)

	var invocationCount int
	require.NoError(t, db.QueryRow(ctx, `SELECT COUNT(*) FROM workflow_invocations WHERE workspace_id=$1`, fixture.WorkspaceID).Scan(&invocationCount))
	require.Zero(t, invocationCount)
	var persistedOccurrence time.Time
	require.NoError(t, db.QueryRow(ctx, `SELECT next_occurrence_at FROM workflow_binding_sources WHERE id=$1`, fixture.SourceID).Scan(&persistedOccurrence))
	require.Equal(t, occurrence, persistedOccurrence.UTC())
}

func TestScheduledFreshOccurrenceBlocksWhenOwnerIsUnavailable(t *testing.T) {
	ctx := context.Background()
	db := startRollbackDB(t, ctx)
	fixture := seedInvocationFixture(t, ctx, db, "schedule_owner_blocked")
	groupID := "group-invocation-owner-blocked"
	occurrence := fixture.Now.Add(-time.Second)
	configureScheduleInvocationFixture(t, ctx, db, fixture, groupID, binding.ScheduleSource{
		Mode: binding.ScheduleModeAbsolute, ScheduledAt: &occurrence,
		Inputs: json.RawMessage(`{"title":"blocked"}`),
	})
	store := invocation.NewWithOptions(db, invocation.StoreOptions{WorkspaceGroupID: groupID})
	candidates, err := store.ListDueScheduleCandidates(ctx, invocationlang.ListDueScheduleInput{Now: fixture.Now, Limit: 10})
	require.NoError(t, err)
	require.Len(t, candidates, 1)
	mustExec(t, ctx, db, `UPDATE members SET status='inactive', updated_at=$1 WHERE id=$2`, fixture.Now, fixture.MemberID)

	result, err := store.AcceptScheduled(ctx, candidates[0], fixture.Now)
	require.NoError(t, err)
	require.Equal(t, invocationlang.ScheduleMaterializeBlocked, result.Disposition)
	require.Empty(t, result.Acceptance.Invocation.ID)

	var invocationCount int
	require.NoError(t, db.QueryRow(ctx, `SELECT COUNT(*) FROM workflow_invocations WHERE workspace_id=$1`, fixture.WorkspaceID).Scan(&invocationCount))
	require.Zero(t, invocationCount)
	var persistedOccurrence time.Time
	require.NoError(t, db.QueryRow(ctx, `SELECT next_occurrence_at FROM workflow_binding_sources WHERE id=$1`, fixture.SourceID).Scan(&persistedOccurrence))
	require.Equal(t, occurrence, persistedOccurrence.UTC())
}

func TestBlockedScheduleBatchBacksOffSoHealthySourceIsNotStarved(t *testing.T) {
	ctx := context.Background()
	db := startRollbackDB(t, ctx)
	groupID := "group-invocation-blocked-fairness"
	now := time.Date(2026, 7, 21, 8, 0, 0, 0, time.UTC)

	blockedOne := seedInvocationFixture(t, ctx, db, "schedule_blocked_fairness_one")
	blockedTwo := seedInvocationFixture(t, ctx, db, "schedule_blocked_fairness_two")
	healthy := seedInvocationFixture(t, ctx, db, "schedule_blocked_fairness_healthy")
	fixtures := []struct {
		fixture    invocationFixture
		occurrence time.Time
		blocked    bool
	}{
		{fixture: blockedOne, occurrence: now.Add(-3 * time.Minute), blocked: true},
		{fixture: blockedTwo, occurrence: now.Add(-2 * time.Minute), blocked: true},
		{fixture: healthy, occurrence: now.Add(-time.Minute)},
	}
	for _, item := range fixtures {
		configureScheduleInvocationFixture(t, ctx, db, item.fixture, groupID, binding.ScheduleSource{
			Mode: binding.ScheduleModeAbsolute, ScheduledAt: &item.occurrence,
			Inputs: json.RawMessage(`{"title":"fair"}`),
		})
		if item.blocked {
			mustExec(t, ctx, db, `UPDATE members SET status='inactive', updated_at=$1 WHERE id=$2`, now, item.fixture.MemberID)
		}
	}

	store := invocation.NewWithOptions(db, invocation.StoreOptions{WorkspaceGroupID: groupID})
	firstBatch, err := store.ListDueScheduleCandidates(ctx, invocationlang.ListDueScheduleInput{Now: now, Limit: 2})
	require.NoError(t, err)
	require.Len(t, firstBatch, 2)
	require.Equal(t, blockedOne.SourceID, firstBatch[0].SourceID())
	require.Equal(t, blockedTwo.SourceID, firstBatch[1].SourceID())
	for _, candidate := range firstBatch {
		result, err := store.MaterializeScheduled(ctx, candidate, now)
		require.NoError(t, err)
		require.Equal(t, invocationlang.ScheduleMaterializeBlocked, result.Disposition)
	}

	// At the same dispatcher instant, both blocked rows are now durably fenced
	// out. The next bounded scan can reach the healthy source that was ordered
	// behind a completely blocked first batch.
	secondBatch, err := store.ListDueScheduleCandidates(ctx, invocationlang.ListDueScheduleInput{Now: now, Limit: 2})
	require.NoError(t, err)
	require.Len(t, secondBatch, 1)
	require.Equal(t, healthy.SourceID, secondBatch[0].SourceID())
	created, err := store.MaterializeScheduled(ctx, secondBatch[0], now)
	require.NoError(t, err)
	require.Equal(t, invocationlang.ScheduleMaterializeCreated, created.Disposition)

	// Even when the deferred prefix becomes eligible again, a newly visible
	// never-deferred source must sort ahead of it. This is the property that
	// prevents a blocked prefix larger than batch*backoff/poll from cycling at
	// the head of every bounded scan forever.
	lateHealthy := seedInvocationFixture(t, ctx, db, "schedule_blocked_fairness_late_healthy")
	lateOccurrence := now.Add(-30 * time.Second)
	configureScheduleInvocationFixture(t, ctx, db, lateHealthy, groupID, binding.ScheduleSource{
		Mode: binding.ScheduleModeAbsolute, ScheduledAt: &lateOccurrence,
		Inputs: json.RawMessage(`{"title":"late-fair"}`),
	})

	for _, fixture := range []invocationFixture{blockedOne, blockedTwo} {
		var retryAt time.Time
		require.NoError(t, db.QueryRow(ctx, `SELECT materialization_retry_at FROM workflow_binding_sources WHERE id=$1`, fixture.SourceID).Scan(&retryAt))
		require.Equal(t, now.Add(30*time.Second), retryAt.UTC())
	}
	retryAt := now.Add(30 * time.Second)
	retryBatch, err := store.ListDueScheduleCandidates(ctx, invocationlang.ListDueScheduleInput{Now: retryAt, Limit: 2})
	require.NoError(t, err)
	require.Len(t, retryBatch, 2)
	require.Equal(t, lateHealthy.SourceID, retryBatch[0].SourceID(), "never-deferred due work must precede eligible blocked retries")
	lateCreated, err := store.MaterializeScheduled(ctx, retryBatch[0], retryAt)
	require.NoError(t, err)
	require.Equal(t, invocationlang.ScheduleMaterializeCreated, lateCreated.Disposition)

	retryBatch, err = store.ListDueScheduleCandidates(ctx, invocationlang.ListDueScheduleInput{Now: retryAt, Limit: 2})
	require.NoError(t, err)
	require.Len(t, retryBatch, 2, "blocked occurrences must remain eligible after never-deferred work drains")
	require.ElementsMatch(t, []string{blockedOne.SourceID, blockedTwo.SourceID}, []string{retryBatch[0].SourceID(), retryBatch[1].SourceID()})
}

func configureScheduleInvocationFixture(
	t *testing.T,
	ctx context.Context,
	db coredb.QueryExecer,
	fixture invocationFixture,
	groupID string,
	schedule binding.ScheduleSource,
) {
	t.Helper()
	artifact, diagnostics, err := plan.Compile([]byte(scheduleInvocationPlanSource))
	require.NoError(t, err)
	require.False(t, plan.HasErrors(diagnostics), diagnostics)
	require.NotNil(t, artifact)
	artifactJSON, err := artifact.DeterministicJSON()
	require.NoError(t, err)

	normalized, err := binding.NormalizeRevision(binding.RevisionSpec{
		PlanID: fixture.PlanID, PlanVersionID: fixture.PlanVersionID,
		InputPolicy: binding.InputPolicy{Defaults: json.RawMessage(`{}`), Constants: json.RawMessage(`{}`)},
		Source:      binding.Source{Type: binding.SourceTypeSchedule, Schedule: &schedule},
	}, artifact.Contracts.Input)
	require.NoError(t, err)
	initial, err := binding.InitialOccurrence(*normalized.Spec.Source.Schedule, fixture.Now.Add(-time.Hour))
	require.NoError(t, err)
	require.NotNil(t, initial)
	// Tests author explicit due instants. Interval's anchor and absolute's
	// scheduledAt make InitialOccurrence identical to the desired first cursor
	// whenever activation precedes it.
	wantOccurrence := schedule.ScheduledAt
	if wantOccurrence == nil {
		wantOccurrence = schedule.AnchorAt
	}
	require.NotNil(t, wantOccurrence)
	require.Equal(t, wantOccurrence.UTC().Truncate(time.Microsecond), initial.UTC())

	mustExec(t, ctx, db, `
UPDATE workflow_plan_versions
SET artifact=$1, compiler_version=$2, digest=$3
WHERE workspace_id=$4 AND plan_id=$5 AND id=$6
`, artifactJSON, artifact.CompilerVersion, artifact.Digest, fixture.WorkspaceID, fixture.PlanID, fixture.PlanVersionID)
	mustExec(t, ctx, db, `
UPDATE workflow_binding_revisions
SET source_kind='schedule', config_json=$1, digest=$2
WHERE workspace_id=$3 AND binding_id=$4 AND id=$5
`, normalized.ConfigJSON, normalized.Digest, fixture.WorkspaceID, fixture.BindingID, fixture.RevisionID)
	mustExec(t, ctx, db, `
UPDATE workflow_binding_sources
SET source_kind='schedule', next_occurrence_at=$1
WHERE workspace_id=$2 AND binding_id=$3 AND id=$4
`, initial, fixture.WorkspaceID, fixture.BindingID, fixture.SourceID)
	mustExec(t, ctx, db, `
INSERT IGNORE INTO workspace_groups (
  workspace_group_id, display_name, state, resource_profile, service_ref, created_at, updated_at
) VALUES ($1,$1,'active','default',$1,$2,$2)
`, groupID, fixture.Now)
	mustExec(t, ctx, db, `
INSERT INTO workspace_placements (
  workspace_id, workspace_group_id, state, placement_epoch, created_at, updated_at
) VALUES ($1,$2,'active',1,$3,$3)
ON DUPLICATE KEY UPDATE
  workspace_group_id=VALUES(workspace_group_id),
  state=VALUES(state),
  placement_epoch=VALUES(placement_epoch),
  updated_at=VALUES(updated_at)
`, fixture.WorkspaceID, groupID, fixture.Now)
}

const scheduleInvocationPlanSource = `apiVersion: loop.plan/v1
kind: PlanVersion
input:
  title: {type: string, required: true}
graph:
  entry: {nodeId: pause, port: in}
  nodes:
  - id: pause
    type: wait
    wait:
      until: {value: PT1M}
      timeout: PT2M
      retry: {maxAttempts: 2, backoff: 1s}
  - id: done
    type: end
    end: {outcome: succeeded}
  - id: failed
    type: end
    end: {outcome: failed, code: wait_failed}
  edges:
  - id: completed
    source: {nodeId: pause, port: completed}
    target: {nodeId: done, port: in}
  - id: timed_out
    source: {nodeId: pause, port: timed_out}
    target: {nodeId: failed, port: in}
  - id: action_failed
    source: {nodeId: pause, port: failed}
    target: {nodeId: failed, port: in}`