package extensionflow

import (
	"context"
	"fmt"
	"sort"
	"sync"
	"testing"
	"time"

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

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

// TestAgentExtensionCatalog_ConcurrentSerializes proves two concurrent changed
// mutations on the same Agent serialize via the catalog row lock (real pooled
// TiDB): BOTH calls succeed and return exactly versions [1, 2], two unique
// update events exist at exactly versions [1, 2], and the final catalog is at
// version 2 — no lost update, no duplicate version.
func TestAgentExtensionCatalog_ConcurrentSerializes(t *testing.T) {
	ctx := context.Background()
	db := startServiceConcurrentDB(t, ctx)
	fixture := seedServiceFixture(t, ctx, db, "cat_concurrent")
	store := extension.New(db)
	ws := fixture.WorkspaceID
	const agent = "agent-1"
	state := func(seed string) string {
		return `[{"extension":"` + seed + `"}]`
	}

	var (
		mu       sync.Mutex
		versions []int64
	)
	apply := func(nextState string) {
		var version int64
		err := coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
			v, e := store.InTx(tx).ApplyExtensionCatalogChange(ctx, ws, agent, time.Now().UTC(), func(current int64, _ string) (string, bool, bool, error) {
				return nextState, true, true, nil
			}, extension.ApplyExtensionCatalogChangeOptions{})
			version = v
			return e
		})
		require.NoError(t, err)
		mu.Lock()
		versions = append(versions, version)
		mu.Unlock()
	}

	var wg sync.WaitGroup
	start := make(chan struct{})
	wg.Add(2)
	go func() { defer wg.Done(); <-start; apply(state("a")) }()
	go func() { defer wg.Done(); <-start; apply(state("b")) }()
	close(start)
	wg.Wait()

	// Both calls succeeded and returned exactly versions 1 and 2.
	sort.Slice(versions, func(i, j int) bool { return versions[i] < versions[j] })
	assert.Equal(t, []int64{1, 2}, versions, "the two concurrent changes return exactly v1 and v2")

	// The catalog is at version 2.
	var catalog extension.AgentExtensionCatalog
	var found bool
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		c, f, err := store.InTx(tx).LoadAgentExtensionCatalogForUpdate(ctx, ws, agent)
		catalog, found = c, f
		return err
	}))
	require.True(t, found)
	assert.Equal(t, int64(2), catalog.Version)

	// Exactly two update events at versions 1 and 2.
	var eventVersions []int64
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		rows, err := tx.Query(ctx, `SELECT version FROM agent_extension_catalog_updates WHERE workspace_id=$1 AND agent_id=$2 ORDER BY version`, ws, agent)
		if err != nil {
			return err
		}
		defer rows.Close()
		for rows.Next() {
			var v int64
			if err := rows.Scan(&v); err != nil {
				return err
			}
			eventVersions = append(eventVersions, v)
		}
		return rows.Err()
	}))
	assert.Equal(t, []int64{1, 2}, eventVersions, "one unique update event per version, at exactly v1 and v2")
}

// TestCatalogReconcile_ConcurrentDefaultAllowIsIdempotent proves two server
// instances can scan the same absent default-allow Catalog concurrently. The
// catalog-row lock plus canonical-state gate materialize exactly v1
// and exactly one update; the stale second candidate is a no-op.
func TestCatalogReconcile_ConcurrentDefaultAllowIsIdempotent(t *testing.T) {
	ctx := context.Background()
	db := startServiceConcurrentDB(t, ctx)
	fixture := seedServiceFixture(t, ctx, db, "cat_reconcile_concurrent")
	require.NoError(t, registerProviderWatchBuiltinsForTest(fixture.Registry))
	serviceA := newService(db, fixture.Registry)
	serviceB := newService(db, fixture.Registry)
	seedProductAppConnection(t, db, fixture, "concurrent")

	var (
		wg      sync.WaitGroup
		start   = make(chan struct{})
		results = make([]int, 2)
		errs    = make([]error, 2)
	)
	for i, service := range []*Service{serviceA, serviceB} {
		i, service := i, service
		wg.Add(1)
		go func() {
			defer wg.Done()
			<-start
			results[i], errs[i] = service.BackfillDefaultAllowCatalogs(ctx)
		}()
	}
	close(start)
	wg.Wait()

	require.NoError(t, errs[0])
	require.NoError(t, errs[1])
	assert.Equal(t, 6, results[0]+results[1], "each of the six Agents materializes once across both workers")

	catalog := readCatalogStateRow(t, ctx, db, fixture.WorkspaceID, "agent-1")
	assert.Equal(t, int64(1), catalog.Version)
	assert.NotEqual(t, "[]", catalog.State)
	assert.True(t, catalog.Enabled)
	assertUpdateCount(t, ctx, db, fixture.WorkspaceID, "agent-1", 1)

	n, err := serviceA.BackfillDefaultAllowCatalogs(ctx)
	require.NoError(t, err)
	assert.Equal(t, 0, n, "a repeated full scan is a no-op")
	assert.Equal(t, int64(1), readCatalogHint(t, ctx, db, fixture.WorkspaceID, "agent-1").Version)
}

func TestWorkspaceCatalogReconcile_OptOutTombstoneFencesStaleMissingRowSnapshot(t *testing.T) {
	// Finish DB initialization and fixture seeding before the fenced body
	// deadline starts. The separate setup context gives the open/ping path a
	// wider budget when CI resource contention delays the initialized DB.
	setupCtx, setupCancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer setupCancel()
	db := startServiceConcurrentDB(t, setupCtx)
	fixture := seedServiceFixture(t, setupCtx, db, "cat_reconcile_optout_race")
	require.NoError(t, registerProviderWatchBuiltinsForTest(fixture.Registry))
	service := newService(db, fixture.Registry)
	installationID, _ := seedProductAppConnection(t, db, fixture, "optoutrace")

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	workerReady := make(chan struct{})
	resumeWorker := make(chan struct{})
	workerErr := make(chan error, 1)
	go func() {
		workerErr <- coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
			extTx := service.store.InTx(tx)
			agentIDs := []string{"agent-1"}
			authorizationsByAgent, err := extTx.ListAgentAppAuthorizationDecisionsForAgents(ctx, fixture.WorkspaceID, agentIDs)
			if err != nil {
				return err
			}
			catalogsByAgent, rowErrors, err := extTx.LoadAgentExtensionCatalogsForUpdate(ctx, fixture.WorkspaceID, agentIDs)
			if err != nil {
				return err
			}
			if len(rowErrors) != 0 {
				return fmt.Errorf("unexpected catalog row errors: %v", rowErrors)
			}
			selections, err := extTx.ListActiveAppHostedExtensionSelections(ctx, fixture.WorkspaceID)
			if err != nil {
				return err
			}
			decisionsByInstallation, err := validateAppAuthorizationDecisions(
				fixture.WorkspaceID,
				"agent-1",
				authorizationsByAgent["agent-1"],
			)
			if err != nil {
				return err
			}
			entries, enabled, err := service.projectAppHostedCatalogEntriesFromSelections(
				ctx,
				tx,
				fixture.WorkspaceID,
				selections,
				decisionsByInstallation,
			)
			if err != nil {
				return err
			}
			state, newEnabled, changed, err := projectedCatalogStateChange(
				entries,
				enabled,
				0,
				"",
			)
			if err != nil {
				return err
			}
			if !changed || !newEnabled {
				return fmt.Errorf("stale worker snapshot should see default-allow catalog enabled")
			}
			close(workerReady)
			select {
			case <-resumeWorker:
			case <-ctx.Done():
				return ctx.Err()
			}
			current, found := catalogsByAgent["agent-1"]
			_, err = extTx.ApplyPrecomputedExtensionCatalogChange(
				ctx,
				fixture.WorkspaceID,
				"agent-1",
				current,
				found,
				state,
				newEnabled,
				time.Now().UTC(),
			)
			return err
		})
	}()

	select {
	case <-workerReady:
	case <-ctx.Done():
		t.Fatal(ctx.Err())
	}
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		now := time.Now().UTC()
		if _, err := tx.Exec(ctx, `
INSERT INTO service_principal_authorizations (
    id, workspace_id, service_definition_id, service_installation_id,
    service_kind, principal_kind, principal_ref, status, disabled_reason,
    scopes_json, authorized_by_member_id, policy_version, metadata,
    created_at, updated_at, disabled_at, active_authorization_key
) VALUES (
    'authz-app-optout-race', $1, 'app_github', $2,
    'app', 'agent', 'agent-1', 'disabled', 'opt out',
    '[]', $3, 1, '{"source":"default_allow_opt_out"}',
    $4, $4, $4, NULL
)`, fixture.WorkspaceID, installationID, fixture.HumanMemberID, now); err != nil {
			return err
		}
		_, err := service.applyAgentCatalogChange(ctx, tx, fixture.WorkspaceID, "agent-1", now)
		return err
	}))
	close(resumeWorker)

	select {
	case err := <-workerErr:
		require.Error(t, err, "stale missing-row snapshot must not insert an enabled v1 over the opt-out tombstone")
	case <-ctx.Done():
		t.Fatal(ctx.Err())
	}
	catalog := readCatalogStateRow(t, ctx, db, fixture.WorkspaceID, "agent-1")
	assert.Equal(t, int64(0), catalog.Version)
	assert.Equal(t, "", catalog.State)
	assert.False(t, catalog.Enabled)
	assert.Equal(t, 1, catalogRowCount(t, ctx, db, fixture.WorkspaceID, "agent-1"))
	assertUpdateCount(t, ctx, db, fixture.WorkspaceID, "agent-1", 0)

	result, err := service.reconcileWorkspaceCatalogJobBatch(ctx, fixture.WorkspaceID, "", 1, time.Now().UTC())
	require.NoError(t, err)
	assert.Equal(t, 1, result.Scanned)
	assert.Equal(t, 0, result.Changed, "retry with fresh auth sees the opt-out tombstone as a no-op")
}

// TestWorkspaceCatalogReconcileWorker_JobDrivenAndCancelled proves production
// lifecycle behavior: startup does not scan Agents without a dirty job, a
// Workspace job materializes/reconciles the Catalog, and cancelling the server
// context stops further reconciliation.
func TestWorkspaceCatalogReconcileWorker_JobDrivenAndCancelled(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	db := startServiceConcurrentDB(t, ctx)
	fixture := seedServiceFixture(t, ctx, db, "cat_reconcile_agentteam")
	require.NoError(t, registerProviderWatchBuiltinsForTest(fixture.Registry))
	service := newService(db, fixture.Registry)
	_, connectionID := seedProductAppConnection(t, db, fixture, "agentteam")

	errs := make(chan error, 1)
	service.StartWorkspaceCatalogReconcileWorker(ctx, 20*time.Millisecond, func(err error) {
		select {
		case errs <- err:
		default:
		}
	})
	require.Never(t, func() bool {
		return readCatalogHint(t, context.Background(), db, fixture.WorkspaceID, "agent-1").Version > 0
	}, 120*time.Millisecond, 20*time.Millisecond, "idle worker must not scan Agents without a Workspace dirty job")

	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return service.EnqueueWorkspaceCatalogReconcile(ctx, tx, fixture.WorkspaceID, time.Now().UTC())
	}))
	require.Eventually(t, func() bool {
		return readCatalogHint(t, context.Background(), db, fixture.WorkspaceID, "agent-1").Version == 1
	}, 5*time.Second, 20*time.Millisecond, "the immediate pass materializes v1")

	mustExecService(t, db, `UPDATE service_connections SET status='disabled', disabled_at=$2, updated_at=$2 WHERE id=$1`,
		connectionID, time.Now().UTC())
	require.NoError(t, coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return service.EnqueueWorkspaceCatalogReconcile(ctx, tx, fixture.WorkspaceID, time.Now().UTC())
	}))
	require.Eventually(t, func() bool {
		catalog := readCatalogStateRow(
			t,
			context.Background(),
			db,
			fixture.WorkspaceID,
			"agent-1",
		)
		return catalog.Version == 2 &&
			catalog.State == "[]" &&
			!catalog.Enabled
	}, 5*time.Second, 20*time.Millisecond, "a Workspace dirty job reconciles disconnect to v2")

	select {
	case err := <-errs:
		require.NoError(t, err)
	default:
	}
	cancel()
	time.Sleep(80 * time.Millisecond)
	mustExecService(t, db, `UPDATE service_connections SET status='active', disabled_at=NULL, updated_at=$2 WHERE id=$1`,
		connectionID, time.Now().UTC())
	require.NoError(t, coredb.WithTx(context.Background(), db, func(tx coredb.Tx) error {
		return service.EnqueueWorkspaceCatalogReconcile(context.Background(), tx, fixture.WorkspaceID, time.Now().UTC())
	}))
	require.Never(t, func() bool {
		return readCatalogHint(t, context.Background(), db, fixture.WorkspaceID, "agent-1").Version > 2
	}, 150*time.Millisecond, 20*time.Millisecond, "the cancelled agentteam must not reconcile v3")
}