package agentops

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

	"github.com/stretchr/testify/require"

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

// startAgentOpsConcurrencyStore brings up an isolated initialized TiDB and wraps
// it in a Store with a small connection pool, mirroring the approval/app store
// concurrency tests. These prove the reserve arbitration under a real database —
// the fakes only cover logic, not transaction semantics.
func startAgentOpsConcurrencyStore(t *testing.T, ctx context.Context) *Store {
	t.Helper()
	dsn := testkit.StartInitializedTiDBWithRootDSN(t, agentOpsTiDB.RootDSN(t))
	sqlDB, err := testkit.OpenTiDBContext(ctx, dsn)
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(4)
	sqlDB.SetMaxIdleConns(4)
	t.Cleanup(func() { _ = sqlDB.Close() })
	return New(coredb.WrapSQLDB(sqlDB), Options{})
}

func receiptRowCount(t *testing.T, ctx context.Context, pool coredb.DB, applicant, request string) int64 {
	t.Helper()
	var count int64
	err := pool.QueryRow(ctx,
		`SELECT COUNT(*) FROM agent_lifecycle_receipts WHERE applicant_agent_id = $1 AND request_id = $2`,
		applicant, request).Scan(&count)
	require.NoError(t, err)
	return count
}

// Two concurrent reserves with the same (applicant, requestId) and the
// same payload digest must both succeed and converge on a single persisted
// receipt: exactly one creates, exactly one replays.
func TestReserveLifecycleReceiptConcurrentSamePayloadConverges(t *testing.T) {
	ctx := context.Background()
	store := startAgentOpsConcurrencyStore(t, ctx)

	intent := LifecycleReceiptIntent{
		WorkspaceID: "wsp-1", ApplicantAgentID: "agt-caller", OwnerUserID: "usr-owner",
		Action: LifecycleActionStart, TargetAgentID: "agt-target", RequestID: "req-same",
		PayloadDigest: "digest-shared", CorrelationID: "corr-1",
	}

	const goroutines = 2
	type reserveResult struct {
		receipt  LifecycleReceipt
		replayed bool
		err      error
	}
	results := make([]reserveResult, goroutines)

	var wg sync.WaitGroup
	wg.Add(goroutines)
	for i := 0; i < goroutines; i++ {
		go func(index int) {
			defer wg.Done()
			results[index].receipt, results[index].replayed, results[index].err =
				store.ReserveLifecycleReceipt(ctx, intent)
		}(i)
	}
	wg.Wait()

	created, replayed := 0, 0
	for _, result := range results {
		require.NoError(t, result.err, "both same-payload reserves succeed")
		require.Equal(t, "req-same", result.receipt.RequestID)
		require.Equal(t, "digest-shared", result.receipt.PayloadDigest)
		if result.replayed {
			replayed++
		} else {
			created++
		}
	}
	require.Equal(t, 1, created, "exactly one reserve creates the receipt")
	require.Equal(t, 1, replayed, "exactly one reserve replays")
	require.Equal(t, int64(1),
		receiptRowCount(t, ctx, store.db, intent.ApplicantAgentID, intent.RequestID),
		"both reserves converge on a single persisted receipt")
}

// Two concurrent reserves with the same key but different payload digests must
// produce exactly one winner and exactly one ErrLifecyclePayloadConflict, with a
// single row left behind.
func TestReserveLifecycleReceiptConcurrentDifferentPayloadConflicts(t *testing.T) {
	ctx := context.Background()
	store := startAgentOpsConcurrencyStore(t, ctx)

	base := LifecycleReceiptIntent{
		WorkspaceID: "wsp-1", ApplicantAgentID: "agt-caller", OwnerUserID: "usr-owner",
		Action: LifecycleActionStart, TargetAgentID: "agt-target", RequestID: "req-conflict",
		CorrelationID: "corr-1",
	}
	intents := []LifecycleReceiptIntent{base, base}
	intents[0].PayloadDigest = "digest-a"
	intents[1].PayloadDigest = "digest-b"

	results := make([]error, len(intents))
	var wg sync.WaitGroup
	wg.Add(len(intents))
	for i := range intents {
		go func(index int) {
			defer wg.Done()
			_, _, results[index] = store.ReserveLifecycleReceipt(ctx, intents[index])
		}(i)
	}
	wg.Wait()

	succeeded, conflicts := 0, 0
	for _, err := range results {
		switch {
		case err == nil:
			succeeded++
		case errors.Is(err, ErrLifecyclePayloadConflict):
			conflicts++
		default:
			t.Fatalf("unexpected reserve error: %v", err)
		}
	}
	require.Equal(t, 1, succeeded, "exactly one different-payload reserve wins")
	require.Equal(t, 1, conflicts, "exactly one different-payload reserve conflicts")
	require.Equal(t, int64(1),
		receiptRowCount(t, ctx, store.db, base.ApplicantAgentID, base.RequestID),
		"only one row survives the conflicting reserves")
}