package store

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

	"github.com/google/uuid"
)

func TestRuntimeClientStateConcurrentUpdatesAndRecovery(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 runtime client state PostgreSQL integration tests")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 60*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)
	}

	suffix := strings.ReplaceAll(uuid.NewString(), "-", "")
	platform, err := db.CreatePlatform(ctx, CreatePlatformInput{
		Provider:    "runtime-state-test",
		PlatformKey: "runtime-state-test-" + suffix,
		Name:        "Runtime State Test " + suffix,
		AuthType:    "none",
		Status:      "enabled",
	})
	if err != nil {
		t.Fatalf("create platform: %v", err)
	}
	candidate := RuntimeModelCandidate{
		PlatformID: platform.ID,
		ClientID:   platform.PlatformKey + ":image_generate:model",
		QueueKey:   platform.PlatformKey + ":image_generate:model",
		Provider:   platform.Provider,
		ModelType:  "image_generate",
	}
	t.Cleanup(func() {
		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cleanupCancel()
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM runtime_client_states WHERE client_id = $1`, candidate.ClientID)
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM integration_platforms WHERE id = $1::uuid`, platform.ID)
	})

	runConcurrent := func(operation func() error) {
		t.Helper()
		const operations = 64
		errs := make(chan error, operations)
		var wg sync.WaitGroup
		for index := 0; index < operations; index++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				errs <- operation()
			}()
		}
		wg.Wait()
		close(errs)
		for err := range errs {
			if err != nil {
				t.Fatalf("runtime client state operation: %v", err)
			}
		}
	}
	runningCount := func() int {
		t.Helper()
		var count int
		if err := db.pool.QueryRow(ctx, `
SELECT running_count
FROM runtime_client_states
WHERE client_id = $1`, candidate.ClientID).Scan(&count); err != nil {
			t.Fatalf("read runtime client state: %v", err)
		}
		return count
	}

	runConcurrent(func() error {
		return db.RecordClientAssignment(ctx, candidate)
	})
	if count := runningCount(); count != 64 {
		t.Fatalf("running count after assignments = %d, want 64", count)
	}
	runConcurrent(func() error {
		return db.RecordClientRelease(ctx, candidate.ClientID, "")
	})
	if count := runningCount(); count != 0 {
		t.Fatalf("running count after releases = %d, want 0", count)
	}
	if _, err := db.pool.Exec(ctx, `
UPDATE runtime_client_states
SET running_count = 99
WHERE client_id = $1`, candidate.ClientID); err != nil {
		t.Fatalf("corrupt runtime client state: %v", err)
	}
	if _, err := db.RecoverInterruptedRuntimeState(ctx); err != nil {
		t.Fatalf("recover runtime client state: %v", err)
	}
	if count := runningCount(); count != 0 {
		t.Fatalf("recovered running count = %d, want 0", count)
	}
	var synchronousCommit string
	if err := db.pool.QueryRow(ctx, `SHOW synchronous_commit`).Scan(&synchronousCommit); err != nil {
		t.Fatalf("read synchronous_commit: %v", err)
	}
	if synchronousCommit != "on" {
		t.Fatalf("session synchronous_commit = %q, want on", synchronousCommit)
	}
}
