package store

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

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

func TestDistributedTaskAdmissionQueueAcrossStores(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 task admission PostgreSQL integration tests")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	applyOIDCJITTestMigrations(t, ctx, databaseURL)

	first, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect first store: %v", err)
	}
	defer first.Close()
	second, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect second store: %v", err)
	}
	defer second.Close()
	var databaseName string
	if err := first.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
		t.Fatalf("read test database name: %v", err)
	}
	if !strings.Contains(strings.ToLower(databaseName), "test") {
		t.Fatalf("refusing to use non-test database %q", databaseName)
	}

	suffix := strings.ReplaceAll(uuid.NewString(), "-", "")
	platform, err := first.CreatePlatform(ctx, CreatePlatformInput{
		Provider:    "queue-test",
		PlatformKey: "queue-test-" + suffix,
		Name:        "Queue Test " + suffix,
		AuthType:    "none",
		Status:      "enabled",
	})
	if err != nil {
		t.Fatalf("create platform: %v", err)
	}
	var platformModelID string
	if err := first.pool.QueryRow(ctx, `
INSERT INTO platform_models (
  platform_id, model_name, provider_model_name, model_alias, model_type,
  display_name, pricing_mode, enabled
)
VALUES ($1::uuid, $2, $2, $2, '["image_generate"]'::jsonb, $2, 'inherit_discount', true)
RETURNING id::text`, platform.ID, "queue-model-"+suffix).Scan(&platformModelID); err != nil {
		t.Fatalf("create platform model: %v", err)
	}
	taskIDs := make([]string, 0, 4)
	t.Cleanup(func() {
		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cleanupCancel()
		for _, taskID := range taskIDs {
			_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = $1::uuid`, taskID)
		}
		_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM platform_models WHERE id = $1::uuid`, platformModelID)
		_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM integration_platforms WHERE id = $1::uuid`, platform.ID)
		_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM gateway_worker_instances WHERE instance_id LIKE $1`, "queue-test-"+suffix+"%")
	})

	createTask := func(async bool) GatewayTask {
		t.Helper()
		task, createErr := first.CreateTask(ctx, CreateTaskInput{
			Kind:    "images.generate",
			Model:   "queue-model-" + suffix,
			RunMode: "production",
			Async:   async,
			Request: map[string]any{"prompt": "queue admission test"},
		}, &auth.User{ID: "queue-test-user-" + suffix, Source: "gateway"})
		if createErr != nil {
			t.Fatalf("create task: %v", createErr)
		}
		taskIDs = append(taskIDs, task.ID)
		return task
	}
	scope := AdmissionScope{
		ScopeType:       "platform_model",
		ScopeKey:        platformModelID,
		ScopeName:       "queue-model",
		ConcurrentLimit: 1,
		Amount:          1,
		LeaseTTLSeconds: 137,
		QueueLimit:      2,
		MaxWaitSeconds:  600,
	}
	inputFor := func(task GatewayTask, priority int, waiterID string) TaskAdmissionInput {
		mode := "sync"
		if task.AsyncMode {
			mode = "async"
		}
		return TaskAdmissionInput{
			TaskID:          task.ID,
			PlatformID:      platform.ID,
			PlatformModelID: platformModelID,
			QueueKey:        "queue-test:" + suffix,
			Mode:            mode,
			Priority:        priority,
			WaiterID:        waiterID,
			Scopes:          []AdmissionScope{scope},
		}
	}

	running := createTask(false)
	result, err := first.TryTaskAdmission(ctx, inputFor(running, 100, "waiter-running"))
	if err != nil || !result.Admitted || len(result.Leases) != 1 {
		t.Fatalf("first task admission = %+v, err=%v", result, err)
	}
	reloadedLeases, err := second.ActiveTaskAdmissionLeases(ctx, running.ID)
	if err != nil || len(reloadedLeases) != 1 || reloadedLeases[0].TTL != 137*time.Second {
		t.Fatalf("reloaded admission leases = %+v, err=%v, want one 137s lease", reloadedLeases, err)
	}
	lowerPriority := createTask(false)
	result, err = second.TryTaskAdmission(ctx, inputFor(lowerPriority, 20, "waiter-low"))
	if err != nil || result.Admitted {
		t.Fatalf("second task should wait: result=%+v err=%v", result, err)
	}
	higherPriorityAsync := createTask(true)
	result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, ""))
	if err != nil || result.Admitted {
		t.Fatalf("third task should wait: result=%+v err=%v", result, err)
	}
	rejected := createTask(true)
	_, err = second.TryTaskAdmission(ctx, inputFor(rejected, 30, ""))
	var limitErr *RateLimitExceededError
	if !errors.As(err, &limitErr) || limitErr.Metric != "queue_size" || limitErr.Reason != "queue_full" ||
		limitErr.QueueDepth != 2 || limitErr.QueueLimit != 2 {
		t.Fatalf("fourth task error = %#v, want queue_full depth 2/2", err)
	}

	var queueCounters int
	if err := first.pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_rate_limit_counters
WHERE metric = 'queue_size'
  AND scope_key = $1`, platformModelID).Scan(&queueCounters); err != nil {
		t.Fatalf("count queue counters: %v", err)
	}
	if queueCounters != 0 {
		t.Fatalf("queue_size wrote %d fixed-window counters", queueCounters)
	}
	var attempts int
	if err := first.pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_task_attempts
WHERE task_id = ANY($1::uuid[])`, taskIDs).Scan(&attempts); err != nil {
		t.Fatalf("count attempts: %v", err)
	}
	if attempts != 0 {
		t.Fatalf("waiting admissions created %d task attempts", attempts)
	}

	if err := first.DeleteTaskAdmission(ctx, running.ID); err != nil {
		t.Fatalf("release running admission: %v", err)
	}
	result, err = second.TryTaskAdmission(ctx, inputFor(lowerPriority, 20, "waiter-low"))
	if err != nil {
		t.Fatalf("try lower-priority task after release: %v", err)
	}
	if result.Admitted {
		t.Fatal("lower-priority task bypassed higher-priority asynchronous waiter")
	}
	result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, ""))
	if err != nil || !result.Admitted {
		t.Fatalf("higher-priority async task was not admitted first: result=%+v err=%v", result, err)
	}
	for _, taskID := range taskIDs {
		_ = first.DeleteTaskAdmission(ctx, taskID)
	}

	concurrentTasks := []GatewayTask{createTask(false), createTask(true), createTask(false), createTask(true)}
	type concurrentResult struct {
		result TaskAdmissionResult
		err    error
	}
	start := make(chan struct{})
	results := make(chan concurrentResult, len(concurrentTasks))
	var waitGroup sync.WaitGroup
	for index, task := range concurrentTasks {
		waitGroup.Add(1)
		go func(index int, task GatewayTask) {
			defer waitGroup.Done()
			<-start
			target := first
			if index%2 == 1 {
				target = second
			}
			waiterID := ""
			if !task.AsyncMode {
				waiterID = "concurrent-waiter-" + task.ID
			}
			admission, admissionErr := target.TryTaskAdmission(ctx, inputFor(task, 100, waiterID))
			results <- concurrentResult{result: admission, err: admissionErr}
		}(index, task)
	}
	close(start)
	waitGroup.Wait()
	close(results)
	admittedCount := 0
	waitingCount := 0
	queueFullCount := 0
	for item := range results {
		if item.err == nil && item.result.Admitted {
			admittedCount++
			continue
		}
		if item.err == nil {
			waitingCount++
			continue
		}
		if errors.As(item.err, &limitErr) && limitErr.Metric == "queue_size" {
			queueFullCount++
			continue
		}
		t.Fatalf("unexpected concurrent admission result: %+v err=%v", item.result, item.err)
	}
	if admittedCount != 1 || waitingCount != 2 || queueFullCount != 1 {
		t.Fatalf("concurrent outcomes admitted=%d waiting=%d queue_full=%d, want 1/2/1", admittedCount, waitingCount, queueFullCount)
	}
	var activeLeases, waitingAdmissions int
	if err := first.pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_concurrency_leases
WHERE scope_type = 'platform_model'
  AND scope_key = $1
  AND released_at IS NULL
  AND expires_at > now()`, platformModelID).Scan(&activeLeases); err != nil {
		t.Fatalf("count active leases: %v", err)
	}
	if err := first.pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_task_admissions
WHERE platform_model_id = $1::uuid
  AND status = 'waiting'`, platformModelID).Scan(&waitingAdmissions); err != nil {
		t.Fatalf("count waiting admissions: %v", err)
	}
	if activeLeases != 1 || waitingAdmissions != 2 {
		t.Fatalf("database bounds active=%d waiting=%d, want 1/2", activeLeases, waitingAdmissions)
	}

	var deadlineTaskID string
	if err := first.pool.QueryRow(ctx, `
SELECT task_id::text
FROM gateway_task_admissions
WHERE platform_model_id = $1::uuid
  AND status = 'waiting'
ORDER BY task_id
LIMIT 1`, platformModelID).Scan(&deadlineTaskID); err != nil {
		t.Fatalf("select deadline waiter: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_task_admissions
SET enqueued_at = now() - interval '2 seconds',
    wait_deadline_at = now() - interval '1 second'
WHERE task_id = $1::uuid`, deadlineTaskID); err != nil {
		t.Fatalf("expire queue deadline: %v", err)
	}
	reaped, err := second.ReapExpiredTaskAdmissions(ctx, 10)
	if err != nil || reaped.ExpiredDeadlines < 1 {
		t.Fatalf("deadline reap = %+v, err=%v", reaped, err)
	}
	var taskStatus, taskErrorCode string
	if err := first.pool.QueryRow(ctx, `
SELECT status, COALESCE(error_code, '')
FROM gateway_tasks
WHERE id = $1::uuid`, deadlineTaskID).Scan(&taskStatus, &taskErrorCode); err != nil {
		t.Fatalf("read deadline task: %v", err)
	}
	if taskStatus != "failed" || taskErrorCode != "queue_timeout" {
		t.Fatalf("deadline task status=%s code=%s, want failed/queue_timeout", taskStatus, taskErrorCode)
	}

	for _, taskID := range taskIDs {
		_ = first.DeleteTaskAdmission(ctx, taskID)
	}
	leaseHolder := createTask(false)
	if result, err := first.TryTaskAdmission(ctx, inputFor(leaseHolder, 100, "lease-holder")); err != nil || !result.Admitted {
		t.Fatalf("admit lease holder: result=%+v err=%v", result, err)
	}
	disconnectedWaiter := createTask(false)
	if result, err := second.TryTaskAdmission(ctx, inputFor(disconnectedWaiter, 100, "disconnected-waiter")); err != nil || result.Admitted {
		t.Fatalf("queue disconnected waiter: result=%+v err=%v", result, err)
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_task_admissions
SET waiter_lease_expires_at = now() - interval '1 second'
WHERE task_id = $1::uuid`, disconnectedWaiter.ID); err != nil {
		t.Fatalf("expire waiter lease: %v", err)
	}
	reaped, err = second.ReapExpiredTaskAdmissions(ctx, 10)
	if err != nil || reaped.ExpiredWaiters < 1 {
		t.Fatalf("waiter reap = %+v, err=%v", reaped, err)
	}
	if err := first.pool.QueryRow(ctx, `
SELECT status, COALESCE(error_code, '')
FROM gateway_tasks
WHERE id = $1::uuid`, disconnectedWaiter.ID).Scan(&taskStatus, &taskErrorCode); err != nil {
		t.Fatalf("read disconnected task: %v", err)
	}
	if taskStatus != "cancelled" || taskErrorCode != "client_disconnected" {
		t.Fatalf("disconnected task status=%s code=%s, want cancelled/client_disconnected", taskStatus, taskErrorCode)
	}
	if err := first.DeleteTaskAdmission(ctx, leaseHolder.ID); err != nil {
		t.Fatalf("release lease holder admission: %v", err)
	}

	preSubmission := createTask(false)
	claimed, err := first.ClaimTaskPreparation(ctx, preSubmission.ID, uuid.NewString(), time.Minute)
	if err != nil {
		t.Fatalf("claim pre-submission task: %v", err)
	}
	result, err = first.TryTaskAdmission(ctx, inputFor(claimed, 100, "pre-submission"))
	if err != nil || !result.Admitted {
		t.Fatalf("admit pre-submission task: result=%+v err=%v", result, err)
	}
	var attemptID string
	if err := first.pool.QueryRow(ctx, `
INSERT INTO gateway_task_attempts (
  task_id, attempt_no, platform_id, platform_model_id, queue_key, status,
  upstream_submission_status
)
VALUES ($1::uuid, 1, $2::uuid, $3::uuid, $4, 'running', 'not_submitted')
RETURNING id::text`,
		preSubmission.ID, platform.ID, platformModelID, "queue-test:"+suffix,
	).Scan(&attemptID); err != nil {
		t.Fatalf("create pre-submission attempt: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
INSERT INTO gateway_task_param_preprocessing_logs (
  task_id, attempt_id, attempt_no, model_type
)
VALUES ($1::uuid, $2::uuid, 1, 'image_generate')`, preSubmission.ID, attemptID); err != nil {
		t.Fatalf("create pre-submission log: %v", err)
	}
	cancelled, changed, err := first.CancelTaskBeforeUpstreamSubmission(
		ctx,
		preSubmission.ID,
		claimed.ExecutionToken,
		"client disconnected before upstream submission",
	)
	if err != nil || !changed || cancelled.Status != "cancelled" || cancelled.ErrorCode != "client_disconnected" {
		t.Fatalf("cancel pre-submission task = %+v changed=%v err=%v", cancelled, changed, err)
	}
	var remainingAdmissions, remainingLeases, remainingAttempts, remainingPreprocessing int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT COUNT(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
  (SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL),
  (SELECT COUNT(*) FROM gateway_task_attempts WHERE task_id = $1::uuid),
  (SELECT COUNT(*) FROM gateway_task_param_preprocessing_logs WHERE task_id = $1::uuid)`,
		preSubmission.ID,
	).Scan(&remainingAdmissions, &remainingLeases, &remainingAttempts, &remainingPreprocessing); err != nil {
		t.Fatalf("read pre-submission cleanup: %v", err)
	}
	if remainingAdmissions != 0 || remainingLeases != 0 || remainingAttempts != 0 || remainingPreprocessing != 0 {
		t.Fatalf(
			"pre-submission cleanup admissions=%d leases=%d attempts=%d preprocessing=%d",
			remainingAdmissions,
			remainingLeases,
			remainingAttempts,
			remainingPreprocessing,
		)
	}

	postSubmission := createTask(false)
	claimed, err = first.ClaimTaskPreparation(ctx, postSubmission.ID, uuid.NewString(), time.Minute)
	if err != nil {
		t.Fatalf("claim post-submission task: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
INSERT INTO gateway_task_attempts (
  task_id, attempt_no, platform_id, platform_model_id, queue_key, status,
  upstream_submission_status
)
VALUES ($1::uuid, 1, $2::uuid, $3::uuid, $4, 'running', 'submitting')`,
		postSubmission.ID, platform.ID, platformModelID, "queue-test:"+suffix,
	); err != nil {
		t.Fatalf("create submitted attempt: %v", err)
	}
	_, changed, err = first.CancelTaskBeforeUpstreamSubmission(
		ctx,
		postSubmission.ID,
		claimed.ExecutionToken,
		"client disconnected after upstream submission",
	)
	if err != nil || changed {
		t.Fatalf("post-submission cancellation changed=%v err=%v, want unchanged", changed, err)
	}

	atomicTask := createTask(true)
	hookFailure := errors.New("synthetic admitted hook failure")
	_, err = first.TryTaskAdmissionWithAdmittedHook(ctx, inputFor(atomicTask, 100, ""), func(pgx.Tx) error {
		return hookFailure
	})
	if !errors.Is(err, hookFailure) {
		t.Fatalf("admitted hook failure = %v, want synthetic failure", err)
	}
	var atomicAdmissions, atomicLeases int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT COUNT(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
  (SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`,
		atomicTask.ID,
	).Scan(&atomicAdmissions, &atomicLeases); err != nil {
		t.Fatalf("read rolled back admitted hook state: %v", err)
	}
	if atomicAdmissions != 0 || atomicLeases != 0 {
		t.Fatalf("failed admitted hook left admissions=%d leases=%d", atomicAdmissions, atomicLeases)
	}
	const syntheticRiverJobID int64 = 987654321
	result, err = first.TryTaskAdmissionWithAdmittedHook(ctx, inputFor(atomicTask, 100, ""), func(tx pgx.Tx) error {
		_, hookErr := tx.Exec(ctx, `
UPDATE gateway_tasks
SET river_job_id = $2
WHERE id = $1::uuid`, atomicTask.ID, syntheticRiverJobID)
		return hookErr
	})
	if err != nil || !result.Admitted {
		t.Fatalf("atomic admitted hook result=%+v err=%v", result, err)
	}
	var riverJobID int64
	if err := first.pool.QueryRow(ctx, `
SELECT river_job_id
FROM gateway_tasks
WHERE id = $1::uuid`, atomicTask.ID).Scan(&riverJobID); err != nil {
		t.Fatalf("read atomic River job marker: %v", err)
	}
	if riverJobID != syntheticRiverJobID {
		t.Fatalf("atomic River job marker=%d, want %d", riverJobID, syntheticRiverJobID)
	}
	if err := first.DeleteTaskAdmission(ctx, atomicTask.ID); err != nil {
		t.Fatalf("release atomic admitted hook task: %v", err)
	}
}

func TestWorkerCapacityAllocationAndFailover(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 worker registry PostgreSQL integration tests")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	applyOIDCJITTestMigrations(t, ctx, databaseURL)
	db, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect store: %v", err)
	}
	defer db.Close()
	var databaseName string
	if err := db.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
		t.Fatalf("read test database name: %v", err)
	}
	if !strings.Contains(strings.ToLower(databaseName), "test") {
		t.Fatalf("refusing to use non-test database %q", databaseName)
	}

	prefix := "queue-test-" + strings.ReplaceAll(uuid.NewString(), "-", "")
	firstID := prefix + "-a"
	secondID := prefix + "-b"
	defer db.pool.Exec(context.Background(), `DELETE FROM gateway_worker_instances WHERE instance_id = ANY($1::text[])`, []string{firstID, secondID})

	first, err := db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
	if err != nil || first.Allocated != 5 || first.ActiveInstances != 1 {
		t.Fatalf("first allocation = %+v, err=%v", first, err)
	}
	second, err := db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: secondID, DesiredCapacity: 5})
	if err != nil || second.Allocated != 2 || second.ActiveInstances != 2 {
		t.Fatalf("second allocation = %+v, err=%v", second, err)
	}
	first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
	if err != nil || first.Allocated != 3 || first.ActiveInstances != 2 {
		t.Fatalf("rebalanced first allocation = %+v, err=%v", first, err)
	}
	if _, err := db.pool.Exec(ctx, `
UPDATE gateway_worker_instances
SET heartbeat_at = now() - interval '16 seconds'
WHERE instance_id = $1`, secondID); err != nil {
		t.Fatalf("expire second heartbeat: %v", err)
	}
	first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
	if err != nil || first.Allocated != 5 || first.ActiveInstances != 1 {
		t.Fatalf("single-worker failover allocation = %+v, err=%v", first, err)
	}
}