package uisession

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func TestStoreKeepsActiveServerWindowLocalAndGenerationMonotonic(t *testing.T) {
	now := time.Date(2026, 7, 14, 10, 0, 0, 0, time.UTC)
	store := NewStore(Options{
		MaxSessions: 8,
		Now:         func() time.Time { return now },
	})
	expiresAt := now.Add(time.Hour)

	first, err := store.Select("window-a", expiresAt, "server-a")
	require.NoError(t, err)
	require.Equal(t, Snapshot{
		ActiveServerID: "server-a",
		Generation:     1,
		ExpiresAt:      expiresAt,
	}, first)

	second, err := store.Select("window-b", expiresAt, "server-b")
	require.NoError(t, err)
	require.Equal(t, uint64(1), second.Generation)

	switched, err := store.Select("window-a", expiresAt, "server-c")
	require.NoError(t, err)
	require.Equal(t, "server-c", switched.ActiveServerID)
	require.Equal(t, uint64(2), switched.Generation)

	windowB, ok := store.Get("window-b")
	require.True(t, ok)
	require.Equal(t, "server-b", windowB.ActiveServerID)
	require.Equal(t, uint64(1), windowB.Generation)

	cleared, ok := store.Clear("window-a")
	require.True(t, ok)
	require.Empty(t, cleared.ActiveServerID)
	require.Equal(t, uint64(3), cleared.Generation)

	reselected, err := store.Select("window-a", expiresAt, "server-a")
	require.NoError(t, err)
	require.Equal(t, uint64(4), reselected.Generation)
}

func TestStoreInvalidatesMatchingServerWithoutClobberingConcurrentSwitch(t *testing.T) {
	now := time.Date(2026, 7, 14, 10, 0, 0, 0, time.UTC)
	store := NewStore(Options{Now: func() time.Time { return now }})
	expiresAt := now.Add(time.Hour)

	windowA, err := store.Select("window-a", expiresAt, "server-a")
	require.NoError(t, err)
	_, err = store.Select("window-b", expiresAt, "server-a")
	require.NoError(t, err)
	_, err = store.Select("window-c", expiresAt, "server-b")
	require.NoError(t, err)

	switched, err := store.Select("window-a", expiresAt, "server-b")
	require.NoError(t, err)
	require.False(t, store.ClearIfGeneration("window-a", windowA.Generation))

	require.Equal(t, 1, store.InvalidateServer("server-a"))

	gotA, ok := store.Get("window-a")
	require.True(t, ok)
	require.Equal(t, switched, gotA)

	gotB, ok := store.Get("window-b")
	require.True(t, ok)
	require.Empty(t, gotB.ActiveServerID)
	require.Equal(t, uint64(2), gotB.Generation)

	gotC, ok := store.Get("window-c")
	require.True(t, ok)
	require.Equal(t, "server-b", gotC.ActiveServerID)
}

func TestStorePrunesExpiredSessionsAndCapsLiveState(t *testing.T) {
	now := time.Date(2026, 7, 14, 10, 0, 0, 0, time.UTC)
	store := NewStore(Options{
		MaxSessions: 2,
		Now:         func() time.Time { return now },
	})

	_, err := store.Select("oldest", now.Add(time.Hour), "server-a")
	require.NoError(t, err)
	now = now.Add(time.Minute)
	_, err = store.Select("middle", now.Add(time.Hour), "server-b")
	require.NoError(t, err)
	now = now.Add(time.Minute)
	_, err = store.Select("newest", now.Add(time.Hour), "server-c")
	require.NoError(t, err)

	_, ok := store.Get("oldest")
	require.False(t, ok)
	require.Equal(t, 2, store.Len())

	now = now.Add(2 * time.Hour)
	require.Equal(t, 0, store.Prune())
	require.Equal(t, 0, store.Len())
}

func TestStoreRejectsExpiredOrIncompleteSelection(t *testing.T) {
	now := time.Date(2026, 7, 14, 10, 0, 0, 0, time.UTC)
	store := NewStore(Options{Now: func() time.Time { return now }})

	_, err := store.Select("", now.Add(time.Hour), "server-a")
	require.ErrorIs(t, err, ErrSessionRequired)
	_, err = store.Select("window-a", now.Add(time.Hour), "")
	require.ErrorIs(t, err, ErrServerRequired)
	_, err = store.Select("window-a", now, "server-a")
	require.ErrorIs(t, err, ErrSessionExpired)
}

func TestStoreConcurrentSwitchesAreAtomic(t *testing.T) {
	now := time.Date(2026, 7, 14, 10, 0, 0, 0, time.UTC)
	store := NewStore(Options{Now: func() time.Time { return now }})
	expiresAt := now.Add(time.Hour)

	const switches = 128
	var wg sync.WaitGroup
	errs := make(chan error, switches)
	for i := 0; i < switches; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			_, err := store.Select("window-a", expiresAt, fmt.Sprintf("server-%d", i%4))
			errs <- err
		}(i)
	}
	wg.Wait()
	close(errs)
	for err := range errs {
		require.NoError(t, err)
	}

	got, ok := store.Get("window-a")
	require.True(t, ok)
	require.Equal(t, uint64(switches), got.Generation)
	require.Contains(t, []string{"server-0", "server-1", "server-2", "server-3"}, got.ActiveServerID)
}

func TestStoreConcurrentReadsWritesAndInvalidationStayConsistent(t *testing.T) {
	now := time.Date(2026, 7, 14, 10, 0, 0, 0, time.UTC)
	store := NewStore(Options{Now: func() time.Time { return now }})
	expiresAt := now.Add(time.Hour)
	_, err := store.Select("window-a", expiresAt, "server-a")
	require.NoError(t, err)

	const operations = 256
	var wg sync.WaitGroup
	for i := 0; i < operations; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			switch i % 4 {
			case 0:
				store.Get("window-a")
			case 1:
				_, selectErr := store.Select(
					"window-a",
					expiresAt,
					fmt.Sprintf("server-%d", i%3),
				)
				require.NoError(t, selectErr)
			case 2:
				store.InvalidateServer(fmt.Sprintf("server-%d", i%3))
			case 3:
				store.Clear("window-a")
			}
		}(i)
	}
	wg.Wait()

	got, ok := store.Get("window-a")
	require.True(t, ok)
	require.Greater(t, got.Generation, uint64(1))
	require.True(t,
		got.ActiveServerID == "" ||
			got.ActiveServerID == "server-0" ||
			got.ActiveServerID == "server-1" ||
			got.ActiveServerID == "server-2",
	)
}

func TestStoreMatchesGenerationRejectsSwitchedOrClearedSelection(t *testing.T) {
	now := time.Date(2026, 7, 14, 10, 0, 0, 0, time.UTC)
	store := NewStore(Options{Now: func() time.Time { return now }})
	expiresAt := now.Add(time.Hour)

	selected, err := store.Select("window-a", expiresAt, "server-a")
	require.NoError(t, err)
	require.True(t, store.MatchesGeneration("window-a", selected.Generation))

	switched, err := store.Select("window-a", expiresAt, "server-b")
	require.NoError(t, err)
	require.False(t, store.MatchesGeneration("window-a", selected.Generation))
	require.True(t, store.MatchesGeneration("window-a", switched.Generation))

	store.Clear("window-a")
	require.False(t, store.MatchesGeneration("window-a", switched.Generation))
}

func TestStoreAdvancesMatchingServerGenerationWithoutClearingSelection(t *testing.T) {
	now := time.Date(2026, 7, 14, 10, 0, 0, 0, time.UTC)
	store := NewStore(Options{Now: func() time.Time { return now }})
	expiresAt := now.Add(time.Hour)
	_, err := store.Select("window-a", expiresAt, "server-a")
	require.NoError(t, err)
	_, err = store.Select("window-b", expiresAt, "server-a")
	require.NoError(t, err)
	serverB, err := store.Select("window-c", expiresAt, "server-b")
	require.NoError(t, err)

	require.Equal(t, 2, store.AdvanceServerGeneration("server-a"))

	for _, sessionID := range []string{"window-a", "window-b"} {
		got, ok := store.Get(sessionID)
		require.True(t, ok)
		require.Equal(t, "server-a", got.ActiveServerID)
		require.Equal(t, uint64(2), got.Generation)
	}
	gotB, ok := store.Get("window-c")
	require.True(t, ok)
	require.Equal(t, serverB, gotB)
}

func TestStoreGenerationWatchClosesOnMutationAndCanBeReleased(t *testing.T) {
	now := time.Date(2026, 7, 14, 10, 0, 0, 0, time.UTC)
	store := NewStore(Options{Now: func() time.Time { return now }})
	expiresAt := now.Add(time.Hour)
	selected, err := store.Select("window-a", expiresAt, "server-a")
	require.NoError(t, err)

	changed, release, ok := store.WatchGeneration("window-a", selected.Generation)
	require.True(t, ok)

	store.AdvanceServerGeneration("server-a")
	select {
	case <-changed:
	case <-time.After(time.Second):
		t.Fatal("generation watcher was not closed")
	}
	release()

	current, ok := store.Get("window-a")
	require.True(t, ok)
	releasedChannel, released, ok := store.WatchGeneration("window-a", current.Generation)
	require.True(t, ok)
	released()
	store.AdvanceServerGeneration("server-a")
	select {
	case <-releasedChannel:
		t.Fatal("released generation watcher must not be retained")
	default:
	}

	_, _, ok = store.WatchGeneration("window-a", selected.Generation)
	require.False(t, ok)
}

func TestStoreServerTransitionFencesBothSidesAndRejectsIntermediateWork(t *testing.T) {
	now := time.Date(2026, 7, 14, 10, 0, 0, 0, time.UTC)
	store := NewStore(Options{Now: func() time.Time { return now }})
	expiresAt := now.Add(time.Hour)
	selected, err := store.Select("window-a", expiresAt, "server-a")
	require.NoError(t, err)

	oldGenerationChanged, releaseOldWatch, ok := store.WatchGeneration(
		"window-a",
		selected.Generation,
	)
	require.True(t, ok)
	defer releaseOldWatch()

	endTransition := store.BeginServerTransition("server-a")
	select {
	case <-oldGenerationChanged:
	case <-time.After(time.Second):
		t.Fatal("transition did not fence the old generation")
	}

	during, ok := store.Get("window-a")
	require.True(t, ok)
	require.Equal(t, uint64(2), during.Generation)
	_, _, ok = store.WatchGeneration("window-a", during.Generation)
	require.False(t, ok, "requests must not start while server identity is transitioning")

	endTransition()
	after, ok := store.Get("window-a")
	require.True(t, ok)
	require.Equal(t, "server-a", after.ActiveServerID)
	require.Equal(t, uint64(3), after.Generation)
	_, releaseAfterWatch, ok := store.WatchGeneration("window-a", after.Generation)
	require.True(t, ok)
	releaseAfterWatch()
}

func TestStoreOverlappingServerTransitionsRemainClosedUntilLastCompletion(t *testing.T) {
	now := time.Date(2026, 7, 14, 10, 0, 0, 0, time.UTC)
	store := NewStore(Options{Now: func() time.Time { return now }})
	expiresAt := now.Add(time.Hour)
	_, err := store.Select("window-a", expiresAt, "server-a")
	require.NoError(t, err)

	finishFirst := store.BeginServerTransition("server-a")
	finishSecond := store.BeginServerTransition("server-a")
	during, ok := store.Get("window-a")
	require.True(t, ok)
	require.Equal(t, uint64(2), during.Generation)

	finishFirst()
	_, _, ok = store.WatchGeneration("window-a", during.Generation)
	require.False(t, ok, "one outstanding transition must keep the server fenced")

	finishSecond()
	after, ok := store.Get("window-a")
	require.True(t, ok)
	require.Equal(t, uint64(3), after.Generation)
	_, release, ok := store.WatchGeneration("window-a", after.Generation)
	require.True(t, ok)
	release()
}

func TestStoreFailClosedInvalidatesEverySelectionAndDoesNotReopen(t *testing.T) {
	now := time.Date(2026, 7, 14, 10, 0, 0, 0, time.UTC)
	store := NewStore(Options{Now: func() time.Time { return now }})
	expiresAt := now.Add(time.Hour)
	for _, sessionID := range []string{"window-a", "window-b"} {
		_, err := store.Select(sessionID, expiresAt, "server-a")
		require.NoError(t, err)
	}

	transition := store.BeginServerTransitionHandle("server-a")
	transition.FailClosed()
	transition.Complete()

	for _, sessionID := range []string{"window-a", "window-b"} {
		snapshot, ok := store.Get(sessionID)
		require.True(t, ok)
		require.Empty(t, snapshot.ActiveServerID)
		require.Equal(t, uint64(3), snapshot.Generation)
		_, _, ok = store.WatchGeneration(sessionID, snapshot.Generation)
		require.False(t, ok)
	}
}