package extensionflow

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

// TestCatalogStateReconcileRepairsMismatchedState proves a syntactically valid
// but non-canonical durable state fails the invocation fence. A Workspace
// reconcile compares the live canonical catalog_state, advances the revision,
// restores the exact state, and makes invocation available again.
func TestCatalogStateReconcileRepairsMismatchedState(t *testing.T) {
	ctx := context.Background()
	db := startServiceRollbackDB(t, ctx)
	fixture := seedServiceFixture(t, ctx, db, "state_repair")
	require.NoError(t, registerProviderWatchBuiltinsForTest(fixture.Registry))
	service := newService(db, fixture.Registry)
	seedProductAppConnection(t, db, fixture, "state_repair")
	now := time.Now().UTC()

	// Materialize a current catalog, then replace its canonical projection with
	// a different valid JSON array. Keeping it non-empty preserves the enabled
	// invariant, so this specifically exercises semantic state mismatch.
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return service.ApplyAgentCatalogChange(
			ctx, tx, fixture.WorkspaceID, "agent-1", now)
	}))
	before := readCatalogStateRow(
		t,
		ctx,
		db,
		fixture.WorkspaceID,
		"agent-1",
	)
	require.Equal(t, int64(1), before.Version)
	corruptState := `[{"corrupt":true}]`
	mustExecService(t, db, `
UPDATE agent_extension_catalogs
SET catalog_state = $1
WHERE workspace_id = $2 AND agent_id = 'agent-1'`,
		corruptState, fixture.WorkspaceID)
	registered, ok := fixture.Registry.LookupByIdentity("github", 1)
	require.True(t, ok)

	// The mismatched state fails the invocation fence.
	_, err := service.Invoke(ctx, invocationParams(
		fixture, "agent-1", "github", "list-watches",
		registered.Definition.Digest, before.Version))
	require.Error(t, err,
		"a mismatched catalog_state must fail the invocation fence")

	// The bounded worker re-derives catalog_state and advances the version.
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return service.EnqueueWorkspaceCatalogReconcile(
			ctx, tx, fixture.WorkspaceID, now)
	}))
	result, err := service.reconcileWorkspaceCatalogJobBatch(
		ctx, fixture.WorkspaceID, "", 1, now)
	require.NoError(t, err)
	require.Equal(t, 1, result.Scanned)

	repaired := readCatalogStateRow(
		t,
		ctx,
		db,
		fixture.WorkspaceID,
		"agent-1",
	)
	require.Equal(t, before.Version+1, repaired.Version,
		"the worker must advance the mismatched catalog revision")
	require.Equal(t, before.State, repaired.State,
		"the worker must restore the live canonical catalog_state")
	require.NotEqual(t, corruptState, repaired.State)
	require.Equal(t, before.Enabled, repaired.Enabled)

	// The reconciled catalog serves invocation again.
	_, err = service.Invoke(ctx, invocationParams(
		fixture, "agent-1", "github", "list-watches",
		registered.Definition.Digest, repaired.Version))
	require.NoError(t, err)
}

// TestEnqueueFencesStaleLeasedHolder proves the generation fence end to end
// through the REAL store enqueue/release paths — the exact semantics
// migration 00348 executes for the scheme upgrade. A pre-change holder
// mid-pass in a leased job is fenced by the enqueue's generation bump: its
// completion neither deletes the job nor carries its cursor forward, and
// the job requeues pending from the beginning.
func TestEnqueueFencesStaleLeasedHolder(t *testing.T) {
	ctx := context.Background()
	db := startServiceRollbackDB(t, ctx)
	fixture := seedServiceFixture(t, ctx, db, "leasefence")
	now := time.Now().UTC()

	// A holder leased generation 9 and is mid-pass.
	mustExecService(t, db, `
INSERT INTO extension_catalog_reconcile_jobs (
  workspace_id, generation, processing_generation, cursor_agent_id,
  status, attempts, last_error_code, last_error, lease_owner,
  lease_expires_at, next_retry_at, created_at, updated_at
) VALUES ($1, 9, 9, 'agent-mid', 'leased', 1, '', '', 'holder-1',
  $2, NULL, $2, $2)`, fixture.WorkspaceID, now)

	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		_, err := service_jobsForTest(db).InTx(tx).
			EnqueueWorkspaceCatalogReconcile(
				ctx, fixture.WorkspaceID, now)
		return err
	}))

	var generation, processing int64
	var status, cursor, leaseOwner string
	require.NoError(t, db.QueryRow(ctx, `
SELECT generation, processing_generation, status, cursor_agent_id, lease_owner
FROM extension_catalog_reconcile_jobs WHERE workspace_id=$1`,
		fixture.WorkspaceID).
		Scan(&generation, &processing, &status, &cursor, &leaseOwner))
	require.Equal(t, int64(10), generation,
		"the enqueue re-arms the generation under the holder")
	require.Equal(t, int64(9), processing,
		"the in-flight holder's processing_generation is preserved")
	require.Equal(t, "leased", status)
	require.Equal(t, "holder-1", leaseOwner)

	// The stale holder finishes with its pre-bump generation: the fence
	// must keep the job and restart the round instead of completing it.
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return service_jobsForTest(db).InTx(tx).
			ReleaseWorkspaceCatalogReconcileBatch(
				ctx, fixture.WorkspaceID, "holder-1", 9, "", false, now)
	}))
	require.NoError(t, db.QueryRow(ctx, `
SELECT generation, processing_generation, status, cursor_agent_id, lease_owner
FROM extension_catalog_reconcile_jobs WHERE workspace_id=$1`,
		fixture.WorkspaceID).
		Scan(&generation, &processing, &status, &cursor, &leaseOwner))
	require.Equal(t, int64(10), generation,
		"the stale holder's completion must not delete the re-armed job")
	require.Equal(t, "pending", status)
	require.Equal(t, "", cursor,
		"the stale holder must not carry its mid-flight cursor forward")
	require.Zero(t, processing,
		"the requeued round starts from the beginning")
	require.Equal(t, "", leaseOwner)

	// The rollback (Down) symmetry with a second-generation holder: a
	// worker leases generation 10 and is mid-pass when the migration's
	// rollback re-arm bumps the job to generation 11. The stale holder
	// releases with processing=10: the fence must keep the job, requeue it
	// pending, clear the cursor, and reset processing for a full pass —
	// the exact state the rolled-back binary's worker then drains.
	mustExecService(t, db, `
UPDATE extension_catalog_reconcile_jobs
SET generation=10, processing_generation=10, status='leased',
    cursor_agent_id='agent-v2-mid', lease_owner='holder-2',
    lease_expires_at=$2, updated_at=$2
WHERE workspace_id=$1`, fixture.WorkspaceID, now)
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		_, err := service_jobsForTest(db).InTx(tx).
			EnqueueWorkspaceCatalogReconcile(
				ctx, fixture.WorkspaceID, now)
		return err
	}))
	require.NoError(t, db.QueryRow(ctx, `
SELECT generation, processing_generation, status, cursor_agent_id, lease_owner
FROM extension_catalog_reconcile_jobs WHERE workspace_id=$1`,
		fixture.WorkspaceID).
		Scan(&generation, &processing, &status, &cursor, &leaseOwner))
	require.Equal(t, int64(11), generation,
		"the rollback re-arm advances generation under the new holder")
	require.Equal(t, int64(10), processing)
	require.Equal(t, "leased", status)
	require.Equal(t, "holder-2", leaseOwner)

	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return service_jobsForTest(db).InTx(tx).
			ReleaseWorkspaceCatalogReconcileBatch(
				ctx, fixture.WorkspaceID, "holder-2", 10, "", false, now)
	}))
	require.NoError(t, db.QueryRow(ctx, `
SELECT generation, processing_generation, status, cursor_agent_id, lease_owner
FROM extension_catalog_reconcile_jobs WHERE workspace_id=$1`,
		fixture.WorkspaceID).
		Scan(&generation, &processing, &status, &cursor, &leaseOwner))
	require.Equal(t, int64(11), generation,
		"the stale v2 holder's completion must not delete the job")
	require.Equal(t, "pending", status)
	require.Equal(t, "", cursor,
		"the stale v2 holder must not carry its cursor forward")
	require.Zero(t, processing,
		"the requeued round starts from the beginning")
	require.Equal(t, "", leaseOwner)
}

func service_jobsForTest(db coredb.DB) *extension.Store {
	return extension.New(db)
}