package execution

import (
	"context"
	"errors"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

func startExecutionConcurrencyStore(t *testing.T) (*Store, coredb.DB) {
	t.Helper()
	ctx := context.Background()
	dsn := testkit.StartInitializedTiDBWithRootDSN(t, roleManagementTiDB.RootDSN(t))
	sqlDB, err := testkit.OpenTiDBContext(ctx, dsn)
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(8)
	sqlDB.SetMaxIdleConns(8)
	t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
	db := coredb.WrapSQLDB(sqlDB)
	var version string
	require.NoError(t, db.QueryRow(ctx, "SELECT VERSION()").Scan(&version))
	require.Contains(t, version, "TiDB", "concurrency gate must run on real TiDB/TiKV")
	t.Logf("real TiDB concurrency gate: %s", version)
	return New(db), db
}

func TestTiDBConcurrentClaimHasSingleWinner(t *testing.T) {
	ctx := context.Background()
	store, db := startExecutionConcurrencyStore(t)
	now := time.Now().UTC().Truncate(time.Microsecond)
	suffix := testIdentifierSuffix(t.Name())
	workspaceID := "wsp_role_claim_" + suffix
	memberID := "mem_role_claim_" + suffix
	approvalID := "appr_role_claim_" + suffix
	seedExecutionWorkspaceAndApproval(t, db, workspaceID, memberID, approvalID, now)
	enqueueExecution(t, ctx, store, workspaceID, approvalID, now.Add(-time.Minute))

	start := make(chan struct{})
	results := make([][]ClaimedExecution, 2)
	errs := make([]error, 2)
	var wg sync.WaitGroup
	wg.Add(2)
	for index := range results {
		go func(index int) {
			defer wg.Done()
			<-start
			results[index], errs[index] = store.ClaimRunnable(ctx, ClaimParams{
				WorkspaceID: workspaceID,
				Now:         now,
				LeaseTTL:    time.Minute,
				Limit:       1,
			})
		}(index)
	}
	close(start)
	wg.Wait()

	winners := 0
	for index := range results {
		require.NoError(t, errs[index])
		require.LessOrEqual(t, len(results[index]), 1)
		if len(results[index]) == 1 {
			winners++
			require.Equal(t, approvalID, results[index][0].Execution.ApprovalID)
			require.NotEmpty(t, results[index][0].Fence.LeaseToken)
		}
	}
	require.Equal(t, 1, winners)

	persisted, err := store.LoadByApprovalID(ctx, workspaceID, approvalID)
	require.NoError(t, err)
	require.Equal(t, ExecutionRunning, persisted.Status)
	require.NotEmpty(t, persisted.LeaseToken)
}

func TestTiDBClaimIsTenantFenced(t *testing.T) {
	ctx := context.Background()
	store, db := startExecutionConcurrencyStore(t)
	now := time.Now().UTC().Truncate(time.Microsecond)
	suffix := testIdentifierSuffix(t.Name())
	workspace1, workspace2 := "wsp_role_tenant_a_"+suffix, "wsp_role_tenant_b_"+suffix
	member1, member2 := "mem_role_tenant_a_"+suffix, "mem_role_tenant_b_"+suffix
	approval1, approval2 := "appr_role_tenant_a_"+suffix, "appr_role_tenant_b_"+suffix
	seedExecutionWorkspaceAndApproval(t, db, workspace1, member1, approval1, now)
	seedExecutionWorkspaceAndApproval(t, db, workspace2, member2, approval2, now)

	enqueueExecution(t, ctx, store, workspace2, approval2, now.Add(-time.Minute))
	none, err := store.ClaimRunnable(ctx, ClaimParams{
		WorkspaceID: workspace1, Now: now, LeaseTTL: time.Minute, Limit: 1,
	})
	require.NoError(t, err)
	require.Empty(t, none)

	enqueueExecution(t, ctx, store, workspace1, approval1, now.Add(-time.Minute))
	for _, expected := range []struct {
		workspaceID string
		approvalID  string
	}{
		{workspaceID: workspace1, approvalID: approval1},
		{workspaceID: workspace2, approvalID: approval2},
	} {
		claimed, err := store.ClaimRunnable(ctx, ClaimParams{
			WorkspaceID: expected.workspaceID, Now: now, LeaseTTL: time.Minute, Limit: 1,
		})
		require.NoError(t, err)
		require.Len(t, claimed, 1)
		require.Equal(t, expected.workspaceID, claimed[0].Execution.WorkspaceID)
		require.Equal(t, expected.approvalID, claimed[0].Execution.ApprovalID)
	}
}

func TestTiDBExpiredReclaimReplacesFenceAndRejectsOldToken(t *testing.T) {
	ctx := context.Background()
	store, db := startExecutionConcurrencyStore(t)
	now := time.Now().UTC().Truncate(time.Microsecond)
	suffix := testIdentifierSuffix(t.Name())
	workspaceID := "wsp_role_reclaim_" + suffix
	memberID := "mem_role_reclaim_" + suffix
	approvalID := "appr_role_reclaim_" + suffix
	seedExecutionWorkspaceAndApproval(t, db, workspaceID, memberID, approvalID, now)
	enqueueExecution(t, ctx, store, workspaceID, approvalID, now.Add(-time.Minute))

	first, err := store.ClaimRunnable(ctx, ClaimParams{
		WorkspaceID: workspaceID, Now: now, LeaseTTL: time.Minute, Limit: 1,
	})
	require.NoError(t, err)
	require.Len(t, first, 1)
	reclaimedAt := now.Add(time.Minute + time.Second)
	second, err := store.ClaimRunnable(ctx, ClaimParams{
		WorkspaceID: workspaceID, Now: reclaimedAt, LeaseTTL: time.Minute, Limit: 1,
	})
	require.NoError(t, err)
	require.Len(t, second, 1)
	require.NotEqual(t, first[0].Fence.LeaseToken, second[0].Fence.LeaseToken)

	_, err = store.Succeed(ctx, first[0].Fence, Success{RoleID: "role_stale"}, reclaimedAt)
	require.True(t, errors.Is(err, ErrConflict))
	_, err = store.FailDeterministic(ctx, first[0].Fence, Failure{Code: "stale"}, reclaimedAt)
	require.True(t, errors.Is(err, ErrConflict))
	_, err = store.ReleaseUncertain(ctx, first[0].Fence, reclaimedAt)
	require.True(t, errors.Is(err, ErrConflict))

	succeeded, err := store.Succeed(ctx, second[0].Fence, Success{
		RoleID: "role_current", RoleVersionID: "rlv_current",
	}, reclaimedAt.Add(time.Second))
	require.NoError(t, err)
	require.Equal(t, ExecutionSucceeded, succeeded.Status)
	require.Equal(t, "role_current", succeeded.ResultRoleID)
}