// Package uisession owns bounded, ephemeral state scoped to one authenticated
// Desktop UI session. It deliberately stores only local control-plane facts;
// persisted server credentials remain owned by profile.
package uisession

import (
	"errors"
	"strings"
	"sync"
	"time"
)

const (
	defaultMaxSessions         = 256
	defaultSubscriptionChannel = 8
)

var (
	ErrSessionRequired = errors.New("UI session is required")
	ErrSessionExpired  = errors.New("UI session has expired")
	ErrServerRequired  = errors.New("server is required")
)

type Snapshot struct {
	ActiveServerID string
	Generation     uint64
	ExpiresAt      time.Time
}

type Options struct {
	MaxSessions int
	Now         func() time.Time
	// SubscriptionChannel is the per-subscriber change-channel capacity
	// (default 8 when zero). Frozen as part of the rev-9 contract scaffold;
	// NewStore does NOT read it yet.
	SubscriptionChannel int
}

type entry struct {
	Snapshot
	revision  uint64
	updatedAt time.Time
}

type generationWatch struct {
	generation uint64
	changed    chan struct{}
	refs       int
}

// Store keeps active-server selection isolated by stable logical UI-session
// identity. It is safe for concurrent use.
type Store struct {
	mu          sync.RWMutex
	entries     map[string]entry
	watches     map[string]*generationWatch
	transitions map[string]int
	maxSessions int
	subCapacity int
	now         func() time.Time
	subscribers map[string]*subscriberRegistry
	tombstones  map[string]time.Time
}

// ServerTransition owns one server identity transition. Complete reopens
// gateway work after a successful mutation or successful rollback. FailClosed
// clears every active selection for the server and ends the transition without
// the normal reopen fence.
type ServerTransition struct {
	store    *Store
	serverID string
	once     sync.Once
}

func NewStore(options Options) *Store {
	maxSessions := options.MaxSessions
	if maxSessions <= 0 {
		maxSessions = defaultMaxSessions
	}
	now := options.Now
	if now == nil {
		now = time.Now
	}
	subCapacity := options.SubscriptionChannel
	if subCapacity <= 0 {
		subCapacity = defaultSubscriptionChannel
	}
	return &Store{
		entries:     make(map[string]entry),
		watches:     make(map[string]*generationWatch),
		transitions: make(map[string]int),
		maxSessions: maxSessions,
		subCapacity: subCapacity,
		now:         now,
		subscribers: make(map[string]*subscriberRegistry),
		tombstones:  make(map[string]time.Time),
	}
}

// Get returns one live UI session snapshot. Expired state is removed lazily.
func (s *Store) Get(sessionID string) (Snapshot, bool) {
	sessionID = strings.TrimSpace(sessionID)
	if sessionID == "" {
		return Snapshot{}, false
	}
	now := s.now()
	s.mu.Lock()
	defer s.mu.Unlock()
	current, ok := s.entries[sessionID]
	if !ok {
		return Snapshot{}, false
	}
	if expired(current.ExpiresAt, now) {
		s.notifyGenerationChangedLocked(sessionID)
		s.closeSubscribersLocked(sessionID)
		delete(s.entries, sessionID)
		return Snapshot{}, false
	}
	return current.Snapshot, true
}

// ExtendLease advances the authentication lease attached to one stable logical
// UI session without changing its active selection or generation. The control
// plane calls this before business/API state is touched after a renderer lease
// is rotated or reissued.
func (s *Store) ExtendLease(sessionID string, expiresAt time.Time) bool {
	sessionID = strings.TrimSpace(sessionID)
	if sessionID == "" || !expiresAt.After(s.now()) {
		return false
	}
	now := s.now()
	s.mu.Lock()
	defer s.mu.Unlock()
	current, ok := s.entries[sessionID]
	if ok {
		current.ExpiresAt = expiresAt
		current.updatedAt = now
		s.entries[sessionID] = current
	}
	if registry := s.subscribers[sessionID]; registry != nil {
		registry.expiresAt = expiresAt
	}
	return ok
}

// MatchesGeneration reports whether the session still owns the same active
// server generation captured before an in-flight gateway request began.
func (s *Store) MatchesGeneration(sessionID string, generation uint64) bool {
	current, ok := s.Get(sessionID)
	return ok && current.ActiveServerID != "" && current.Generation == generation
}

// WatchGeneration atomically verifies a live selection generation and returns
// a channel closed by the next selection mutation. Call release when the
// associated request ends so idle generations do not retain waiters.
func (s *Store) WatchGeneration(
	sessionID string,
	generation uint64,
) (<-chan struct{}, func(), bool) {
	sessionID = strings.TrimSpace(sessionID)
	if sessionID == "" {
		return nil, func() {}, false
	}
	now := s.now()
	s.mu.Lock()
	defer s.mu.Unlock()
	current, ok := s.entries[sessionID]
	if !ok || expired(current.ExpiresAt, now) {
		s.notifyGenerationChangedLocked(sessionID)
		s.closeSubscribersLocked(sessionID)
		delete(s.entries, sessionID)
		return nil, func() {}, false
	}
	if current.ActiveServerID == "" || current.Generation != generation {
		return nil, func() {}, false
	}
	if s.transitions[current.ActiveServerID] > 0 {
		return nil, func() {}, false
	}
	watch := s.watches[sessionID]
	if watch == nil || watch.generation != generation {
		s.notifyGenerationChangedLocked(sessionID)
		watch = &generationWatch{
			generation: generation,
			changed:    make(chan struct{}),
		}
		s.watches[sessionID] = watch
	}
	watch.refs++
	var once sync.Once
	release := func() {
		once.Do(func() {
			s.releaseGenerationWatch(sessionID, generation, watch.changed)
		})
	}
	return watch.changed, release, true
}

// Select atomically changes one UI session's active server and advances only
// that session's generation.
func (s *Store) Select(sessionID string, expiresAt time.Time, serverID string) (Snapshot, error) {
	sessionID = strings.TrimSpace(sessionID)
	if sessionID == "" {
		return Snapshot{}, ErrSessionRequired
	}
	serverID = strings.TrimSpace(serverID)
	if serverID == "" {
		return Snapshot{}, ErrServerRequired
	}
	now := s.now()
	if expired(expiresAt, now) {
		return Snapshot{}, ErrSessionExpired
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	s.pruneLocked(now)
	if _, closed := s.tombstones[sessionID]; closed {
		return Snapshot{}, ErrSessionClosed
	}
	current := s.entries[sessionID]
	s.notifyGenerationChangedLocked(sessionID)
	current.ActiveServerID = serverID
	current.Generation++
	current.revision++
	current.ExpiresAt = expiresAt
	current.updatedAt = now
	s.entries[sessionID] = current
	s.publishLocked(sessionID, Change{
		Kind:           ChangeKindSelected,
		Revision:       current.revision,
		State:          SubscriptionStateSelected,
		ActiveServerID: current.ActiveServerID,
		Generation:     current.Generation,
	})
	s.enforceLimitLocked()
	return current.Snapshot, nil
}

// Clear removes the selected server while preserving the UI session's
// generation history so a later selection cannot reuse an old generation.
func (s *Store) Clear(sessionID string) (Snapshot, bool) {
	sessionID = strings.TrimSpace(sessionID)
	if sessionID == "" {
		return Snapshot{}, false
	}
	now := s.now()
	s.mu.Lock()
	defer s.mu.Unlock()
	current, ok := s.entries[sessionID]
	if !ok || expired(current.ExpiresAt, now) {
		s.notifyGenerationChangedLocked(sessionID)
		s.closeSubscribersLocked(sessionID)
		delete(s.entries, sessionID)
		return Snapshot{}, false
	}
	s.notifyGenerationChangedLocked(sessionID)
	current.ActiveServerID = ""
	current.Generation++
	current.revision++
	current.updatedAt = now
	s.entries[sessionID] = current
	s.publishLocked(sessionID, Change{
		Kind:       ChangeKindCleared,
		Revision:   current.revision,
		State:      SubscriptionStateCleared,
		Generation: current.Generation,
	})
	return current.Snapshot, true
}

// ClearIfGeneration clears a stale selection only when no concurrent switch
// has advanced the session since the caller took its snapshot.
func (s *Store) ClearIfGeneration(sessionID string, generation uint64) bool {
	sessionID = strings.TrimSpace(sessionID)
	if sessionID == "" {
		return false
	}
	now := s.now()
	s.mu.Lock()
	defer s.mu.Unlock()
	current, ok := s.entries[sessionID]
	if !ok || expired(current.ExpiresAt, now) {
		s.notifyGenerationChangedLocked(sessionID)
		s.closeSubscribersLocked(sessionID)
		delete(s.entries, sessionID)
		return false
	}
	if current.Generation != generation {
		return false
	}
	s.notifyGenerationChangedLocked(sessionID)
	current.ActiveServerID = ""
	current.Generation++
	current.revision++
	current.updatedAt = now
	s.entries[sessionID] = current
	s.publishLocked(sessionID, Change{
		Kind:       ChangeKindCleared,
		Revision:   current.revision,
		State:      SubscriptionStateCleared,
		Generation: current.Generation,
	})
	return true
}

// Delete removes every local fact for a revoked UI session.
func (s *Store) Delete(sessionID string) {
	sessionID = strings.TrimSpace(sessionID)
	s.mu.Lock()
	s.notifyGenerationChangedLocked(sessionID)
	s.closeSubscribersLocked(sessionID)
	delete(s.entries, sessionID)
	s.mu.Unlock()
}

// InvalidateServer clears all live selections pointing at a deleted or
// otherwise invalid server.
func (s *Store) InvalidateServer(serverID string) int {
	serverID = strings.TrimSpace(serverID)
	if serverID == "" {
		return 0
	}
	now := s.now()
	s.mu.Lock()
	defer s.mu.Unlock()
	s.pruneLocked(now)
	return s.invalidateServerLocked(serverID, now)
}

// InvalidateServersMatching clears all live selections whose target ID matches
// predicate. The predicate runs while the store lock is held and must not call
// back into Store.
func (s *Store) InvalidateServersMatching(predicate func(string) bool) int {
	if predicate == nil {
		return 0
	}
	now := s.now()
	s.mu.Lock()
	defer s.mu.Unlock()
	s.pruneLocked(now)
	serverIDs := make(map[string]struct{})
	for _, current := range s.entries {
		if current.ActiveServerID != "" && predicate(current.ActiveServerID) {
			serverIDs[current.ActiveServerID] = struct{}{}
		}
	}
	cleared := 0
	for serverID := range serverIDs {
		cleared += s.invalidateServerLocked(serverID, now)
	}
	return cleared
}

// AdvanceServerGeneration fences in-flight work for every UI session using a
// server whose upstream identity or binding changed, without forcing users to
// reselect the same server.
func (s *Store) AdvanceServerGeneration(serverID string) int {
	serverID = strings.TrimSpace(serverID)
	if serverID == "" {
		return 0
	}
	now := s.now()
	s.mu.Lock()
	defer s.mu.Unlock()
	s.pruneLocked(now)
	return s.advanceServerGenerationLocked(serverID, now)
}

// BeginServerTransition fences both sides of one registered server identity
// update. Existing work is canceled before the registry or running connection
// can change, and work cannot begin until the returned completion callback
// advances generation again after the transition finishes.
func (s *Store) BeginServerTransition(serverID string) func() {
	transition := s.BeginServerTransitionHandle(serverID)
	return transition.Complete
}

// BeginServerTransitionHandle starts a transition with explicit success and
// fail-closed completion paths.
func (s *Store) BeginServerTransitionHandle(serverID string) *ServerTransition {
	serverID = strings.TrimSpace(serverID)
	if serverID == "" {
		return &ServerTransition{}
	}
	now := s.now()
	s.mu.Lock()
	s.pruneLocked(now)
	s.transitions[serverID]++
	if s.transitions[serverID] == 1 {
		s.advanceServerGenerationLocked(serverID, now)
	}
	s.mu.Unlock()

	return &ServerTransition{store: s, serverID: serverID}
}

// Complete ends a successful transition and advances generation once more
// before new work can begin.
func (t *ServerTransition) Complete() {
	if t == nil || t.store == nil || t.serverID == "" {
		return
	}
	t.once.Do(func() {
		now := t.store.now()
		t.store.mu.Lock()
		defer t.store.mu.Unlock()
		t.store.pruneLocked(now)
		if t.store.transitions[t.serverID] <= 1 {
			delete(t.store.transitions, t.serverID)
			t.store.advanceServerGenerationLocked(t.serverID, now)
		} else {
			t.store.transitions[t.serverID]--
		}
	})
}

// FailClosed ends a failed transition by clearing every selection for the
// server. It deliberately does not execute the normal reopen generation fence.
func (t *ServerTransition) FailClosed() {
	if t == nil || t.store == nil || t.serverID == "" {
		return
	}
	t.once.Do(func() {
		now := t.store.now()
		t.store.mu.Lock()
		defer t.store.mu.Unlock()
		t.store.pruneLocked(now)
		if t.store.transitions[t.serverID] <= 1 {
			delete(t.store.transitions, t.serverID)
		} else {
			t.store.transitions[t.serverID]--
		}
		t.store.invalidateServerLocked(t.serverID, now)
	})
}

// Prune drops expired UI sessions and returns the number of live entries.
func (s *Store) Prune() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.pruneLocked(s.now())
	return len(s.entries)
}

// Len returns the number of retained UI sessions after pruning.
func (s *Store) Len() int {
	return s.Prune()
}

func (s *Store) pruneLocked(now time.Time) {
	for sessionID, current := range s.entries {
		if expired(current.ExpiresAt, now) {
			s.notifyGenerationChangedLocked(sessionID)
			s.closeSubscribersLocked(sessionID)
			delete(s.entries, sessionID)
		}
	}
	for sessionID, registry := range s.subscribers {
		if expired(registry.expiresAt, now) {
			s.closeSubscribersLocked(sessionID)
		}
	}
	for sessionID, expiresAt := range s.tombstones {
		if expired(expiresAt, now) {
			delete(s.tombstones, sessionID)
		}
	}
}

func (s *Store) enforceLimitLocked() {
	for len(s.entries) > s.maxSessions {
		var oldestID string
		var oldest time.Time
		for sessionID, current := range s.entries {
			if oldestID == "" || current.updatedAt.Before(oldest) {
				oldestID = sessionID
				oldest = current.updatedAt
			}
		}
		if oldestID == "" {
			return
		}
		s.notifyGenerationChangedLocked(oldestID)
		s.closeSubscribersLocked(oldestID)
		delete(s.entries, oldestID)
	}
}

func (s *Store) notifyGenerationChangedLocked(sessionID string) {
	watch := s.watches[sessionID]
	if watch == nil {
		return
	}
	delete(s.watches, sessionID)
	close(watch.changed)
}

func (s *Store) advanceServerGenerationLocked(serverID string, now time.Time) int {
	advanced := 0
	for sessionID, current := range s.entries {
		if current.ActiveServerID != serverID {
			continue
		}
		s.notifyGenerationChangedLocked(sessionID)
		current.Generation++
		current.revision++
		current.updatedAt = now
		s.entries[sessionID] = current
		s.publishLocked(sessionID, Change{
			Kind:           ChangeKindAdvanced,
			Revision:       current.revision,
			State:          SubscriptionStateSelected,
			ActiveServerID: current.ActiveServerID,
			Generation:     current.Generation,
		})
		advanced++
	}
	return advanced
}

func (s *Store) invalidateServerLocked(serverID string, now time.Time) int {
	cleared := 0
	for sessionID, current := range s.entries {
		if current.ActiveServerID != serverID {
			continue
		}
		s.notifyGenerationChangedLocked(sessionID)
		current.ActiveServerID = ""
		current.Generation++
		current.revision++
		current.updatedAt = now
		s.entries[sessionID] = current
		s.publishLocked(sessionID, Change{
			Kind:       ChangeKindCleared,
			Revision:   current.revision,
			State:      SubscriptionStateCleared,
			Generation: current.Generation,
		})
		cleared++
	}
	return cleared
}

func (s *Store) releaseGenerationWatch(
	sessionID string,
	generation uint64,
	changed chan struct{},
) {
	s.mu.Lock()
	defer s.mu.Unlock()
	watch := s.watches[sessionID]
	if watch == nil || watch.generation != generation || watch.changed != changed {
		return
	}
	watch.refs--
	if watch.refs <= 0 {
		delete(s.watches, sessionID)
	}
}

func expired(expiresAt, now time.Time) bool {
	return expiresAt.IsZero() || !expiresAt.After(now)
}