package store

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

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

func TestQueuedAdmissionRegistrationDoesNotWaitForExecutionScopeLock(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 admission registration PostgreSQL integration tests")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*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:    "registration-lock-test",
		PlatformKey: "registration-lock-test-" + suffix,
		Name:        "Registration Lock Test " + suffix,
		AuthType:    "none",
		Status:      "enabled",
	})
	if err != nil {
		t.Fatalf("create platform: %v", err)
	}
	modelName := "registration-lock-model-" + suffix
	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, modelName).Scan(&platformModelID); err != nil {
		t.Fatalf("create platform model: %v", err)
	}
	task, err := first.CreateTask(ctx, CreateTaskInput{
		Kind:    "images.generate",
		Model:   modelName,
		RunMode: "production",
		Async:   true,
		Request: map[string]any{"prompt": "registration lock isolation"},
	}, &auth.User{ID: "registration-lock-test-" + suffix, Source: "gateway"})
	if err != nil {
		t.Fatalf("create task: %v", err)
	}
	t.Cleanup(func() {
		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cleanupCancel()
		_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = $1::uuid`, task.ID)
		_, _ = 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)
	})
	scope := AdmissionScope{
		ScopeType:       "platform_model",
		ScopeKey:        platformModelID,
		ScopeName:       modelName,
		ConcurrentLimit: 24,
		Amount:          1,
		LeaseTTLSeconds: 120,
		QueueLimit:      100,
		MaxWaitSeconds:  600,
	}
	input := TaskAdmissionInput{
		TaskID:          task.ID,
		PlatformID:      platform.ID,
		PlatformModelID: platformModelID,
		QueueKey:        platform.PlatformKey + ":image_generate:" + modelName,
		Mode:            "async",
		Priority:        100,
		Scopes:          []AdmissionScope{scope},
	}

	executionLockTx, err := first.pool.Begin(ctx)
	if err != nil {
		t.Fatalf("begin execution lock transaction: %v", err)
	}
	if err := tryAdmissionTransactionLock(ctx, executionLockTx, admissionLockKey(scope)); err != nil {
		_ = executionLockTx.Rollback(ctx)
		t.Fatalf("hold execution scope lock: %v", err)
	}
	registrationCtx, registrationCancel := context.WithTimeout(ctx, 2*time.Second)
	_, registrationErr := second.QueueTaskAdmissionWithHook(registrationCtx, input, nil)
	registrationCancel()
	if rollbackErr := executionLockTx.Rollback(ctx); rollbackErr != nil {
		t.Fatalf("release execution scope lock: %v", rollbackErr)
	}
	if registrationErr != nil {
		t.Fatalf("queue registration blocked behind execution scope lock: %v", registrationErr)
	}

	// Simulate a legacy synchronous transaction that is still holding the
	// database object lock used by NOTIFY. The admission deletion must commit
	// and become visible before the best-effort wake-up hint can acquire that
	// lock.
	notificationLockTx, err := first.pool.Begin(ctx)
	if err != nil {
		t.Fatalf("begin notification lock transaction: %v", err)
	}
	if _, err := notificationLockTx.Exec(ctx, `SELECT pg_notify($1, '*')`, admissionNotifyChannel); err != nil {
		_ = notificationLockTx.Rollback(ctx)
		t.Fatalf("hold notification object lock: %v", err)
	}
	deleteResult := make(chan error, 1)
	go func() {
		deleteResult <- second.DeleteTaskAdmission(ctx, task.ID)
	}()
	deleteVisible := false
	visibilityDeadline := time.Now().Add(time.Second)
	for time.Now().Before(visibilityDeadline) {
		var admissionCount int
		if err := first.pool.QueryRow(ctx, `
SELECT count(*)
FROM gateway_task_admissions
WHERE task_id = $1::uuid`, task.ID).Scan(&admissionCount); err != nil {
			_ = notificationLockTx.Rollback(ctx)
			t.Fatalf("read admission visibility: %v", err)
		}
		if admissionCount == 0 {
			deleteVisible = true
			break
		}
		time.Sleep(20 * time.Millisecond)
	}
	if !deleteVisible {
		_ = notificationLockTx.Rollback(ctx)
		t.Fatal("admission state commit waited behind notification object lock")
	}
	if err := notificationLockTx.Rollback(ctx); err != nil {
		t.Fatalf("release notification object lock: %v", err)
	}
	select {
	case err := <-deleteResult:
		if err != nil {
			t.Fatalf("delete admission after notification lock release: %v", err)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("delete admission did not finish after notification lock release")
	}
}
