package extension

import (
	"context"
	"math"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

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

func validCatalogState(seed string) string {
	return `[{"extensionName":"` + seed + `"}]`
}

// seedCatalogDB seeds a workspace + agent and returns the DB, workspace, agent.
func seedCatalogDB(t *testing.T, ctx context.Context, suffix string) (coredb.DB, string, string) {
	t.Helper()
	db := startRollbackDB(t, ctx)
	now := time.Now().UTC()
	fixture := seedExtensionStoreFixture(t, db, "cat_"+suffix, now)
	agent := "agt_ext_cat_" + suffix
	seedExtensionStoreAgent(t, db, fixture.WorkspaceID, fixture.UserID, agent, now)
	return db, fixture.WorkspaceID, agent
}

func changedMutator(state string) CatalogChangeMutator {
	return func(current int64, storedState string) (string, bool, bool, error) {
		return state, true, true, nil
	}
}

// TestAgentExtensionCatalog_FirstCreateAdvanceNoOp proves the lifecycle: missing
// -> v1, v1 -> v2 on change, and a no-op leaves version + state unchanged.
func TestAgentExtensionCatalog_FirstCreateAdvanceNoOp(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	db, ws, agent := seedCatalogDB(t, ctx, "lifecycle")
	store := New(db)

	v, err := runCatalogChange(ctx, db, store, ws, agent, changedMutator(validCatalogState("a")))
	require.NoError(t, err)
	assert.Equal(t, int64(1), v)
	c := loadCatalog(t, ctx, db, ws, agent)
	assert.Equal(t, int64(1), c.Version)
	assert.Equal(t, validCatalogState("a"), c.State)
	assert.True(t, c.Enabled)

	v, err = runCatalogChange(ctx, db, store, ws, agent, changedMutator(validCatalogState("b")))
	require.NoError(t, err)
	assert.Equal(t, int64(2), v)

	noOp := func(current int64, storedState string) (string, bool, bool, error) {
		return "", false, false, nil
	}
	v, err = runCatalogChange(ctx, db, store, ws, agent, noOp)
	require.NoError(t, err)
	assert.Equal(t, int64(2), v, "no-op must not advance the version")
	c = loadCatalog(t, ctx, db, ws, agent)
	assert.Equal(t, validCatalogState("b"), c.State, "no-op must not rewrite the state")
}

// TestAgentExtensionCatalog_MissingRowNoOp proves a no-op on a never-touched
// Agent stays sparse: no v0 row is created and no update event is written.
func TestAgentExtensionCatalog_MissingRowNoOp(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	db, ws, agent := seedCatalogDB(t, ctx, "missingnoop")
	store := New(db)

	noOp := func(current int64, storedState string) (string, bool, bool, error) {
		return "", false, false, nil
	}
	v, err := runCatalogChange(ctx, db, store, ws, agent, noOp)
	require.NoError(t, err)
	assert.Equal(t, int64(0), v, "no-op on a fresh Agent stays at v0")
	var n int
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return tx.QueryRow(ctx, `SELECT COUNT(*) FROM agent_extension_catalogs WHERE workspace_id=$1 AND agent_id=$2`, ws, agent).Scan(&n)
	}))
	assert.Equal(t, 0, n, "fresh no-op must not create an empty v0 row")
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return tx.QueryRow(ctx, `SELECT COUNT(*) FROM agent_extension_catalog_updates WHERE workspace_id=$1 AND agent_id=$2`, ws, agent).Scan(&n)
	}))
	assert.Equal(t, 0, n, "no update event on a no-op")
}

// TestAgentExtensionCatalog_MissingRowNoOpCanKeepTombstone proves explicit
// opt-out callers can keep a v0 row as a concurrency fence without advancing
// the version or writing a delivery event.
func TestAgentExtensionCatalog_MissingRowNoOpCanKeepTombstone(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	db, ws, agent := seedCatalogDB(t, ctx, "tombstone")
	store := New(db)

	noOp := func(current int64, storedState string) (string, bool, bool, error) {
		return "", false, false, nil
	}
	var version int64
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		v, err := store.InTx(tx).ApplyExtensionCatalogChange(ctx, ws, agent, time.Now().UTC(), noOp, ApplyExtensionCatalogChangeOptions{
			KeepEmptyRowOnNoop: true,
		})
		version = v
		return err
	}))
	assert.Equal(t, int64(0), version)

	c := loadCatalog(t, ctx, db, ws, agent)
	assert.Equal(t, int64(0), c.Version)
	assert.Empty(t, c.State)
	assert.False(t, c.Enabled)
	var n int
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return tx.QueryRow(ctx, `SELECT COUNT(*) FROM agent_extension_catalog_updates WHERE workspace_id=$1 AND agent_id=$2`, ws, agent).Scan(&n)
	}))
	assert.Equal(t, 0, n, "tombstone is not a deliverable Catalog change")
}

// TestAgentExtensionCatalog_LastRevokeKeepsRow proves enabled=false with the row
// retained and version incremented.
func TestAgentExtensionCatalog_LastRevokeKeepsRow(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	db, ws, agent := seedCatalogDB(t, ctx, "lastrt")
	store := New(db)
	_, err := runCatalogChange(ctx, db, store, ws, agent, changedMutator(validCatalogState("a")))
	require.NoError(t, err)
	lastRevoke := func(current int64, storedState string) (string, bool, bool, error) {
		return "[]", false, true, nil
	}
	v, err := runCatalogChange(ctx, db, store, ws, agent, lastRevoke)
	require.NoError(t, err)
	assert.Equal(t, int64(2), v)
	c := loadCatalog(t, ctx, db, ws, agent)
	assert.False(t, c.Enabled)
	assert.Equal(t, int64(2), c.Version)
}

// TestAgentExtensionCatalog_StateValidation proves positive versions require a
// JSON array whose empty/non-empty state agrees with enabled.
func TestAgentExtensionCatalog_StateValidation(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	db, ws, agent := seedCatalogDB(t, ctx, "digest")
	store := New(db)
	cases := []string{
		"",
		"not-json",
		"{}",
		"[",
		"[true",
		" [true]",
	}
	for _, state := range cases {
		_, err := runCatalogChange(ctx, db, store, ws, agent, changedMutator(state))
		require.Error(t, err, "state %q must be rejected", state)
	}
	_, err := runCatalogChange(ctx, db, store, ws, agent, changedMutator(validCatalogState("a")))
	require.NoError(t, err)
}

// TestAgentExtensionCatalog_SequentialChangesOneEventEach proves two changed
// mutations produce v1, v2 with exactly one pending update event per version.
func TestAgentExtensionCatalog_SequentialChangesOneEventEach(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	db, ws, agent := seedCatalogDB(t, ctx, "seq")
	store := New(db)
	_, err := runCatalogChange(ctx, db, store, ws, agent, changedMutator(validCatalogState("a")))
	require.NoError(t, err)
	_, err = runCatalogChange(ctx, db, store, ws, agent, changedMutator(validCatalogState("b")))
	require.NoError(t, err)

	var n int
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return tx.QueryRow(ctx, `SELECT COUNT(*) FROM agent_extension_catalog_updates WHERE workspace_id=$1 AND agent_id=$2`, ws, agent).Scan(&n)
	}))
	assert.Equal(t, 2, n, "one update event per changed mutation")
}

// TestAgentExtensionCatalog_Overflow proves MaxInt64 version is rejected.
func TestAgentExtensionCatalog_Overflow(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	db, ws, agent := seedCatalogDB(t, ctx, "overflow")
	store := New(db)
	// Seed a catalog row already at MaxInt64.
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		_, err := tx.Exec(ctx, `INSERT INTO agent_extension_catalogs (workspace_id, agent_id, version, catalog_state, enabled, created_at, updated_at) VALUES ($1,$2,$3,$4,1,$5,$5)`, ws, agent, math.MaxInt64, validCatalogState("a"), time.Now().UTC())
		return err
	}))
	_, err := runCatalogChange(ctx, db, store, ws, agent, changedMutator(validCatalogState("b")))
	require.Error(t, err, "overflow must fail closed")
}

// TestAgentExtensionCatalog_ForeignWorkspaceAgent proves the single entry rejects
// an Agent that does not belong to the workspace.
func TestAgentExtensionCatalog_ForeignWorkspaceAgent(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	db := startRollbackDB(t, ctx)
	now := time.Now().UTC()
	fixtureA := seedExtensionStoreFixture(t, db, "foreign_a", now)
	fixtureB := seedExtensionStoreFixture(t, db, "foreign_b", now)
	const agent = "agt_ext_foreign"
	seedExtensionStoreAgent(t, db, fixtureA.WorkspaceID, fixtureA.UserID, agent, now) // agent in wsA

	// Apply against wsB (agent is not in wsB) -> forbidden before any row write.
	err := coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		_, err := New(db).InTx(tx).ApplyExtensionCatalogChange(ctx, fixtureB.WorkspaceID, agent, now, changedMutator(validCatalogState("a")), ApplyExtensionCatalogChangeOptions{})
		return err
	})
	require.Error(t, err)
}

// TestAgentExtensionCatalog_MutatorErrorRollback proves a mutator error rolls
// back the whole transaction: no catalog row, no update event.
func TestAgentExtensionCatalog_MutatorErrorRollback(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	db, ws, agent := seedCatalogDB(t, ctx, "rollback")
	store := New(db)
	errMutator := func(current int64, storedState string) (string, bool, bool, error) {
		return "", false, false, assert.AnError
	}
	_, err := runCatalogChange(ctx, db, store, ws, agent, errMutator)
	require.Error(t, err)

	var n int
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return tx.QueryRow(ctx, `SELECT COUNT(*) FROM agent_extension_catalogs WHERE workspace_id=$1 AND agent_id=$2`, ws, agent).Scan(&n)
	}))
	assert.Equal(t, 0, n, "mutator error rolls back — no catalog row")
}

// TestAgentExtensionCatalog_NegativeVersionRejected proves a corrupt negative
// version fails closed at Load (and Apply), via the unified validateCatalogState.
func TestAgentExtensionCatalog_NegativeVersionRejected(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	db, ws, agent := seedCatalogDB(t, ctx, "negver")
	now := time.Now().UTC()
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		_, err := tx.Exec(ctx, `INSERT INTO agent_extension_catalogs (workspace_id, agent_id, version, catalog_state, enabled, created_at, updated_at) VALUES ($1,$2,-1,$3,1,$4,$4)`, ws, agent, validCatalogState("a"), now)
		return err
	}))
	// Load must fail closed on the negative version.
	_, _, err := loadCatalogRaw(ctx, db, ws, agent)
	require.Error(t, err, "negative version must fail closed at Load")
	// Apply must also fail closed.
	_, err = runCatalogChange(ctx, db, New(db), ws, agent, changedMutator(validCatalogState("b")))
	require.Error(t, err, "negative version must fail closed at Apply")
}

// TestAgentExtensionCatalog_ForeignPairLoadRejected proves Load rejects a
// (workspace B, agent A) row whose agent belongs to a different workspace — the
// ownership fence runs on Load too (the worker depends on it).
func TestAgentExtensionCatalog_ForeignPairLoadRejected(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	db := startRollbackDB(t, ctx)
	now := time.Now().UTC()
	fixtureA := seedExtensionStoreFixture(t, db, "fpa", now)
	fixtureB := seedExtensionStoreFixture(t, db, "fpb", now)
	const agent = "agt_ext_fp"
	seedExtensionStoreAgent(t, db, fixtureA.WorkspaceID, fixtureA.UserID, agent, now) // agent belongs to wsA
	// Directly persist a catalog row under wsB for the wsA-owned agent (the two
	// independent FKs allow it).
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		_, err := tx.Exec(ctx, `INSERT INTO agent_extension_catalogs (workspace_id, agent_id, version, catalog_state, enabled, created_at, updated_at) VALUES ($1,$2,1,$3,1,$4,$4)`, fixtureB.WorkspaceID, agent, validCatalogState("a"), now)
		return err
	}))
	_, _, err := loadCatalogRaw(ctx, db, fixtureB.WorkspaceID, agent)
	require.Error(t, err, "Load must reject a foreign (workspace,agent) pair")
}

// (Schema-CHECK-rejection and store-level concurrency tests are intentionally
// omitted here: this TiDB test env silently ignores CHECK constraints — SHOW
// CREATE carries only FK constraints — so the update-ledger iff invariants
// (leased/delivered/version>0/status) are enforced in the app layer: the store
// only ever writes valid pending rows, and the delivery worker validates every
// state transition in Go. The migration still DECLARES the CHECKs as the schema
// contract, effective when tidb_enable_check_constraint=ON. The concurrent-
// serialization guarantee is covered at the service level with a pooled DB
// (startServiceConcurrentDB) once grant/revoke wiring lands — the store-level
// rollback DB is single-connection and cannot exercise true concurrency.)

func runCatalogChange(ctx context.Context, db coredb.DB, store *Store, ws, agent string, mutate CatalogChangeMutator) (int64, error) {
	var version int64
	if err := coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		v, err := store.InTx(tx).ApplyExtensionCatalogChange(ctx, ws, agent, time.Now().UTC(), mutate, ApplyExtensionCatalogChangeOptions{})
		version = v
		return err
	}); err != nil {
		return 0, err
	}
	return version, nil
}

// loadCatalogRaw loads the catalog without require, returning the error so
// adversarial tests can assert fail-closed behavior.
func loadCatalogRaw(ctx context.Context, db coredb.DB, ws, agent string) (AgentExtensionCatalog, bool, error) {
	store := New(db)
	var catalog AgentExtensionCatalog
	var found bool
	err := coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		c, f, e := store.InTx(tx).LoadAgentExtensionCatalogForUpdate(ctx, ws, agent)
		catalog, found = c, f
		return e
	})
	return catalog, found, err
}

func loadCatalog(t *testing.T, ctx context.Context, db coredb.DB, ws, agent string) AgentExtensionCatalog {
	t.Helper()
	store := New(db)
	var catalog AgentExtensionCatalog
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		c, _, err := store.InTx(tx).LoadAgentExtensionCatalogForUpdate(ctx, ws, agent)
		catalog = c
		return err
	}))
	return catalog
}