package control

import (
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
	"github.com/tidbcloud/AgentTeam/desktop/internal/uisession"
)

func TestUISessionControlRenewPreservesSelectionDuringBearerOverlap(t *testing.T) {
	selections := uisession.NewStore(uisession.Options{})
	control := NewUISessionControl("instance-a", selections)
	created, err := control.Create(CreateUISessionRequest{
		OwnerProof: OwnerProof{InstanceID: "instance-a"},
		WindowID:   "window-a",
	})
	require.NoError(t, err)
	createdMetadata, ok := token.UISessionInfo(created.SessionToken)
	require.True(t, ok)
	selected, err := selections.Select(createdMetadata.LogicalSessionID, created.ExpiresAt, "server-a")
	require.NoError(t, err)
	subscription, err := selections.Subscribe(createdMetadata.LogicalSessionID, created.ExpiresAt)
	require.NoError(t, err)
	generationChanged, releaseGeneration, ok := selections.WatchGeneration(
		createdMetadata.LogicalSessionID,
		selected.Generation,
	)
	require.True(t, ok)
	defer releaseGeneration()

	renewed, err := control.Renew(RenewUISessionRequest{
		OwnerProof:   OwnerProof{InstanceID: "instance-a"},
		WindowID:     "window-a",
		SessionToken: created.SessionToken,
	})
	require.NoError(t, err)
	require.NotEqual(t, created.SessionToken, renewed.SessionToken)
	require.True(t, token.ValidateSession(created.SessionToken))
	require.True(t, token.ValidateSession(renewed.SessionToken))
	renewedMetadata, ok := token.UISessionInfo(renewed.SessionToken)
	require.True(t, ok)
	selection, ok := selections.Get(renewedMetadata.LogicalSessionID)
	require.True(t, ok)
	require.Equal(t, "server-a", selection.ActiveServerID)
	require.Equal(t, uint64(1), selection.Generation)
	require.Equal(t, renewed.ExpiresAt, selection.ExpiresAt)
	_, ok = selections.Get(createdMetadata.LogicalSessionID)
	require.True(t, ok)
	renewedSubscription, err := selections.Subscribe(renewedMetadata.LogicalSessionID, renewed.ExpiresAt)
	require.NoError(t, err)
	require.Equal(t, uisession.SubscriptionStateSelected, renewedSubscription.Snapshot.State)
	require.Equal(t, "server-a", renewedSubscription.Snapshot.ActiveServerID)
	require.Equal(t, uint64(1), renewedSubscription.Snapshot.Generation)
	renewedSubscription.Release()
	_, releaseRenewedGeneration, ok := selections.WatchGeneration(
		renewedMetadata.LogicalSessionID,
		selection.Generation,
	)
	require.True(t, ok, "renewed session must admit gateway work for the preserved selection")
	releaseRenewedGeneration()
	select {
	case <-generationChanged:
		t.Fatal("lease renewal must not change the stable selection generation")
	default:
	}
	select {
	case _, open := <-subscription.Changes:
		require.True(t, open, "lease renewal must keep the logical session subscriber open")
	default:
	}
}

func TestUISessionControlWindowRevokeClosesAllWindowSubscribers(t *testing.T) {
	selections := uisession.NewStore(uisession.Options{})
	control := NewUISessionControl("instance-a", selections)
	created, err := control.Create(CreateUISessionRequest{
		OwnerProof: OwnerProof{InstanceID: "instance-a"},
		WindowID:   "window-a",
	})
	require.NoError(t, err)
	duplicate, err := control.Create(CreateUISessionRequest{
		OwnerProof: OwnerProof{InstanceID: "instance-a"},
		WindowID:   "window-a",
	})
	require.NoError(t, err)
	require.Equal(t, created.SessionToken, duplicate.SessionToken)
	metadata, ok := token.UISessionInfo(created.SessionToken)
	require.True(t, ok)
	_, err = selections.Select(metadata.LogicalSessionID, created.ExpiresAt, "server-a")
	require.NoError(t, err)
	subscription, err := selections.Subscribe(metadata.LogicalSessionID, created.ExpiresAt)
	require.NoError(t, err)
	require.NoError(t, control.Revoke(RevokeUISessionRequest{
		OwnerProof: OwnerProof{InstanceID: "instance-a"},
		WindowID:   "window-a",
	}))
	select {
	case _, open := <-subscription.Changes:
		require.False(t, open)
	case <-time.After(time.Second):
		t.Fatal("window revoke left subscriber open")
	}
}