package executor

import (
	"testing"
	"time"

	"go.uber.org/zap"
)

// newTestPreparedStore builds a store with an injectable clock so TTL behavior can
// be tested deterministically without sleeping. The clock reads *now on every call.
func newTestPreparedStore(max int, idleTTL time.Duration, now *time.Time) *preparedStore {
	s := newPreparedStoreWithConfig(zap.NewNop(), max, idleTTL)
	s.nowFn = func() time.Time { return *now }
	return s
}

func testStmt(id string) *PreparedStatement {
	return &PreparedStatement{ID: id, Query: "q-" + id}
}

func mustGet(t *testing.T, s *preparedStore, id string) {
	t.Helper()
	if _, ok := s.get(id); !ok {
		t.Fatalf("expected %q present", id)
	}
}

func mustMiss(t *testing.T, s *preparedStore, id string) {
	t.Helper()
	if _, ok := s.get(id); ok {
		t.Fatalf("expected %q absent", id)
	}
}

func TestPreparedStoreEvictsLRUOverMax(t *testing.T) {
	now := time.Unix(0, 0)
	s := newTestPreparedStore(3, time.Hour, &now)

	for _, id := range []string{"a", "b", "c"} {
		now = now.Add(time.Second)
		s.putIfAbsent(testStmt(id))
	}
	// Inserting a fourth distinct statement must evict the least-recently-used ("a").
	now = now.Add(time.Second)
	s.putIfAbsent(testStmt("d"))

	if got := s.size(); got != 3 {
		t.Fatalf("expected size 3 after eviction, got %d", got)
	}
	mustMiss(t, s, "a")
	for _, id := range []string{"b", "c", "d"} {
		mustGet(t, s, id)
	}
}

func TestPreparedStoreHitRefreshesRecency(t *testing.T) {
	now := time.Unix(0, 0)
	s := newTestPreparedStore(3, time.Hour, &now)

	for _, id := range []string{"a", "b", "c"} {
		now = now.Add(time.Second)
		s.putIfAbsent(testStmt(id))
	}
	// Touch "a" so it becomes the most-recently-used; "b" is now the LRU victim.
	now = now.Add(time.Second)
	mustGet(t, s, "a")

	now = now.Add(time.Second)
	s.putIfAbsent(testStmt("d"))

	mustMiss(t, s, "b")
	for _, id := range []string{"a", "c", "d"} {
		mustGet(t, s, id)
	}
}

func TestPreparedStorePutIfAbsentReturnsExisting(t *testing.T) {
	now := time.Unix(0, 0)
	s := newTestPreparedStore(10, time.Hour, &now)

	first := testStmt("a")
	got, stored := s.putIfAbsent(first)
	if !stored || got != first {
		t.Fatalf("first insert should store and return itself")
	}
	second := testStmt("a")
	got, stored = s.putIfAbsent(second)
	if stored {
		t.Fatalf("second insert with same id should not store")
	}
	if got != first {
		t.Fatalf("should return the originally stored statement, not the duplicate")
	}
}

func TestPreparedStoreGetExpiredRemovesAndMisses(t *testing.T) {
	now := time.Unix(0, 0)
	s := newTestPreparedStore(100, 10*time.Minute, &now)
	s.putIfAbsent(testStmt("a"))

	now = now.Add(11 * time.Minute) // past the idle TTL
	if _, ok := s.get("a"); ok {
		t.Fatalf("expired statement should report a miss")
	}
	if got := s.size(); got != 0 {
		t.Fatalf("expired statement should be removed on get, size=%d", got)
	}
}

func TestPreparedStoreNotExpiredWithinTTL(t *testing.T) {
	now := time.Unix(0, 0)
	s := newTestPreparedStore(100, 10*time.Minute, &now)
	s.putIfAbsent(testStmt("a"))

	now = now.Add(9 * time.Minute)
	mustGet(t, s, "a") // refreshes lastUse
	now = now.Add(9 * time.Minute)
	mustGet(t, s, "a") // still alive because the previous get refreshed recency
}

// A re-prepare arriving before get's lazy expiry must refresh recency, otherwise the
// stale entry would keep expiring and the client would loop re-preparing forever.
func TestPreparedStorePutIfAbsentRefreshesExpiredEntry(t *testing.T) {
	now := time.Unix(0, 0)
	s := newTestPreparedStore(100, 10*time.Minute, &now)
	s.putIfAbsent(testStmt("a"))

	now = now.Add(11 * time.Minute) // "a" is expired but still present (no get yet)
	if _, stored := s.putIfAbsent(testStmt("a")); stored {
		t.Fatalf("same id should hit the existing branch, not insert")
	}
	// The hit must have refreshed recency, so the entry is no longer expired.
	mustGet(t, s, "a")
}

// After get evicts an expired entry, the client's re-prepare is a genuine fresh insert.
func TestPreparedStoreReprepareAfterExpiryIsFreshInsert(t *testing.T) {
	now := time.Unix(0, 0)
	s := newTestPreparedStore(100, 10*time.Minute, &now)
	s.putIfAbsent(testStmt("a"))

	now = now.Add(11 * time.Minute)
	mustMiss(t, s, "a") // removes the expired entry

	fresh := testStmt("a")
	got, stored := s.putIfAbsent(fresh)
	if !stored || got != fresh {
		t.Fatalf("re-prepare after expiry should be a fresh insert")
	}
	mustGet(t, s, "a")
}

func TestPreparedStoreRemove(t *testing.T) {
	s := newPreparedStoreWithConfig(zap.NewNop(), 10, time.Hour)
	s.putIfAbsent(testStmt("a"))
	s.remove("a")
	mustMiss(t, s, "a")
	if got := s.size(); got != 0 {
		t.Fatalf("expected empty store after remove, got %d", got)
	}
}

func TestPreparedStoreClear(t *testing.T) {
	s := newPreparedStoreWithConfig(zap.NewNop(), 10, time.Hour)
	for _, id := range []string{"a", "b", "c"} {
		s.putIfAbsent(testStmt(id))
	}
	s.Clear()
	if got := s.size(); got != 0 {
		t.Fatalf("expected empty store after Clear, got %d", got)
	}
	// The store must remain usable after Clear.
	if _, stored := s.putIfAbsent(testStmt("a")); !stored {
		t.Fatalf("store should accept inserts after Clear")
	}
}

func TestPreparedStoreCleanupByIdle(t *testing.T) {
	now := time.Unix(0, 0)
	s := newTestPreparedStore(100, 0, &now) // idle TTL disabled; exercise cleanup explicitly
	s.putIfAbsent(testStmt("old"))
	now = now.Add(time.Hour)
	s.putIfAbsent(testStmt("new"))

	s.cleanup(30 * time.Minute) // drop entries idle > 30m

	mustMiss(t, s, "old")
	mustGet(t, s, "new")
}

func TestPreparedStoreDefaultsClampUnboundedMax(t *testing.T) {
	if s := newPreparedStoreWithConfig(zap.NewNop(), 0, time.Hour); s.max != defaultPreparedStoreMax {
		t.Fatalf("max<=0 should clamp to default %d, got %d", defaultPreparedStoreMax, s.max)
	}
	s := newPreparedStoreWithConfig(zap.NewNop(), -5, -1)
	if s.max != defaultPreparedStoreMax {
		t.Fatalf("negative max should clamp to default, got %d", s.max)
	}
	if s.idleTTL != 0 {
		t.Fatalf("negative idle TTL should disable expiry (0), got %v", s.idleTTL)
	}
}

func TestPreparedStoreNilSafe(t *testing.T) {
	var s *preparedStore
	if _, ok := s.get("x"); ok {
		t.Fatalf("nil store get should miss")
	}
	if got, stored := s.putIfAbsent(testStmt("x")); !stored || got == nil {
		t.Fatalf("nil store putIfAbsent should pass the statement through as stored")
	}
	s.remove("x")
	s.Clear()
	s.cleanup(time.Minute)
	if got := s.size(); got != 0 {
		t.Fatalf("nil store size should be 0, got %d", got)
	}
}
