package store

import (
	"context"
	"encoding/json"
	"errors"
	"os"
	"strings"
	"testing"
	"time"

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

func TestAcceptanceTrafficGateAndCASPromotion(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 acceptance traffic integration test")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()
	applyOIDCJITTestMigrations(t, ctx, databaseURL)
	db, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect store: %v", err)
	}
	defer db.Close()
	resetTrafficMode := func() {
		_, _ = db.pool.Exec(context.Background(), `
UPDATE system_settings
SET value = '{"mode":"live","revision":0}'::jsonb, updated_at = now()
WHERE setting_key = $1`, SystemSettingGatewayTrafficMode)
	}
	resetTrafficMode()
	t.Cleanup(resetTrafficMode)

	token := "acceptance-token-" + time.Now().Format("150405.000000000")
	run, err := db.CreateAcceptanceRun(ctx, CreateAcceptanceRunInput{
		ReleaseSHA:        strings.Repeat("a", 40),
		APIImageDigest:    "sha256:" + strings.Repeat("b", 64),
		WorkerImageDigest: "sha256:" + strings.Repeat("c", 64),
		APIKeyID:          "acceptance-api-key",
		UserID:            "acceptance-user",
		Token:             token,
		EmulatorBaseURL:   "http://acceptance-emulator:8090",
		CallbackURL:       "http://acceptance-emulator:8090/callbacks",
		CapacityProfile:   "P24",
		Config: map[string]any{"participants": []any{
			map[string]any{"apiKeyId": "acceptance-shard-key", "userId": "acceptance-shard-user"},
		}},
	})
	if err != nil {
		t.Fatalf("create acceptance run: %v", err)
	}
	encoded, _ := json.Marshal(run)
	if strings.Contains(string(encoded), token) {
		t.Fatal("acceptance run response exposed the raw token")
	}
	mode, err := db.ActivateAcceptanceRun(ctx, run.ID)
	if err != nil {
		t.Fatalf("activate acceptance run: %v", err)
	}
	if mode.Mode != "validation" || mode.RunID != run.ID || mode.Revision != 1 {
		t.Fatalf("unexpected validation mode: %+v", mode)
	}
	user := &auth.User{ID: "acceptance-user", APIKeyID: "acceptance-api-key"}
	if _, err := db.AuthorizeAcceptanceTask(ctx, "", "", user); !errors.Is(err, ErrProductionTrafficPaused) {
		t.Fatalf("production request error=%v, want traffic paused", err)
	}
	if _, err := db.AuthorizeAcceptanceTask(ctx, run.ID, "wrong-token", user); !errors.Is(err, ErrAcceptanceNotAuthorized) {
		t.Fatalf("wrong acceptance token error=%v", err)
	}
	if authorizedRunID, err := db.AuthorizeAcceptanceTask(ctx, run.ID, token, user); err != nil || authorizedRunID != run.ID {
		t.Fatalf("authorize acceptance task run=%q err=%v", authorizedRunID, err)
	}
	shardUser := &auth.User{ID: "acceptance-shard-user", APIKeyID: "acceptance-shard-key"}
	if authorizedRunID, err := db.AuthorizeAcceptanceTask(ctx, run.ID, token, shardUser); err != nil || authorizedRunID != run.ID {
		t.Fatalf("authorize acceptance shard run=%q err=%v", authorizedRunID, err)
	}
	unpairedShard := &auth.User{ID: "acceptance-shard-user", APIKeyID: "acceptance-api-key"}
	if _, err := db.AuthorizeAcceptanceTask(ctx, run.ID, token, unpairedShard); !errors.Is(err, ErrAcceptanceNotAuthorized) {
		t.Fatalf("unpaired acceptance shard error=%v, want unauthorized", err)
	}
	if baseURL, credential, err := db.AcceptanceCandidateOverride(ctx, run.ID); err != nil ||
		baseURL != "http://acceptance-emulator:8090" || credential == "" || credential == token {
		t.Fatalf("candidate override base=%q credential=%q err=%v", baseURL, credential, err)
	}
	if _, err := db.FinishAcceptanceRun(ctx, FinishAcceptanceRunInput{
		RunID: run.ID, Passed: true, Report: map[string]any{"queueFinal": 0},
	}); err != nil {
		t.Fatalf("finish acceptance run: %v", err)
	}
	promotion := PromoteAcceptanceRunInput{
		RunID: run.ID, Revision: mode.Revision,
		ReleaseSHA: mode.ReleaseSHA, APIImageDigest: mode.APIImageDigest, WorkerImageDigest: mode.WorkerImageDigest,
	}
	stale := promotion
	stale.Revision++
	if _, err := db.PromoteAcceptanceRun(ctx, stale); !errors.Is(err, ErrAcceptanceStateConflict) {
		t.Fatalf("stale promotion error=%v, want state conflict", err)
	}
	live, err := db.PromoteAcceptanceRun(ctx, promotion)
	if err != nil {
		t.Fatalf("promote acceptance run: %v", err)
	}
	if live.Mode != "live" || live.Revision != 2 {
		t.Fatalf("unexpected live mode: %+v", live)
	}

	failedRun, err := db.CreateAcceptanceRun(ctx, CreateAcceptanceRunInput{
		ReleaseSHA:        strings.Repeat("d", 40),
		APIImageDigest:    "sha256:" + strings.Repeat("e", 64),
		WorkerImageDigest: "sha256:" + strings.Repeat("f", 64),
		APIKeyID:          "acceptance-api-key",
		UserID:            "acceptance-user",
		Token:             token + "-retry",
		EmulatorBaseURL:   "http://acceptance-emulator:8090",
		CallbackURL:       "http://acceptance-emulator:8090/callbacks",
		CapacityProfile:   "P24",
	})
	if err != nil {
		t.Fatalf("create retry acceptance run: %v", err)
	}
	failedMode, err := db.ActivateAcceptanceRun(ctx, failedRun.ID)
	if err != nil {
		t.Fatalf("activate retry acceptance run: %v", err)
	}
	createAcceptanceTask := func(label string) GatewayTask {
		t.Helper()
		task, createErr := db.CreateTask(ctx, CreateTaskInput{
			Kind: "images.edits", Model: "acceptance-cleanup-" + label,
			RunMode: "acceptance", AcceptanceRunID: failedRun.ID,
			Request: map[string]any{"prompt": "acceptance cleanup integration test"},
		}, &auth.User{ID: "acceptance-user", Source: "gateway"})
		if createErr != nil {
			t.Fatalf("create acceptance cleanup task: %v", createErr)
		}
		return task
	}
	queuedTask := createAcceptanceTask("queued")
	runningTask := createAcceptanceTask("running")
	notSubmittedTask := createAcceptanceTask("not-submitted")
	submittingTask := createAcceptanceTask("submitting")
	var billingUserID string
	if err := db.pool.QueryRow(ctx, `
INSERT INTO gateway_users (user_key, username)
VALUES ($1, $1)
RETURNING id::text`, "acceptance-cleanup-"+time.Now().Format("150405.000000000")).Scan(&billingUserID); err != nil {
		t.Fatalf("create acceptance cleanup billing user: %v", err)
	}
	if _, err := db.pool.Exec(ctx, `
UPDATE gateway_tasks
SET gateway_user_id = $2::uuid,
    reservation_amount = 1,
    billing_status = 'not_started'
WHERE id = $1::uuid`, queuedTask.ID, billingUserID); err != nil {
		t.Fatalf("reserve acceptance cleanup task billing: %v", err)
	}
	if _, err := db.ClaimTaskExecution(ctx, runningTask.ID, "10000000-0000-4000-8000-000000000001", time.Minute); err != nil {
		t.Fatalf("claim safe running acceptance task: %v", err)
	}
	notSubmittedClaim, err := db.ClaimTaskExecution(ctx, notSubmittedTask.ID, "10000000-0000-4000-8000-000000000002", time.Minute)
	if err != nil {
		t.Fatalf("claim not-submitted acceptance task: %v", err)
	}
	notSubmittedAttemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{
		TaskID: notSubmittedTask.ID, ExecutionToken: notSubmittedClaim.ExecutionToken,
		AttemptNo: 1, Status: "running", Simulated: true,
	})
	if err != nil {
		t.Fatalf("create not-submitted acceptance task attempt: %v", err)
	}
	submittingClaim, err := db.ClaimTaskExecution(ctx, submittingTask.ID, "10000000-0000-4000-8000-000000000003", time.Minute)
	if err != nil {
		t.Fatalf("claim submitting acceptance task: %v", err)
	}
	submittingAttemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{
		TaskID: submittingTask.ID, ExecutionToken: submittingClaim.ExecutionToken,
		AttemptNo: 1, Status: "running", Simulated: true,
	})
	if err != nil {
		t.Fatalf("create submitting acceptance task attempt: %v", err)
	}
	if _, err := db.pool.Exec(ctx, `
UPDATE gateway_task_attempts
SET upstream_submission_status = 'submitting', upstream_submission_updated_at = now()
WHERE id = $1::uuid`, submittingAttemptID); err != nil {
		t.Fatalf("mark acceptance task attempt submitting: %v", err)
	}
	if _, err := db.FinishAcceptanceRun(ctx, FinishAcceptanceRunInput{
		RunID: failedRun.ID, Passed: false, FailureReason: "load failed",
	}); err != nil {
		t.Fatalf("fail acceptance run: %v", err)
	}
	if retried, err := db.RetryAcceptanceRun(ctx, failedRun.ID); err != nil || retried.Status != "running" {
		t.Fatalf("retry acceptance run=%+v err=%v", retried, err)
	}
	aborted, err := db.AbortAcceptanceRun(ctx, PromoteAcceptanceRunInput{
		RunID: failedRun.ID, Revision: failedMode.Revision,
		ReleaseSHA: failedMode.ReleaseSHA, APIImageDigest: failedMode.APIImageDigest,
		WorkerImageDigest: failedMode.WorkerImageDigest,
	})
	if err != nil {
		t.Fatalf("abort acceptance run: %v", err)
	}
	if aborted.Mode != "live" || aborted.Revision != failedMode.Revision+1 {
		t.Fatalf("unexpected aborted mode: %+v", aborted)
	}
	var safeCancelled, submittingRunning, notSubmittedAttempts, releaseEvents int
	if err := db.pool.QueryRow(ctx, `
SELECT
  count(*) FILTER (
    WHERE id = ANY($1::uuid[])
      AND status = 'cancelled'
      AND error_code = 'acceptance_run_aborted'
  ),
  count(*) FILTER (
    WHERE id = $2::uuid
      AND status = 'running'
  ),
  (SELECT count(*) FROM gateway_task_attempts WHERE id = $3::uuid),
  (SELECT count(*)
     FROM settlement_outbox
    WHERE task_id = $4::uuid
      AND action = 'release'
      AND status = 'pending')
FROM gateway_tasks`,
		[]string{queuedTask.ID, runningTask.ID, notSubmittedTask.ID},
		submittingTask.ID,
		notSubmittedAttemptID,
		queuedTask.ID,
	).Scan(&safeCancelled, &submittingRunning, &notSubmittedAttempts, &releaseEvents); err != nil {
		t.Fatalf("read aborted acceptance task states: %v", err)
	}
	if safeCancelled != 3 || submittingRunning != 1 || notSubmittedAttempts != 0 || releaseEvents != 1 {
		t.Fatalf(
			"abort cleanup safe_cancelled=%d submitting_running=%d not_submitted_attempts=%d release_events=%d",
			safeCancelled,
			submittingRunning,
			notSubmittedAttempts,
			releaseEvents,
		)
	}
	if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_tasks WHERE id = ANY($1::uuid[])`,
		[]string{queuedTask.ID, runningTask.ID, notSubmittedTask.ID, submittingTask.ID}); err != nil {
		t.Fatalf("delete acceptance cleanup tasks: %v", err)
	}
	if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_users WHERE id = $1::uuid`, billingUserID); err != nil {
		t.Fatalf("delete acceptance cleanup billing user: %v", err)
	}
}