package auth

import (
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"

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

// recordingTokenStore is an in-package seam that records the order of
// SessionExpiresAt / DeleteSession calls and returns a verified lease so
// the LogoutHandler can read it BEFORE deleting the token.
type recordingTokenStore struct {
	mu              sync.Mutex
	expiresAtByTok  map[string]time.Time
	callOrder       []string
	recordedExpiree time.Time
	onDelete        func()
}

func newRecordingTokenStore() *recordingTokenStore {
	return &recordingTokenStore{expiresAtByTok: map[string]time.Time{}}
}

func (r *recordingTokenStore) Create(tok string, expiresAt time.Time) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.expiresAtByTok[tok] = expiresAt
}

func (r *recordingTokenStore) SessionInfo(tok string) (token.SessionMetadata, bool) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.callOrder = append(r.callOrder, "SessionInfo")
	exp, ok := r.expiresAtByTok[tok]
	r.recordedExpiree = exp
	return token.SessionMetadata{LogicalSessionID: tok, ExpiresAt: exp}, ok
}

func (r *recordingTokenStore) DeleteSession(tok string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.callOrder = append(r.callOrder, "DeleteSession")
	if r.onDelete != nil {
		r.onDelete()
	}
	delete(r.expiresAtByTok, tok)
}

func (r *recordingTokenStore) Order() []string {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := make([]string, len(r.callOrder))
	copy(out, r.callOrder)
	return out
}

func (r *recordingTokenStore) RecordedExpiresAt() time.Time {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.recordedExpiree
}

// TestLogoutHandlerHandoffCallsRevokeWithVerifiedLeaseBeforeDeleteSession
// covers P4-BE-04c: the auth.LogoutHandler owns the handoff. It reads the
// verified expiresAt via tokenStore.SessionExpiresAt(authToken) BEFORE
// calling tokenStore.DeleteSession(authToken), then invokes the upgraded
// func(sessionID string, expiresAt time.Time) callback, which routes to
// uiSessions.Revoke(sessionID, expiresAt).
//
// On main head the callback signature is still func(string); this test
// references the upgraded signature and the new SessionExpiresAt-before-
// DeleteSession ordering, so it fails-red on the missing seam.
func TestLogoutHandlerHandoffCallsRevokeWithVerifiedLeaseBeforeDeleteSession(t *testing.T) {
	gin.SetMode(gin.TestMode)

	store := newRecordingTokenStore()
	sessionToken := token.CreateSession()
	verifiedExpiresAt := time.Now().Add(time.Hour).UTC()
	store.Create(sessionToken, verifiedExpiresAt)

	var revokedSessionID string
	var revokedExpiresAt time.Time
	var revokeCalled bool

	router := gin.New()
	// The lease-aware callback signature is func(sessionID string, expiresAt
	// time.Time). The in-package constructor allows a recording token store.
	router.POST("/logout", newLogoutHandlerWithLease(store, func(sessionID string, expiresAt time.Time) {
		require.NotContains(t, store.Order(), "DeleteSession",
			"uisession revoke must happen before deleting token state")
		revokeCalled = true
		revokedSessionID = sessionID
		revokedExpiresAt = expiresAt
	}))

	req := httptest.NewRequest(http.MethodPost, "/logout", nil)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, req)

	require.Equal(t, http.StatusOK, recorder.Code)
	require.True(t, revokeCalled, "upgraded logout callback MUST fire")

	// Token store calls remain SessionInfo → DeleteSession, while the
	// callback assertion above proves revoke ran between them.
	order := store.Order()
	require.Equal(t, []string{"SessionInfo", "DeleteSession"}, order,
		"call order MUST be SessionInfo then DeleteSession")

	require.Equal(t, sessionToken, revokedSessionID,
		"Revoke MUST receive the original authToken as sessionID")
	require.Equal(t, verifiedExpiresAt, revokedExpiresAt,
		"Revoke MUST receive the verified lease, NOT a zero time or re-derived value")
	require.False(t, verifiedExpiresAt.IsZero(),
		"verifiedExpiresAt is a real lease, not zero")

	// Second logout for the same authToken does NOT fire the callback —
	// the token is already deleted, SessionExpiresAt fails, the handler
	// short-circuits.
	revokeCalled = false
	req2 := httptest.NewRequest(http.MethodPost, "/logout", nil)
	req2.Header.Set("Authorization", "Bearer "+sessionToken)
	recorder2 := httptest.NewRecorder()
	router.ServeHTTP(recorder2, req2)

	require.False(t, revokeCalled,
		"second logout MUST NOT fire the callback — token already deleted")
}

func TestLogoutHandlerRevokesStableLogicalUISessionBeforeDeletingBearer(t *testing.T) {
	gin.SetMode(gin.TestMode)

	issue, err := token.EnsureUISession("instance-a", "window-a")
	require.NoError(t, err)
	rotated, err := token.RenewUISession(issue.Token, "instance-a", "window-a")
	require.NoError(t, err)
	store := token.SessionStore{}
	var revokedSessionID string
	var revokedExpiresAt time.Time

	router := gin.New()
	router.POST("/logout", newLogoutHandlerWithLease(store, func(sessionID string, expiresAt time.Time) {
		require.True(t, token.ValidateSession(issue.Token), "bearer must remain valid during revoke handoff")
		revokedSessionID = sessionID
		revokedExpiresAt = expiresAt
	}))

	req := httptest.NewRequest(http.MethodPost, "/logout", nil)
	req.Header.Set("Authorization", "Bearer "+rotated.Token)
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, req)

	require.Equal(t, http.StatusOK, recorder.Code)
	require.Equal(t, issue.Metadata.LogicalSessionID, revokedSessionID)
	require.Equal(t, rotated.Metadata.ExpiresAt, revokedExpiresAt)
	require.False(t, token.ValidateSession(issue.Token))
	require.False(t, token.ValidateSession(rotated.Token))
}