// Package token provides API key generation and validation using Ed25519 asymmetric signing.
//
// Keys are generated in the format: <prefix>.<payload>.<signature>
//
//   - prefix: fixed identifier "sk-dev"
//   - payload: base64url-encoded creation timestamp (nanoseconds)
//   - signature: base64url-encoded Ed25519 signature of the payload
//
// Keys expire after one week. The key pair is stored in the security section
// of the main config file and is auto-generated on first use.
package token

import (
	"crypto/ed25519"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"

	"github.com/tidbcloud/AgentTeam/desktop/internal/crypto"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
)

const tokenAnalysisPayloadFingerprintDomain = "agentteam-token-analysis-task-payload/v1"

// Prefix is the fixed identifier at the beginning of every generated API key.
const Prefix = "sk-dev"

// TempKeyPrefix identifies short-lived local CLI control credentials.
const TempKeyPrefix = "sk-tmp"

const maxDuration = 7 * 24 * time.Hour

const tempMaxDuration = 5 * time.Minute

const sessionMaxDuration = 30 * 24 * time.Hour
const uiSessionMaxDuration = 15 * time.Minute
const uiSessionRecoveryDuration = 24 * time.Hour
const maxSessionKeys = 256
const agentTokenMaxDuration = 30 * 24 * time.Hour
const maxAgentTokens = 512

// AgentTokenScopeTokenAnalysis marks read-only self-scoped Token analysis
// credentials injected into runtime processes.
const AgentTokenScopeTokenAnalysis = "token_analysis"

const agentTokenScopeAny = "\x00any"

var (
	// ErrInvalidFormat indicates the key does not conform to the expected prefix.payload.signature format.
	ErrInvalidFormat = errors.New("invalid api key format")
	// ErrExpired indicates the key was created more than one week ago.
	ErrExpired = errors.New("api key expired")
	// ErrInvalidSignature indicates the key signature does not match the payload.
	ErrInvalidSignature = errors.New("invalid api key signature")
	// ErrAgentMismatch indicates the token was signed for a different agent.
	ErrAgentMismatch = errors.New("agent id mismatch")
	// ErrAgentTokenInactive indicates the token is signed but is not an active issued token.
	ErrAgentTokenInactive = errors.New("agent token inactive")
	// ErrAgentTokenRevoked indicates the token was explicitly revoked.
	ErrAgentTokenRevoked = errors.New("agent token revoked")
	// ErrAgentTokenServerMismatch indicates the token was issued for a different MCP server route.
	ErrAgentTokenServerMismatch = errors.New("agent token server mismatch")
	// ErrAgentTokenScopeMismatch indicates the token was issued for a different MCP route scope.
	ErrAgentTokenScopeMismatch = errors.New("agent token scope mismatch")
)

var (
	defaultManager *KeyManager
	defaultOnce    sync.Once
)

// Default returns the global singleton KeyManager.
// It initializes on first call by loading or generating the key pair
// from the main config file. It panics if initialization fails.
func Default() *KeyManager {
	defaultOnce.Do(func() {
		m, err := NewKeyManager("")
		if err != nil {
			panic(fmt.Sprintf("apikey: initialize default manager: %v", err))
		}
		defaultManager = m
	})
	return defaultManager
}

// KeyManager manages the Ed25519 key pair used to sign and verify API keys.
// It is safe for concurrent use.
type KeyManager struct {
	mu          sync.RWMutex
	pubKey      ed25519.PublicKey
	priKey      ed25519.PrivateKey
	agentTokens map[string]agentTokenRecord
	execTokens  map[string]execTokenRecord
	now         func() time.Time
}

// NewKeyManager creates a KeyManager by loading or generating an Ed25519 key pair
// from the security section of the config file at cfgPath. If cfgPath is empty,
// the default config path is used. If no keys exist in the config, a new
// key pair is generated and saved automatically.
func NewKeyManager(cfgPath string) (*KeyManager, error) {
	m := &KeyManager{now: time.Now}

	if err := m.load(cfgPath); err != nil {
		return nil, fmt.Errorf("load api key pair: %w", err)
	}

	return m, nil
}

func resolvePath(cfgPath string) string {
	if cfgPath == "" {
		return profile.Path()
	}
	return cfgPath
}

// load reads the key pair from the security section of the config file.
// If the keys do not exist, it generates a new key pair and saves them.
func (m *KeyManager) load(cfgPath string) error {
	p := resolvePath(cfgPath)
	cfg, err := profile.Load(p)
	if err != nil {
		return fmt.Errorf("read config: %w", err)
	}

	if len(cfg.Security.PublicKey) == ed25519.PublicKeySize && len(cfg.Security.PrivateKey) == ed25519.PrivateKeySize {
		return m.installKeys(ed25519.PublicKey(cfg.Security.PublicKey), ed25519.PrivateKey(cfg.Security.PrivateKey))
	}

	return m.generate(cfgPath)
}

func (m *KeyManager) installKeys(pubKey ed25519.PublicKey, priKey ed25519.PrivateKey) error {
	if err := crypto.SetKeys(pubKey, priKey); err != nil {
		return err
	}
	m.mu.Lock()
	m.pubKey = pubKey
	m.priKey = priKey
	m.mu.Unlock()
	return nil
}

// generate creates a new Ed25519 key pair and saves it to the config file.
func (m *KeyManager) generate(cfgPath string) error {
	p := resolvePath(cfgPath)
	cfg, err := profile.Load(p)
	if err != nil {
		return fmt.Errorf("read config: %w", err)
	}
	if len(cfg.Security.PublicKey) == ed25519.PublicKeySize && len(cfg.Security.PrivateKey) == ed25519.PrivateKeySize {
		return m.installKeys(ed25519.PublicKey(cfg.Security.PublicKey), ed25519.PrivateKey(cfg.Security.PrivateKey))
	}

	pubKey, priKey, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return fmt.Errorf("generate ed25519 key pair: %w", err)
	}

	cfg.Security = profile.Security{
		PublicKey:  profile.Base64Bytes(pubKey),
		PrivateKey: profile.Base64Bytes(priKey),
	}

	if err := profile.Save(cfg, p); err != nil {
		return fmt.Errorf("save config: %w", err)
	}
	saved, err := profile.Load(p)
	if err != nil {
		return fmt.Errorf("read saved config: %w", err)
	}
	if len(saved.Security.PublicKey) == ed25519.PublicKeySize && len(saved.Security.PrivateKey) == ed25519.PrivateKeySize {
		pubKey = ed25519.PublicKey(saved.Security.PublicKey)
		priKey = ed25519.PrivateKey(saved.Security.PrivateKey)
	}

	return m.installKeys(pubKey, priKey)
}

// Generate creates a new API key signed with the manager's private key.
// The key format is: sk-dev.<base64url timestamp>.<base64url signature>.
// The returned key is valid for up to one week from the time of generation.
func (m *KeyManager) Generate() (string, error) {
	m.mu.RLock()
	priKey := m.priKey
	m.mu.RUnlock()

	if priKey == nil {
		return "", errors.New("api key manager not initialized")
	}

	ts := time.Now().UnixNano()
	payload := fmt.Sprintf("%d", ts)
	sig := ed25519.Sign(priKey, []byte(payload))

	key := fmt.Sprintf("%s.%s.%s",
		Prefix,
		base64.RawURLEncoding.EncodeToString([]byte(payload)),
		base64.RawURLEncoding.EncodeToString(sig),
	)

	return key, nil
}

// FingerprintTokenAnalysisPayload returns a privacy-safe, machine-stable
// fingerprint for controlled Token comparison. The persisted signing key is
// used only as HMAC key material; the payload and key never enter the report.
func (m *KeyManager) FingerprintTokenAnalysisPayload(payload []byte) string {
	if len(payload) == 0 {
		return ""
	}
	m.mu.RLock()
	key := append([]byte(nil), m.priKey...)
	m.mu.RUnlock()
	if len(key) == 0 {
		return ""
	}
	mac := hmac.New(sha256.New, key)
	_, _ = mac.Write([]byte(tokenAnalysisPayloadFingerprintDomain + "\x00"))
	_, _ = mac.Write(payload)
	return hex.EncodeToString(mac.Sum(nil)[:16])
}

// Validate checks an API key for correct format, valid Ed25519 signature,
// and non-expired creation time (within one week).
func (m *KeyManager) Validate(key string) error {
	m.mu.RLock()
	pubKey := m.pubKey
	m.mu.RUnlock()

	if pubKey == nil {
		return errors.New("api key manager not initialized")
	}

	payload, sig, err := parseKey(key)
	if err != nil {
		return err
	}

	if !ed25519.Verify(pubKey, payload, sig) {
		return ErrInvalidSignature
	}

	ts, err := parseTimestamp(payload)
	if err != nil {
		return fmt.Errorf("%w: %v", ErrInvalidFormat, err)
	}

	created := time.Unix(0, ts)
	if time.Since(created) > maxDuration {
		return ErrExpired
	}

	return nil
}

// GenerateTempKey creates a short-lived local CLI control credential (5 minute expiry).
// Format: sk-tmp.<base64url timestamp>.<base64url signature>
func (m *KeyManager) GenerateTempKey() (string, error) {
	m.mu.RLock()
	priKey := m.priKey
	m.mu.RUnlock()

	if priKey == nil {
		return "", errors.New("api key manager not initialized")
	}

	ts := time.Now().UnixNano()
	payload := fmt.Sprintf("%d", ts)
	sig := ed25519.Sign(priKey, []byte(payload))

	key := fmt.Sprintf("%s.%s.%s",
		TempKeyPrefix,
		base64.RawURLEncoding.EncodeToString([]byte(payload)),
		base64.RawURLEncoding.EncodeToString(sig),
	)

	return key, nil
}

// ValidateTempKey checks a local CLI control credential for correct format,
// valid signature, and non-expired creation time (within 5 minutes).
func (m *KeyManager) ValidateTempKey(key string) error {
	m.mu.RLock()
	pubKey := m.pubKey
	m.mu.RUnlock()

	if pubKey == nil {
		return errors.New("api key manager not initialized")
	}

	if !strings.HasPrefix(key, TempKeyPrefix+".") {
		return ErrInvalidFormat
	}

	payload, sig, err := parseKeyWithPrefix(key, TempKeyPrefix)
	if err != nil {
		return err
	}

	if !ed25519.Verify(pubKey, payload, sig) {
		return ErrInvalidSignature
	}

	ts, err := parseTimestamp(payload)
	if err != nil {
		return fmt.Errorf("%w: %v", ErrInvalidFormat, err)
	}

	created := time.Unix(0, ts)
	if time.Since(created) > tempMaxDuration {
		return ErrExpired
	}

	return nil
}

// parseKey splits a key string into its payload and signature components.
// Expected format: <prefix>.<base64url payload>.<base64url signature>.
func parseKey(key string) (payload []byte, sig []byte, err error) {
	return parseKeyWithPrefix(key, Prefix)
}

// parseKeyWithPrefix splits a token string using the given prefix.
// Expected format: <prefix>.<base64url payload>.<base64url signature>.
func parseKeyWithPrefix(key string, prefix string) (payload []byte, sig []byte, err error) {
	if len(key) < len(prefix)+2 || !strings.HasPrefix(key, prefix+".") {
		return nil, nil, ErrInvalidFormat
	}

	rest := key[len(prefix):]
	rest = rest[1:]

	idx := -1
	for i, c := range rest {
		if c == '.' {
			idx = i
			break
		}
	}
	if idx < 0 {
		return nil, nil, ErrInvalidFormat
	}

	encodedPayload := rest[:idx]
	payload, err = base64.RawURLEncoding.DecodeString(encodedPayload)
	if err != nil {
		return nil, nil, fmt.Errorf("%w: invalid payload encoding", ErrInvalidFormat)
	}
	if base64.RawURLEncoding.EncodeToString(payload) != encodedPayload {
		return nil, nil, fmt.Errorf("%w: non-canonical payload encoding", ErrInvalidFormat)
	}

	encodedSignature := rest[idx+1:]
	sig, err = base64.RawURLEncoding.DecodeString(encodedSignature)
	if err != nil {
		return nil, nil, fmt.Errorf("%w: invalid signature encoding", ErrInvalidFormat)
	}
	if base64.RawURLEncoding.EncodeToString(sig) != encodedSignature {
		return nil, nil, fmt.Errorf("%w: non-canonical signature encoding", ErrInvalidFormat)
	}

	return payload, sig, nil
}

// parseTimestamp extracts a nanosecond Unix timestamp from the raw payload bytes.
func parseTimestamp(payload []byte) (int64, error) {
	var ts int64
	_, err := fmt.Sscanf(string(payload), "%d", &ts)
	return ts, err
}

// agentTokenPrefix is the fixed identifier at the beginning of every agent token.
const agentTokenPrefix = "at-dev"

type agentTokenPayload struct {
	ID        string `json:"id"`
	AgentID   string `json:"agent_id"`
	ServerID  string `json:"server_id,omitempty"`
	Scope     string `json:"scope,omitempty"`
	CreatedAt int64  `json:"created_at"`
}

type agentTokenRecord struct {
	agentID   string
	serverID  string
	scope     string
	createdAt time.Time
	expiresAt time.Time
	revokedAt time.Time
}

// AgentTokenInfo describes the active token record matched during validation.
type AgentTokenInfo struct {
	AgentID  string
	ServerID string
	Scope    string
}

// GenerateAgentTokenForServer creates a signed active agent token bound to a
// specific MCP server route. Server-bound tokens are still valid for leader-only MCP routes for the same
// agent, but server MCP validation rejects them on any
// other server route.
func (m *KeyManager) GenerateAgentTokenForServer(agentID, serverID string) (string, error) {
	return m.generateAgentToken(agentID, serverID, "")
}

func (m *KeyManager) GenerateTokenAnalysisAgentToken(agentID string) (string, error) {
	return m.generateAgentToken(agentID, "", AgentTokenScopeTokenAnalysis)
}

func (m *KeyManager) generateAgentToken(agentID, serverID, scope string) (string, error) {
	agentID = strings.TrimSpace(agentID)
	if agentID == "" {
		return "", errors.New("agent id is required")
	}
	serverID = strings.TrimSpace(serverID)
	scope = strings.TrimSpace(scope)

	now := m.currentTime()
	payload, err := json.Marshal(agentTokenPayload{
		ID:        uuid.NewString(),
		AgentID:   agentID,
		ServerID:  serverID,
		Scope:     scope,
		CreatedAt: now.UnixNano(),
	})
	if err != nil {
		return "", fmt.Errorf("marshal agent token payload: %w", err)
	}

	m.mu.Lock()
	defer m.mu.Unlock()

	if m.priKey == nil {
		return "", errors.New("api key manager not initialized")
	}

	sig := ed25519.Sign(m.priKey, payload)
	token := fmt.Sprintf("%s.%s.%s",
		agentTokenPrefix,
		base64.RawURLEncoding.EncodeToString(payload),
		base64.RawURLEncoding.EncodeToString(sig),
	)

	m.pruneAgentTokensLocked(now)
	if m.agentTokens == nil {
		m.agentTokens = make(map[string]agentTokenRecord)
	}
	m.agentTokens[agentTokenHash(token)] = agentTokenRecord{
		agentID:   agentID,
		serverID:  serverID,
		scope:     scope,
		createdAt: now,
		expiresAt: now.Add(agentTokenMaxDuration),
	}
	m.enforceAgentTokenLimitLocked()

	return token, nil
}

// ValidateAgentTokenInfo verifies a signed active ordinary Agent token for
// generic routes. Read-only Token-analysis credentials are intentionally
// excluded so they cannot inherit MCP or tool capabilities.
func (m *KeyManager) ValidateAgentTokenInfo(agentID, token string) (AgentTokenInfo, error) {
	info, err := m.validateAgentToken(agentID, "", token, false, agentTokenScopeAny)
	if err != nil {
		return AgentTokenInfo{}, err
	}
	if info.Scope == AgentTokenScopeTokenAnalysis {
		return AgentTokenInfo{}, ErrAgentTokenScopeMismatch
	}
	return info, nil
}

// ValidateAgentTokenForServerInfo verifies a signed active agent token against
// a concrete server MCP route and returns the matched token record.
func (m *KeyManager) ValidateAgentTokenForServerInfo(agentID, serverID, token string) (AgentTokenInfo, error) {
	info, err := m.validateAgentToken(agentID, serverID, token, true, "")
	if err != nil {
		return AgentTokenInfo{}, err
	}
	if info.Scope == AgentTokenScopeTokenAnalysis {
		return AgentTokenInfo{}, ErrAgentTokenScopeMismatch
	}
	return info, nil
}

// ValidateUnscopedAgentTokenInfo verifies a token for full AgentTeam MCP routes
// and returns the matched token record.
func (m *KeyManager) ValidateUnscopedAgentTokenInfo(agentID, token string) (AgentTokenInfo, error) {
	return m.validateAgentToken(agentID, "", token, false, "")
}

func (m *KeyManager) ValidateTokenAnalysisAgentTokenInfo(agentID, token string) (AgentTokenInfo, error) {
	return m.validateAgentToken(agentID, "", token, false, AgentTokenScopeTokenAnalysis)
}

// RevokeAgentToken marks a previously issued agent token inactive.
func (m *KeyManager) RevokeAgentToken(token string) bool {
	token = strings.TrimSpace(token)
	if token == "" {
		return false
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.agentTokens == nil {
		return false
	}
	key := agentTokenHash(token)
	record, ok := m.agentTokens[key]
	if !ok {
		return false
	}
	if record.revokedAt.IsZero() {
		record.revokedAt = m.currentTime()
		m.agentTokens[key] = record
	}
	return true
}

// parseAgentToken splits an agent token into its payload and signature components.
// Expected format: at-dev.<base64url payload>.<base64url signature>.
func parseAgentToken(token string) (payload []byte, sig []byte, err error) {
	return parseKeyWithPrefix(token, agentTokenPrefix)
}

func (m *KeyManager) validateAgentToken(agentID, serverID, token string, requireServer bool, requiredScope string) (AgentTokenInfo, error) {
	agentID = strings.TrimSpace(agentID)
	if agentID == "" {
		return AgentTokenInfo{}, errors.New("agent id is required")
	}
	serverID = strings.TrimSpace(serverID)
	if requireServer && serverID == "" {
		return AgentTokenInfo{}, errors.New("server id is required")
	}

	payloadBytes, sig, err := parseAgentToken(token)
	if err != nil {
		return AgentTokenInfo{}, err
	}

	m.mu.RLock()
	pubKey := m.pubKey
	m.mu.RUnlock()

	if pubKey == nil {
		return AgentTokenInfo{}, errors.New("api key manager not initialized")
	}
	if !ed25519.Verify(pubKey, payloadBytes, sig) {
		return AgentTokenInfo{}, ErrInvalidSignature
	}

	var payload agentTokenPayload
	if err := json.Unmarshal(payloadBytes, &payload); err != nil {
		return AgentTokenInfo{}, fmt.Errorf("%w: invalid agent token payload", ErrInvalidFormat)
	}
	if strings.TrimSpace(payload.AgentID) != agentID {
		return AgentTokenInfo{}, ErrAgentMismatch
	}
	payload.Scope = strings.TrimSpace(payload.Scope)

	now := m.currentTime()
	m.mu.Lock()
	defer m.mu.Unlock()

	key := agentTokenHash(token)
	record, ok := m.agentTokens[key]
	if !ok {
		return AgentTokenInfo{}, ErrAgentTokenInactive
	}
	if !record.revokedAt.IsZero() {
		return AgentTokenInfo{}, ErrAgentTokenRevoked
	}
	if now.After(record.expiresAt) {
		delete(m.agentTokens, key)
		return AgentTokenInfo{}, ErrExpired
	}
	m.pruneAgentTokensLocked(now)
	if record.agentID != agentID {
		return AgentTokenInfo{}, ErrAgentMismatch
	}
	if record.serverID != strings.TrimSpace(payload.ServerID) {
		return AgentTokenInfo{}, ErrInvalidSignature
	}
	if record.scope != payload.Scope {
		return AgentTokenInfo{}, ErrInvalidSignature
	}
	if requireServer && record.serverID != serverID {
		return AgentTokenInfo{}, ErrAgentTokenServerMismatch
	}
	if requiredScope != agentTokenScopeAny && strings.TrimSpace(requiredScope) != record.scope {
		return AgentTokenInfo{}, ErrAgentTokenScopeMismatch
	}

	return AgentTokenInfo{
		AgentID:  record.agentID,
		ServerID: record.serverID,
		Scope:    record.scope,
	}, nil
}

func agentTokenHash(token string) string {
	sum := sha256.Sum256([]byte(strings.TrimSpace(token)))
	return base64.RawURLEncoding.EncodeToString(sum[:])
}

func (m *KeyManager) currentTime() time.Time {
	if m.now != nil {
		return m.now()
	}
	return time.Now()
}

func (m *KeyManager) pruneAgentTokensLocked(now time.Time) {
	for key, record := range m.agentTokens {
		if record.expiresAt.IsZero() || now.After(record.expiresAt) {
			delete(m.agentTokens, key)
		}
	}
}

func (m *KeyManager) enforceAgentTokenLimitLocked() {
	for len(m.agentTokens) > maxAgentTokens {
		var oldestKey string
		var oldestCreated time.Time
		for key, record := range m.agentTokens {
			if record.revokedAt.IsZero() {
				continue
			}
			if oldestKey == "" || record.createdAt.Before(oldestCreated) {
				oldestKey = key
				oldestCreated = record.createdAt
			}
		}
		if oldestKey != "" {
			delete(m.agentTokens, oldestKey)
			continue
		}
		for key, record := range m.agentTokens {
			if oldestKey == "" || record.createdAt.Before(oldestCreated) {
				oldestKey = key
				oldestCreated = record.createdAt
			}
		}
		if oldestKey == "" {
			return
		}
		delete(m.agentTokens, oldestKey)
	}
}

var (
	sessionMu   sync.RWMutex
	sessionKeys = make(map[string]sessionRecord)
)

type sessionRecord struct {
	CreatedAt        time.Time
	ExpiresAt        time.Time
	LogicalExpiresAt time.Time
	InstanceID       string
	WindowID         string
	LogicalSessionID string
	Generation       uint64
}

// CreateSession generates a new session token and stores it in memory.
func CreateSession() string {
	tok := uuid.New().String()
	now := time.Now()
	sessionMu.Lock()
	pruneSessionsLocked(now)
	sessionKeys[tok] = sessionRecord{
		CreatedAt: now,
		ExpiresAt: now.Add(sessionMaxDuration),
	}
	enforceSessionLimitLocked()
	sessionMu.Unlock()
	return tok
}

// ValidateSession checks whether a session token exists and has not expired.
func ValidateSession(tok string) bool {
	_, ok := SessionExpiresAt(tok)
	return ok
}

// SessionExpiresAt returns the validated session's absolute expiration time.
// Expired sessions are removed lazily before returning.
func SessionExpiresAt(tok string) (time.Time, bool) {
	metadata, ok := SessionInfo(tok)
	return metadata.ExpiresAt, ok
}

// SessionMetadata is the validated identity carried by one bearer token.
// Legacy long-lived sessions use the token itself as their logical identity;
// renderer UI leases use a stable daemon-generated logical session ID.
type SessionMetadata struct {
	LogicalSessionID string
	ExpiresAt        time.Time
}

// SessionInfo returns the validated logical identity and absolute expiration
// carried by a session token. Expired sessions are removed lazily.
func SessionInfo(tok string) (SessionMetadata, bool) {
	tok = strings.TrimSpace(tok)
	now := time.Now()
	sessionMu.Lock()
	defer sessionMu.Unlock()
	record, ok := sessionKeys[tok]
	if !ok || record.CreatedAt.IsZero() || record.ExpiresAt.IsZero() || !record.ExpiresAt.After(now) {
		if ok && shouldPruneSessionRecord(record, now) {
			delete(sessionKeys, tok)
		}
		return SessionMetadata{}, false
	}
	logicalSessionID := strings.TrimSpace(record.LogicalSessionID)
	if logicalSessionID == "" {
		logicalSessionID = tok
	}
	return SessionMetadata{
		LogicalSessionID: logicalSessionID,
		ExpiresAt:        logicalSessionExpiresAt(record),
	}, true
}

// DeleteSession removes a legacy token or every bearer generation belonging to
// the same logical renderer UI session.
func DeleteSession(tok string) {
	sessionMu.Lock()
	record, ok := sessionKeys[tok]
	if !ok || strings.TrimSpace(record.LogicalSessionID) == "" {
		delete(sessionKeys, tok)
		sessionMu.Unlock()
		return
	}
	for candidate, current := range sessionKeys {
		if current.LogicalSessionID == record.LogicalSessionID {
			delete(sessionKeys, candidate)
		}
	}
	sessionMu.Unlock()
}

// SessionStore exposes the session-token operations needed by auth handlers.
type SessionStore struct{}

func (SessionStore) SessionInfo(tok string) (SessionMetadata, bool) {
	return SessionInfo(tok)
}

func (SessionStore) DeleteSession(tok string) {
	DeleteSession(tok)
}

// SetSession stores a pre-generated session token for trusted native hosts.
func SetSession(tok string) {
	now := time.Now()
	sessionMu.Lock()
	pruneSessionsLocked(now)
	sessionKeys[tok] = sessionRecord{
		CreatedAt: now,
		ExpiresAt: now.Add(sessionMaxDuration),
	}
	enforceSessionLimitLocked()
	sessionMu.Unlock()
}

func pruneSessionsLocked(now time.Time) {
	for tok, record := range sessionKeys {
		if shouldPruneSessionRecord(record, now) {
			delete(sessionKeys, tok)
		}
	}
}

func shouldPruneSessionRecord(record sessionRecord, now time.Time) bool {
	if record.CreatedAt.IsZero() || record.ExpiresAt.IsZero() {
		return true
	}
	if record.ExpiresAt.After(now) {
		return false
	}
	return strings.TrimSpace(record.LogicalSessionID) == "" ||
		!record.ExpiresAt.Add(uiSessionRecoveryDuration).After(now)
}

func logicalSessionExpiresAt(record sessionRecord) time.Time {
	if record.LogicalExpiresAt.After(record.ExpiresAt) {
		return record.LogicalExpiresAt
	}
	return record.ExpiresAt
}

func enforceSessionLimitLocked() {
	for len(sessionKeys) > maxSessionKeys {
		var oldestToken string
		var oldestCreated time.Time
		for tok, record := range sessionKeys {
			if oldestToken == "" || record.CreatedAt.Before(oldestCreated) {
				oldestToken = tok
				oldestCreated = record.CreatedAt
			}
		}
		if oldestToken == "" {
			return
		}
		delete(sessionKeys, oldestToken)
	}
}