package store

import (
	"context"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
	"github.com/jackc/pgx/v5/pgxpool"
)

func TestExecutionPoolQueriesUseTimestampCutoff(t *testing.T) {
	db := billingV2IntegrationStore(t)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	if _, err := db.ListWorkers(ctx, time.Now()); err != nil {
		t.Fatalf("list workers with timestamp cutoff: %v", err)
	}
	if _, err := db.ListCapacity(ctx, time.Now()); err != nil {
		t.Fatalf("list capacity with timestamp cutoff: %v", err)
	}
}

func TestRouteProbeLeaseDoesNotPinDatabaseConnection(t *testing.T) {
	db := billingV2IntegrationStore(t)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	if err := db.UpsertExecutionPool(ctx, executionpool.ExecutionPool{ID: "probe-pool", State: executionpool.PoolActive}); err != nil {
		t.Fatalf("upsert execution pool: %v", err)
	}
	token, claimed, err := db.TryClaimRouteProbe(ctx, "probe-pool", "route-profile", 15*time.Second)
	if err != nil {
		t.Fatalf("claim route probe: %v", err)
	}
	if !claimed || token == "" {
		t.Fatal("route probe was not claimed")
	}

	connections := make([]*pgxpool.Conn, 0, db.Pool().Stat().MaxConns())
	defer func() {
		for _, connection := range connections {
			connection.Release()
		}
	}()
	for range db.Pool().Stat().MaxConns() {
		connection, err := db.Pool().Acquire(ctx)
		if err != nil {
			t.Fatalf("acquire all pool connections while lease is active: %v", err)
		}
		connections = append(connections, connection)
	}
	for _, connection := range connections {
		connection.Release()
	}
	connections = connections[:0]

	if err := db.ReleaseRouteProbe(ctx, "probe-pool", "route-profile", token); err != nil {
		t.Fatalf("release route probe: %v", err)
	}
}