package agentops

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

	"github.com/stretchr/testify/require"

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

func TestValidateEnqueueInput(t *testing.T) {
	t.Parallel()

	valid := EnqueueInput{
		WorkspaceID: "w", ApplicantAgentID: "a", OwnerUserID: "u",
		Action: "create", ApprovalID: "ap", RequestID: "r",
		Payload: json.RawMessage(`{}`), Items: []json.RawMessage{json.RawMessage(`{}`)},
	}
	tests := []struct {
		name    string
		mutate  func(EnqueueInput) EnqueueInput
		wantErr bool
	}{
		{"valid", func(e EnqueueInput) EnqueueInput { return e }, false},
		{"missing workspace", func(e EnqueueInput) EnqueueInput { e.WorkspaceID = ""; return e }, true},
		{"missing approval", func(e EnqueueInput) EnqueueInput { e.ApprovalID = ""; return e }, true},
		{"invalid action", func(e EnqueueInput) EnqueueInput { e.Action = "explode"; return e }, true},
		{"no items", func(e EnqueueInput) EnqueueInput { e.Items = nil; return e }, true},
		{"invalid payload JSON", func(e EnqueueInput) EnqueueInput { e.Payload = json.RawMessage(`bad`); return e }, true},
		{"invalid item JSON", func(e EnqueueInput) EnqueueInput { e.Items = []json.RawMessage{json.RawMessage(`bad`)}; return e }, true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			err := validateEnqueueInput(tt.mutate(valid))
			if tt.wantErr {
				require.Error(t, err)
			} else {
				require.NoError(t, err)
			}
		})
	}
}

func TestOperationEnqueueClaimComplete(t *testing.T) {
	ctx := context.Background()
	store := startAgentOpsConcurrencyStore(t, ctx)

	// Enqueue operation + 2 items in a tx.
	var opID int64
	err := coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		var err error
		opID, err = store.InTx(tx).EnqueueOperationInTx(ctx, EnqueueInput{
			WorkspaceID: "wsp-1", ApplicantAgentID: "agt-caller", OwnerUserID: "usr-owner",
			Action: "duplicate", Payload: json.RawMessage(`{"source":"agt_src","count":2}`),
			ApprovalID: "appr-1", RequestID: "req-1",
			Items: []json.RawMessage{json.RawMessage(`{"index":0}`), json.RawMessage(`{"index":1}`)},
		})
		return err
	})
	require.NoError(t, err)
	require.NotZero(t, opID)

	// Claim the queued operation.
	op, err := store.ClaimNextOperation(ctx, "")
	require.NoError(t, err)
	require.NotNil(t, op)
	require.Equal(t, "running", op.Status)

	// No more queued.
	op2, err := store.ClaimNextOperation(ctx, "")
	require.NoError(t, err)
	require.Nil(t, op2)

	// Load pending items in strict index order.
	items, err := store.LoadItems(ctx, opID, true)
	require.NoError(t, err)
	require.Len(t, items, 2)
	require.Equal(t, 0, items[0].ItemIndex)
	require.Equal(t, 1, items[1].ItemIndex)

	// Complete item 0 (applied) + item 1 (failed).
	require.NoError(t, store.CompleteItem(ctx, opID, 0, "applied", json.RawMessage(`{"agentId":"agt_new1"}`), ""))
	require.NoError(t, store.CompleteItem(ctx, opID, 1, "failed", nil, "name conflict"))
	allItems, err := store.LoadItems(ctx, opID, false)
	require.NoError(t, err)
	require.Len(t, allItems, 2)
	require.Equal(t, []string{"applied", "failed"}, []string{allItems[0].Status, allItems[1].Status})

	// Invalid item status rejected.
	require.Error(t, store.CompleteItem(ctx, opID, 0, "skipped", nil, ""))

	// CAS miss: item 0 is no longer pending.
	require.Error(t, store.CompleteItem(ctx, opID, 0, "applied", nil, ""))

	// Complete operation as partial (1 applied + 1 failed).
	require.NoError(t, store.CompleteOperation(ctx, opID, "partial"))

	// CAS miss: operation is no longer running.
	require.Error(t, store.CompleteOperation(ctx, opID, "applied"))

	// Invalid terminal status rejected.
	require.Error(t, store.CompleteOperation(ctx, opID, "queued"))
}

func TestClaimNextQueuedRunningEligibility(t *testing.T) {
	ctx := context.Background()
	store := startAgentOpsConcurrencyStore(t, ctx)
	pendingID := enqueueTestOperation(t, ctx, store, "wsp-pending", "pending")
	readyID := enqueueTestOperation(t, ctx, store, "wsp-ready", "ready")
	mustSetOperationRunning(t, ctx, store, pendingID)
	mustSetOperationRunning(t, ctx, store, readyID)
	require.NoError(t, store.CompleteItem(ctx, readyID, 0, "applied", json.RawMessage(`{"agentId":"agt-ready"}`), ""))

	claimed, err := store.ClaimNextOperation(ctx, "")
	require.NoError(t, err)
	require.NotNil(t, claimed)
	require.Equal(t, readyID, claimed.ID)
	require.Equal(t, "running", claimed.Status)

	claimed, err = store.ClaimNextOperation(ctx, "")
	require.NoError(t, err)
	require.NotNil(t, claimed)
	require.Equal(t, readyID, claimed.ID,
		"finalization-ready remains retryable until operation completion")
	require.Equal(t, "running", operationStatus(t, store, pendingID),
		"running operation with pending items is never claimed")
}

func TestClaimNextQueuedPrefersNewerQueuedOverOlderFinalizationReady(t *testing.T) {
	ctx := context.Background()
	store := startAgentOpsConcurrencyStore(t, ctx)
	poisonID := enqueueTestOperation(t, ctx, store, "wsp-poison", "poison")
	mustSetOperationRunning(t, ctx, store, poisonID)
	require.NoError(t, store.CompleteItem(
		ctx, poisonID, 0, "failed", nil, "projection keeps failing",
	))
	queuedID := enqueueTestOperation(t, ctx, store, "wsp-queued", "queued")

	claimed, err := store.ClaimNextOperation(ctx, "")
	require.NoError(t, err)
	require.NotNil(t, claimed)
	require.Equal(t, queuedID, claimed.ID,
		"an older finalization retry must not starve queued work")
	require.Equal(t, "running", operationStatus(t, store, poisonID))
}

func TestClaimNextQueuedForWorkspaceGroupSkipsForeignFinalizationReady(t *testing.T) {
	ctx := context.Background()
	store := startAgentOpsConcurrencyStore(t, ctx)
	seedOperationWorkspacePlacements(t, ctx, store, []operationWorkspacePlacement{
		{workspaceID: "wsp-final-foreign", groupID: "group-b", state: "active"},
		{workspaceID: "wsp-final-local", groupID: "group-a", state: "active"},
	})
	foreignID := enqueueTestOperation(t, ctx, store, "wsp-final-foreign", "final-foreign")
	localID := enqueueTestOperation(t, ctx, store, "wsp-final-local", "final-local")
	for _, id := range []int64{foreignID, localID} {
		mustSetOperationRunning(t, ctx, store, id)
		require.NoError(t, store.CompleteItem(ctx, id, 0, "applied", json.RawMessage(`{"agentId":"agt-ready"}`), ""))
	}

	claimed, err := store.ClaimNextOperation(ctx, "group-a")
	require.NoError(t, err)
	require.NotNil(t, claimed)
	require.Equal(t, localID, claimed.ID)
	require.Equal(t, "running", operationStatus(t, store, foreignID))
}

func TestClaimNextQueuedForWorkspaceGroupPrefersLocalQueuedOverOlderFinalizationReady(t *testing.T) {
	ctx := context.Background()
	store := startAgentOpsConcurrencyStore(t, ctx)
	seedOperationWorkspacePlacements(t, ctx, store, []operationWorkspacePlacement{
		{workspaceID: "wsp-priority-poison", groupID: "group-a", state: "active"},
		{workspaceID: "wsp-priority-foreign", groupID: "group-b", state: "active"},
		{workspaceID: "wsp-priority-local", groupID: "group-a", state: "active"},
	})
	poisonID := enqueueTestOperation(t, ctx, store, "wsp-priority-poison", "priority-poison")
	mustSetOperationRunning(t, ctx, store, poisonID)
	require.NoError(t, store.CompleteItem(
		ctx, poisonID, 0, "failed", nil, "projection keeps failing",
	))
	foreignID := enqueueTestOperation(t, ctx, store, "wsp-priority-foreign", "priority-foreign")
	localID := enqueueTestOperation(t, ctx, store, "wsp-priority-local", "priority-local")

	claimed, err := store.ClaimNextOperation(ctx, "group-a")
	require.NoError(t, err)
	require.NotNil(t, claimed)
	require.Equal(t, localID, claimed.ID,
		"an older local finalization retry must not starve queued local work")
	require.Equal(t, "running", operationStatus(t, store, poisonID))
	require.Equal(t, "queued", operationStatus(t, store, foreignID),
		"workspace-group claim must not mutate foreign queued work")
}

func mustSetOperationRunning(t *testing.T, ctx context.Context, store *Store, operationID int64) {
	t.Helper()
	result, err := store.db.Exec(ctx,
		`UPDATE agent_operations SET status = 'running' WHERE id = $1`,
		operationID,
	)
	require.NoError(t, err)
	require.Equal(t, int64(1), result.RowsAffected())
}

func TestClaimNextQueuedForWorkspaceGroupSkipsOlderForeignOperation(t *testing.T) {
	ctx := context.Background()
	store := startAgentOpsConcurrencyStore(t, ctx)
	seedOperationWorkspacePlacements(t, ctx, store, []operationWorkspacePlacement{
		{workspaceID: "wsp_foreign", groupID: "group-b", state: "active"},
		{workspaceID: "wsp_local_active", groupID: "group-a", state: "active"},
		{workspaceID: "wsp_local_draining", groupID: "group-a", state: "draining"},
		{workspaceID: "wsp_local_moving", groupID: "group-a", state: "moving"},
		{workspaceID: "wsp_local_disabled", groupID: "group-a", state: "disabled"},
	})

	foreignID := enqueueTestOperation(t, ctx, store, "wsp_foreign", "foreign")
	localActiveID := enqueueTestOperation(t, ctx, store, "wsp_local_active", "local-active")
	localDrainingID := enqueueTestOperation(t, ctx, store, "wsp_local_draining", "local-draining")
	localMovingID := enqueueTestOperation(t, ctx, store, "wsp_local_moving", "local-moving")
	disabledID := enqueueTestOperation(t, ctx, store, "wsp_local_disabled", "local-disabled")

	for _, wantID := range []int64{localActiveID, localDrainingID, localMovingID} {
		claimed, err := store.ClaimNextOperation(ctx, " group-a ")
		require.NoError(t, err)
		require.NotNil(t, claimed)
		require.Equal(t, wantID, claimed.ID)
	}
	claimed, err := store.ClaimNextOperation(ctx, "group-a")
	require.NoError(t, err)
	require.Nil(t, claimed, "non-routable and foreign operations are not eligible")

	require.Equal(t, "queued", operationStatus(t, store, foreignID),
		"workspace-profile claim must not mutate the older foreign operation")
	require.Equal(t, "queued", operationStatus(t, store, disabledID),
		"non-routable local placement must remain queued")
}

func TestClaimNextQueuedForEmptyWorkspaceGroupUsesGlobalOldest(t *testing.T) {
	ctx := context.Background()
	store := startAgentOpsConcurrencyStore(t, ctx)
	seedOperationWorkspacePlacements(t, ctx, store, []operationWorkspacePlacement{
		{workspaceID: "wsp_global_foreign", groupID: "group-b", state: "active"},
		{workspaceID: "wsp_global_local", groupID: "group-a", state: "active"},
	})

	foreignID := enqueueTestOperation(t, ctx, store, "wsp_global_foreign", "global-foreign")
	localID := enqueueTestOperation(t, ctx, store, "wsp_global_local", "global-local")

	claimed, err := store.ClaimNextOperation(ctx, " ")
	require.NoError(t, err)
	require.NotNil(t, claimed)
	require.Equal(t, foreignID, claimed.ID, "empty group preserves global oldest-first behavior")
	require.Equal(t, "running", operationStatus(t, store, foreignID))
	require.Equal(t, "queued", operationStatus(t, store, localID))
}

type operationWorkspacePlacement struct {
	workspaceID string
	groupID     string
	state       string
}

func seedOperationWorkspacePlacements(t *testing.T, ctx context.Context, store *Store, placements []operationWorkspacePlacement) {
	t.Helper()
	now := time.Date(2026, 8, 2, 12, 0, 0, 0, time.UTC)
	_, err := store.db.Exec(ctx, `
INSERT INTO users (id, email, name, display_name, password_hash, created_at, updated_at)
VALUES ('usr_operation_scope', 'operation-scope@example.com', 'Operation Scope', 'Operation Scope', 'hash', $1, $1)
`, now)
	require.NoError(t, err)
	_, err = store.db.Exec(ctx, `
INSERT INTO organizations (id, slug, name, created_by_user_id, created_at, updated_at)
VALUES ('org_operation_scope', 'operation-scope', 'Operation Scope', 'usr_operation_scope', $1, $1)
`, now)
	require.NoError(t, err)
	for _, placement := range placements {
		_, err = store.db.Exec(ctx, `
INSERT IGNORE INTO workspace_groups (
  workspace_group_id, display_name, state, resource_profile, service_ref, created_at, updated_at
) VALUES ($1, $1, 'active', '', $1, $2, $2)
`, placement.groupID, now)
		require.NoError(t, err)
		_, err = store.db.Exec(ctx, `
INSERT INTO workspaces (id, organization_id, slug, name, created_at, updated_at)
VALUES ($1, 'org_operation_scope', $1, $1, $2, $2)
`, placement.workspaceID, now)
		require.NoError(t, err)
		_, err = store.db.Exec(ctx, `
INSERT INTO workspace_placements (
  workspace_id, workspace_group_id, state, placement_epoch, created_at, updated_at
) VALUES ($1, $2, $3, 1, $4, $4)
`, placement.workspaceID, placement.groupID, placement.state, now)
		require.NoError(t, err)
	}
}

func enqueueTestOperation(t *testing.T, ctx context.Context, store *Store, workspaceID, suffix string) int64 {
	t.Helper()
	var operationID int64
	require.NoError(t, coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		var err error
		operationID, err = store.InTx(tx).EnqueueOperationInTx(ctx, EnqueueInput{
			WorkspaceID: workspaceID, ApplicantAgentID: "agt-" + suffix, OwnerUserID: "usr_operation_scope",
			Action: "delete", Payload: json.RawMessage(`{}`),
			ApprovalID: "appr-" + suffix, RequestID: "req-" + suffix,
			Items: []json.RawMessage{json.RawMessage(`{"agentId":"agt-target"}`)},
		})
		return err
	}))
	return operationID
}