package store

import (
	"context"
	"errors"
	"os"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
)

func TestReserveTaskBillingSerializesConcurrentWalletReservations(t *testing.T) {
	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
	if databaseURL == "" {
		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the wallet reservation integration test")
	}

	ctx := context.Background()
	db, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect store: %v", err)
	}
	defer db.Close()

	tenantID, userID := seedWalletReservationUser(t, ctx, db)
	if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
		GatewayUserID: userID,
		Currency:      "resource",
		Balance:       10,
		Reason:        "seed wallet reservation test",
	}); err != nil {
		t.Fatalf("seed wallet balance: %v", err)
	}

	firstTaskID := newWalletReservationTestUUID(t, ctx, db)
	secondTaskID := newWalletReservationTestUUID(t, ctx, db)
	billings := []any{map[string]any{"currency": "resource", "amount": float64(10)}}
	user := &auth.User{GatewayUserID: userID, GatewayTenantID: tenantID}
	tasks := []GatewayTask{
		{ID: firstTaskID, GatewayUserID: userID, GatewayTenantID: tenantID, Kind: "images.generations", Model: "mock-image"},
		{ID: secondTaskID, GatewayUserID: userID, GatewayTenantID: tenantID, Kind: "videos.generations", Model: "mock-video"},
	}

	type reserveResult struct {
		reservations []WalletBillingReservation
		err          error
	}
	results := make(chan reserveResult, len(tasks))
	var wg sync.WaitGroup
	for _, task := range tasks {
		task := task
		wg.Add(1)
		go func() {
			defer wg.Done()
			reservations, err := db.ReserveTaskBilling(ctx, task, user, billings)
			results <- reserveResult{reservations: reservations, err: err}
		}()
	}
	wg.Wait()
	close(results)

	var successReservations []WalletBillingReservation
	successCount := 0
	insufficientCount := 0
	for result := range results {
		if result.err == nil {
			successCount++
			successReservations = result.reservations
			continue
		}
		if errors.Is(result.err, ErrInsufficientWalletBalance) {
			insufficientCount++
			continue
		}
		t.Fatalf("unexpected reservation error: %v", result.err)
	}
	if successCount != 1 || insufficientCount != 1 {
		t.Fatalf("expected one successful reservation and one insufficient balance rejection, got success=%d insufficient=%d", successCount, insufficientCount)
	}
	if len(successReservations) != 1 || !walletFloatNear(successReservations[0].Amount, 10) {
		t.Fatalf("unexpected successful reservations: %+v", successReservations)
	}
	reusedReservations, err := db.ReserveTaskBilling(ctx, GatewayTask{
		ID:              successReservations[0].TaskID,
		GatewayUserID:   userID,
		GatewayTenantID: tenantID,
		Kind:            "images.generations",
		Model:           "mock-image",
	}, user, billings)
	if err != nil || len(reusedReservations) != 1 || reusedReservations[0].IdempotencyKey != successReservations[0].IdempotencyKey {
		t.Fatalf("active reservation was not reused: reservations=%+v err=%v", reusedReservations, err)
	}
	var reserveTransactions int
	if err := db.pool.QueryRow(ctx, `
SELECT count(*)
FROM gateway_wallet_transactions
WHERE reference_type='gateway_task'
  AND reference_id=$1
  AND transaction_type='reserve'`, successReservations[0].TaskID).Scan(&reserveTransactions); err != nil {
		t.Fatalf("count reused wallet reservations: %v", err)
	}
	if reserveTransactions != 1 {
		t.Fatalf("active reservation reuse created %d reserve transactions, want 1", reserveTransactions)
	}

	balance, frozen, spent := readWalletReservationAccount(t, ctx, db, userID)
	if !walletFloatNear(balance, 10) || !walletFloatNear(frozen, 10) || !walletFloatNear(spent, 0) {
		t.Fatalf("reservation should freeze balance without spending it, balance=%f frozen=%f spent=%f", balance, frozen, spent)
	}

	settleTask := GatewayTask{
		ID:                successReservations[0].TaskID,
		GatewayUserID:     userID,
		GatewayTenantID:   tenantID,
		Kind:              "images.generations",
		Model:             "mock-image",
		ResolvedModel:     "mock-image",
		Billings:          billings,
		BillingSummary:    map[string]any{"currency": "resource", "totalAmount": float64(10)},
		FinalChargeAmount: 10,
	}
	if err := db.SettleTaskBilling(ctx, settleTask); err != nil {
		t.Fatalf("settle reserved task billing: %v", err)
	}
	if err := db.SettleTaskBilling(ctx, settleTask); err != nil {
		t.Fatalf("settle reserved task billing should be idempotent: %v", err)
	}
	balance, frozen, spent = readWalletReservationAccount(t, ctx, db, userID)
	if !walletFloatNear(balance, 0) || !walletFloatNear(frozen, 0) || !walletFloatNear(spent, 10) {
		t.Fatalf("settlement should release reservation and debit once, balance=%f frozen=%f spent=%f", balance, frozen, spent)
	}
}

func TestWalletReservationLockSetSerializesSameWallet(t *testing.T) {
	var locks walletReservationLockSet
	unlock := locks.lock("wallet-user")
	acquired := make(chan struct{})
	released := make(chan struct{})
	go func() {
		defer close(released)
		unlockSecond := locks.lock("wallet-user")
		close(acquired)
		unlockSecond()
	}()
	select {
	case <-acquired:
		t.Fatal("same-wallet reservation lock was acquired concurrently")
	case <-time.After(20 * time.Millisecond):
	}
	unlock()
	select {
	case <-acquired:
	case <-time.After(time.Second):
		t.Fatal("same-wallet reservation lock did not unblock")
	}
	<-released
}

func TestCancelQueuedTaskReleasesRetainedBillingReservation(t *testing.T) {
	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
	if databaseURL == "" {
		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the queued cancellation billing integration test")
	}

	ctx := context.Background()
	db, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect store: %v", err)
	}
	defer db.Close()

	tenantID, userID := seedWalletReservationUser(t, ctx, db)
	if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
		GatewayUserID: userID,
		Currency:      "resource",
		Balance:       10,
		Reason:        "seed queued cancellation reservation test",
	}); err != nil {
		t.Fatalf("seed wallet balance: %v", err)
	}
	user := &auth.User{GatewayUserID: userID, GatewayTenantID: tenantID}
	task, err := db.CreateTask(ctx, CreateTaskInput{
		Kind: "videos.generations", Model: "queued-cancellation", RunMode: "production", Async: true,
		Request: map[string]any{"model": "queued-cancellation"},
	}, user)
	if err != nil {
		t.Fatalf("create queued task: %v", err)
	}
	t.Cleanup(func() {
		cleanupCtx := context.Background()
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM settlement_outbox WHERE task_id=$1::uuid`, task.ID)
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID)
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_transactions WHERE gateway_user_id=$1::uuid`, userID)
		_, _ = db.pool.Exec(cleanupCtx, `
DELETE FROM gateway_wallet_account_audit_guards
WHERE account_id IN (SELECT id FROM gateway_wallet_accounts WHERE gateway_user_id=$1::uuid)`, userID)
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_accounts WHERE gateway_user_id=$1::uuid`, userID)
	})

	reservations, err := db.ReserveTaskBilling(ctx, task, user, []any{map[string]any{"currency": "resource", "amount": float64(1)}})
	if err != nil || len(reservations) != 1 {
		t.Fatalf("reserve queued task billing: reservations=%+v err=%v", reservations, err)
	}
	if _, changed, err := db.CancelQueuedTask(ctx, task.ID, "cancel retained reservation"); err != nil || !changed {
		t.Fatalf("cancel queued task: changed=%v err=%v", changed, err)
	}

	claims, err := db.ClaimBillingSettlements(ctx, "queued-cancellation-test", BillingSettlementBatchSize, BillingSettlementLockTimeout)
	if err != nil {
		t.Fatalf("claim cancellation release: %v", err)
	}
	settlement := settlementForTask(t, claims, task.ID)
	if settlement.Action != "release" {
		t.Fatalf("queued cancellation settlement=%+v, want release", settlement)
	}
	if err := db.ProcessBillingSettlement(ctx, settlement); err != nil {
		t.Fatalf("process cancellation release: %v", err)
	}

	_, frozen, spent := readWalletReservationAccount(t, ctx, db, userID)
	if !walletFloatNear(frozen, 0) || !walletFloatNear(spent, 0) {
		t.Fatalf("queued cancellation did not release hold: frozen=%f spent=%f", frozen, spent)
	}
	var reserveCount int
	var releaseCount int
	if err := db.pool.QueryRow(ctx, `
SELECT count(*) FILTER (WHERE transaction_type='reserve'),
       count(*) FILTER (WHERE transaction_type='release')
FROM gateway_wallet_transactions
WHERE reference_type='gateway_task' AND reference_id=$1`, task.ID).Scan(&reserveCount, &releaseCount); err != nil {
		t.Fatalf("count queued cancellation transactions: %v", err)
	}
	if reserveCount != 1 || releaseCount != 1 {
		t.Fatalf("queued cancellation transactions reserve/release=%d/%d, want 1/1", reserveCount, releaseCount)
	}
}

func seedWalletReservationUser(t *testing.T, ctx context.Context, db *Store) (string, string) {
	t.Helper()
	suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
	var tenantID string
	if err := db.pool.QueryRow(ctx, `
INSERT INTO gateway_tenants (tenant_key, name)
VALUES ($1, $2)
RETURNING id::text`, "wallet-reservation-"+suffix, "Wallet Reservation Test "+suffix).Scan(&tenantID); err != nil {
		t.Fatalf("insert test tenant: %v", err)
	}
	var userID string
	if err := db.pool.QueryRow(ctx, `
INSERT INTO gateway_users (user_key, username, gateway_tenant_id, tenant_key, roles)
VALUES ($1, $2, $3::uuid, $4, '["basic"]'::jsonb)
RETURNING id::text`, "wallet-reservation-user-"+suffix, "wallet_reservation_"+suffix, tenantID, "wallet-reservation-"+suffix).Scan(&userID); err != nil {
		t.Fatalf("insert test user: %v", err)
	}
	t.Cleanup(func() {
		cleanupCtx := context.Background()
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_transactions WHERE gateway_user_id = $1::uuid`, userID)
		_, _ = db.pool.Exec(cleanupCtx, `
DELETE FROM gateway_wallet_account_audit_guards
WHERE account_id IN (SELECT id FROM gateway_wallet_accounts WHERE gateway_user_id = $1::uuid)`, userID)
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_users WHERE id = $1::uuid`, userID)
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tenants WHERE id = $1::uuid`, tenantID)
	})
	return tenantID, userID
}

func newWalletReservationTestUUID(t *testing.T, ctx context.Context, db *Store) string {
	t.Helper()
	var id string
	if err := db.pool.QueryRow(ctx, `SELECT gen_random_uuid()::text`).Scan(&id); err != nil {
		t.Fatalf("generate uuid: %v", err)
	}
	return id
}

func readWalletReservationAccount(t *testing.T, ctx context.Context, db *Store, userID string) (float64, float64, float64) {
	t.Helper()
	var balance float64
	var frozen float64
	var spent float64
	if err := db.pool.QueryRow(ctx, `
SELECT balance::float8, frozen_balance::float8, total_spent::float8
FROM gateway_wallet_accounts
WHERE gateway_user_id = $1::uuid
  AND currency = 'resource'`, userID).Scan(&balance, &frozen, &spent); err != nil {
		t.Fatalf("read wallet account: %v", err)
	}
	return balance, frozen, spent
}

func walletFloatNear(a float64, b float64) bool {
	delta := a - b
	if delta < 0 {
		delta = -delta
	}
	return delta < 0.000001
}