package dispatcher

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

const sessionKeyVersion = "sk1"

var errSessionManagerClosed = errors.New("dispatcher session manager is closed")

// SessionKey is the dispatcher-owned identity for special runtime sessions.
// Normal interactive sessions are intentionally scoped only by agent ID so a
// local agent keeps one long-lived process and workspace as it moves across
// AgentTeam channels, threads, DMs, and task discussion scopes. Extra fields remain
// available for explicitly isolated internal workflows.
type SessionKey struct {
	AgentID         string
	ChannelID       string
	ConversationKey string
	Kind            string
}

func NewSessionKey(agentID, channelID, conversationKey string) SessionKey {
	return SessionKey{
		AgentID:         strings.TrimSpace(agentID),
		ChannelID:       strings.TrimSpace(channelID),
		ConversationKey: strings.TrimSpace(conversationKey),
	}
}

func (k SessionKey) WithKind(kind string) SessionKey {
	k.Kind = strings.TrimSpace(kind)
	return k
}

func (k SessionKey) String() string {
	if strings.TrimSpace(k.AgentID) != "" &&
		strings.TrimSpace(k.ChannelID) == "" &&
		strings.TrimSpace(k.ConversationKey) == "" &&
		strings.TrimSpace(k.Kind) == "" {
		return strings.TrimSpace(k.AgentID)
	}
	parts := []string{
		sessionKeyVersion,
		encodeSessionKeyPart(k.AgentID),
		encodeSessionKeyPart(k.ChannelID),
		encodeSessionKeyPart(k.ConversationKey),
		encodeSessionKeyPart(k.Kind),
	}
	return strings.Join(parts, ":")
}

func (k SessionKey) IsZero() bool {
	return strings.TrimSpace(k.AgentID) == ""
}

func encodeSessionKeyPart(value string) string {
	return base64.RawURLEncoding.EncodeToString([]byte(strings.TrimSpace(value)))
}

func parseSessionKey(value string) (SessionKey, bool) {
	parts := strings.Split(strings.TrimSpace(value), ":")
	if len(parts) != 5 || parts[0] != sessionKeyVersion {
		return SessionKey{}, false
	}
	decoded := make([]string, 4)
	for i := 1; i < len(parts); i++ {
		bytes, err := base64.RawURLEncoding.DecodeString(parts[i])
		if err != nil {
			return SessionKey{}, false
		}
		decoded[i-1] = string(bytes)
	}
	return SessionKey{
		AgentID:         decoded[0],
		ChannelID:       decoded[1],
		ConversationKey: decoded[2],
		Kind:            decoded[3],
	}, true
}

func conversationKeyForMetadata(meta types.Metadata) string {
	if meta == nil {
		return ""
	}
	if agentteamMeta, ok := meta.(*types.AgentTeamMeta); ok {
		if target := strings.TrimSpace(agentteamMeta.Target); target != "" {
			return target
		}
	}
	return types.NormalizeMetadata(meta).ConversationKey
}

func sessionKeyForBinding(agentID string, _ conversationBinding) SessionKey {
	return NewSessionKey(agentID, "", "")
}

func sessionWorkDir(agentID string) string {
	return filepath.Join(profile.ConfigDir(), "agents", strings.TrimSpace(agentID))
}

func agentRuntimeStateDir(agentID string) string {
	return filepath.Join(sessionWorkDir(agentID), ".runtime")
}

func sessionWorkDirForKey(agentID, sessionKey string, _ bool) string {
	agentDir := sessionWorkDir(agentID)
	trimmedKey := strings.TrimSpace(sessionKey)
	if trimmedKey == "" || trimmedKey == strings.TrimSpace(agentID) {
		return agentDir
	}
	return filepath.Join(agentDir, "sessions", sessionKeyHash(trimmedKey))
}

func sessionKeyHash(sessionKey string) string {
	sum := sha256.Sum256([]byte(strings.TrimSpace(sessionKey)))
	return hex.EncodeToString(sum[:16])
}

// SessionManager tracks active agent sessions by key, providing creation, lookup, and lifecycle management including graceful shutdown.
type SessionManager struct {
	mu              sync.RWMutex
	sessions        map[string]*interactiveSession
	starting        map[string]*sessionStartReservation
	startOwners     map[*sessionStartReservation]struct{}
	startGeneration map[string]uint64
	closed          bool
}

// sessionStartReservation is the per-session-key generation fence shared by
// every Process start path. expected is nil for a new key and points at the
// currently published session for a restart/continuation. A Process may only
// become visible through PublishStart while this exact generation is current.
//
// All fields except done are read or written while SessionManager.mu is held.
type sessionStartReservation struct {
	key        string
	generation uint64
	expected   *interactiveSession
	done       chan struct{}
	ownerDone  chan struct{}
	finishOnce sync.Once
	published  *interactiveSession
}

func NewSessionManager() *SessionManager {
	return &SessionManager{
		sessions:        make(map[string]*interactiveSession),
		starting:        make(map[string]*sessionStartReservation),
		startOwners:     make(map[*sessionStartReservation]struct{}),
		startGeneration: make(map[string]uint64),
	}
}

func (sm *SessionManager) Get(key string) (*interactiveSession, bool) {
	if sm == nil {
		return nil, false
	}
	sm.mu.RLock()
	sess, ok := sm.sessions[strings.TrimSpace(key)]
	sm.mu.RUnlock()
	return sess, ok
}

func (sm *SessionManager) Set(key string, sess *interactiveSession) {
	if sm == nil || sess == nil {
		return
	}
	key = strings.TrimSpace(key)
	var previous *interactiveSession
	sm.mu.Lock()
	if sm.closed {
		sess.retired.Store(true)
		sm.mu.Unlock()
		revokeSessionMCPToken(sess)
		return
	}
	if current := sm.sessions[key]; current != nil && current != sess {
		previous = current
	}
	if previous != nil {
		previous.retired.Store(true)
	}
	sm.sessions[key] = sess
	sm.mu.Unlock()
	revokeSessionMCPToken(previous)
}

// startReservationCurrent reports whether the start reservation is still the
// live generation for its key. An invalidated (deleted or superseded)
// reservation makes every launch-window scope bound to it fail closed.
func (sm *SessionManager) startReservationCurrent(res *sessionStartReservation) bool {
	if sm == nil || res == nil {
		return false
	}
	sm.mu.RLock()
	defer sm.mu.RUnlock()
	// Currency requires BOTH the generation to be live AND this exact
	// reservation to still be the active start owner. Invalidation (Delete,
	// Close, supersede) removes the entry from starting without necessarily
	// advancing the generation, so the generation check alone would keep
	// reporting an invalidated reservation as current.
	return sm.startGeneration[res.key] == res.generation && sm.starting[res.key] == res
}

func (sm *SessionManager) Delete(key string) {
	if sm == nil {
		return
	}
	key = strings.TrimSpace(key)
	var removed *interactiveSession
	sm.mu.Lock()
	sm.invalidateStartLocked(key)
	removed = sm.sessions[key]
	if removed != nil {
		removed.retired.Store(true)
	}
	delete(sm.sessions, key)
	sm.mu.Unlock()
	revokeSessionMCPToken(removed)
}

// Close atomically fences every in-flight start generation and prevents new
// generations from being reserved. A generation that already spawned a
// Process will fail PublishStart and tear that Process down instead of becoming
// visible after Dispatcher shutdown. Existing sessions remain snapshot-able so
// the Dispatcher can perform their ordinary teardown after closing the fence.
//
// The returned channels are a snapshot of reservation owners that have not yet
// called FinishStart. They are intentionally distinct from reservation.done:
// done is closed immediately to wake delivery waiters when shutdown invalidates
// a generation, while ownerDone closes only after the launch/recovery owner has
// unwound. SessionManager itself never waits for these channels, so an owner may
// safely initiate Close without synchronously waiting on itself. Top-level
// Dispatcher shutdown paths join the snapshot with a bounded context.
func (sm *SessionManager) Close() []<-chan struct{} {
	owners, _ := sm.CloseAndSnapshot(true)
	return owners
}

// CloseAndSnapshot closes start/publication admission and captures the sessions
// visible at that exact manager boundary under one lock. Shutdown uses this
// instead of a separate Close then Snapshot pair: a concurrent remover may
// retire/delete a session immediately after admission closes, but its already-
// published memory watcher must remain in the shutdown drain set.
func (sm *SessionManager) CloseAndSnapshot(includeEphemeral bool) ([]<-chan struct{}, []*interactiveSession) {
	if sm == nil {
		return nil, nil
	}
	sm.mu.Lock()
	if !sm.closed {
		sm.closed = true
		for key := range sm.starting {
			sm.invalidateStartLocked(key)
		}
		for _, sess := range sm.sessions {
			if sess != nil {
				sess.retired.Store(true)
			}
		}
	}
	owners := make([]<-chan struct{}, 0, len(sm.startOwners))
	for reservation := range sm.startOwners {
		if reservation != nil && reservation.ownerDone != nil {
			owners = append(owners, reservation.ownerDone)
		}
	}
	sessions := make([]*interactiveSession, 0, len(sm.sessions))
	for _, sess := range sm.sessions {
		if sess == nil || (!includeEphemeral && sess.ephemeral) {
			continue
		}
		sessions = append(sessions, sess)
	}
	sm.mu.Unlock()
	return owners, sessions
}

func (sm *SessionManager) Closed() bool {
	if sm == nil {
		return true
	}
	sm.mu.RLock()
	closed := sm.closed
	sm.mu.RUnlock()
	return closed
}

// invalidateStartLocked invalidates exactly the currently registered
// generation for key. FinishStart checks identity before closing, so an owner
// finishing an invalidated generation cannot double-close this channel or a
// later generation's channel.
func (sm *SessionManager) invalidateStartLocked(key string) {
	reservation, ok := sm.starting[key]
	if !ok || reservation == nil {
		return
	}
	delete(sm.starting, key)
	close(reservation.done)
}

func (sm *SessionManager) Snapshot(includeEphemeral bool) []*interactiveSession {
	if sm == nil {
		return nil
	}
	sm.mu.RLock()
	sessions := make([]*interactiveSession, 0, len(sm.sessions))
	for _, sess := range sm.sessions {
		if sess == nil {
			continue
		}
		if !includeEphemeral && sess.ephemeral {
			continue
		}
		sessions = append(sessions, sess)
	}
	sm.mu.RUnlock()
	return sessions
}

func (sm *SessionManager) FindByLookupID(id string) (*interactiveSession, bool) {
	id = strings.TrimSpace(id)
	if id == "" || sm == nil {
		return nil, false
	}
	if sess, ok := sm.Get(id); ok {
		return sess, true
	}

	sm.mu.RLock()
	defer sm.mu.RUnlock()
	var newest *interactiveSession
	var newestStartedAt time.Time
	for _, sess := range sm.sessions {
		if sess == nil || sess.ephemeral {
			continue
		}
		if !sessionMatchesLookupID(sess, id) {
			continue
		}
		startedAt := sess.startedAt
		if newest == nil || startedAt.After(newestStartedAt) {
			newest = sess
			newestStartedAt = startedAt
		}
	}
	return newest, newest != nil
}

func sessionMatchesLookupID(sess *interactiveSession, id string) bool {
	return strings.TrimSpace(sess.sessionKey) == id ||
		strings.TrimSpace(sess.agentID) == id ||
		strings.TrimSpace(sess.serverAgentID) == id ||
		strings.TrimSpace(sess.agent.ID) == id
}

// StartInFlight reports the current start generation for key. Callers that can
// observe a Process published before its post-launch durable handoff completes
// must wait on this channel before delivering another message.
func (sm *SessionManager) StartInFlight(key string) (<-chan struct{}, bool) {
	if sm == nil {
		return nil, false
	}
	key = strings.TrimSpace(key)
	sm.mu.RLock()
	reservation, ok := sm.starting[key]
	sm.mu.RUnlock()
	if !ok || reservation == nil {
		return nil, false
	}
	return reservation.done, true
}

// CommitDeliveryIfCurrent atomically linearizes a completed Process handoff
// against session start/replacement publication. The callback may mutate sess
// fields that are protected by sess.mu; it must not call back into
// SessionManager or block on Process/runtime work.
func (sm *SessionManager) CommitDeliveryIfCurrent(
	key string,
	sess *interactiveSession,
	process *agent.Process,
	commit func(),
) bool {
	if sm == nil || sess == nil || process == nil {
		return false
	}
	key = strings.TrimSpace(key)
	sm.mu.Lock()
	defer sm.mu.Unlock()
	if sm.closed || sm.starting[key] != nil || sm.sessions[key] != sess {
		return false
	}
	sess.mu.Lock()
	defer sess.mu.Unlock()
	if sess.retired.Load() || sess.process != process || !process.AcceptingDeliveries() {
		return false
	}
	if commit != nil {
		commit()
	}
	return true
}

// BeginStart reserves the next start generation for key. expected must be nil
// when the key is currently unpublished, or the exact current session for a
// restart/continuation. This check and reservation creation are atomic with
// PublishStart and RemoveIfCurrent.
func (sm *SessionManager) BeginStart(key string, expected *interactiveSession) (*sessionStartReservation, <-chan struct{}, bool) {
	return sm.beginStart(key, expected, nil, false)
}

// BeginStartForProcess reserves a restart only while both the session and its
// exact Process generation remain current.
func (sm *SessionManager) BeginStartForProcess(key string, expected *interactiveSession, expectedProcess *agent.Process) (*sessionStartReservation, <-chan struct{}, bool) {
	return sm.beginStart(key, expected, expectedProcess, true)
}

func (sm *SessionManager) beginStart(key string, expected *interactiveSession, expectedProcess *agent.Process, requireProcessMatch bool) (*sessionStartReservation, <-chan struct{}, bool) {
	if sm == nil {
		closed := make(chan struct{})
		close(closed)
		return nil, closed, false
	}
	key = strings.TrimSpace(key)
	sm.mu.Lock()
	defer sm.mu.Unlock()
	if sm.closed {
		closed := make(chan struct{})
		close(closed)
		return nil, closed, false
	}
	if current, ok := sm.starting[key]; ok && current != nil {
		return nil, current.done, false
	}
	current, exists := sm.sessions[key]
	if expected == nil {
		if exists {
			closed := make(chan struct{})
			close(closed)
			return nil, closed, false
		}
	} else if !exists || current != expected || expected.retired.Load() {
		closed := make(chan struct{})
		close(closed)
		return nil, closed, false
	}
	if requireProcessMatch {
		expected.mu.Lock()
		matches := expected.process == expectedProcess
		expected.mu.Unlock()
		if !matches {
			closed := make(chan struct{})
			close(closed)
			return nil, closed, false
		}
	}
	sm.startGeneration[key]++
	reservation := &sessionStartReservation{
		key:        key,
		generation: sm.startGeneration[key],
		expected:   expected,
		done:       make(chan struct{}),
		ownerDone:  make(chan struct{}),
	}
	sm.starting[key] = reservation
	sm.startOwners[reservation] = struct{}{}
	return reservation, reservation.done, true
}

// BeginRecovery selects the current recovery target and reserves the default
// production session key in one SessionManager critical section. The returned
// reservation deliberately uses the same starting map/generation as
// BeginStart/PublishStart: normal delivery, configuration replacement, and
// agent:start must therefore settle wholly before or wholly after recovery.
//
// preferredKey is the best production key known before target selection
// (normally the local agent id). match is both a server-to-local resolution
// fallback and compatibility support for an in-memory legacy session published
// under an older conversation-scoped key. When a matching target identifies a
// different local agent id, recovery reserves that local agent's default
// production key so its normal delivery/start generation cannot cross restore.
func (sm *SessionManager) BeginRecovery(
	preferredKey string,
	match func(*interactiveSession) bool,
) (*interactiveSession, *sessionStartReservation, <-chan struct{}, bool) {
	if sm == nil {
		closed := make(chan struct{})
		close(closed)
		return nil, nil, closed, false
	}
	preferredKey = strings.TrimSpace(preferredKey)
	sm.mu.Lock()
	defer sm.mu.Unlock()
	if sm.closed {
		closed := make(chan struct{})
		close(closed)
		return nil, nil, closed, false
	}
	if current, ok := sm.starting[preferredKey]; ok && current != nil {
		return nil, nil, current.done, false
	}

	target := sm.sessions[preferredKey]
	if target == nil && match != nil {
		var newestStartedAt time.Time
		for _, candidate := range sm.sessions {
			if candidate == nil || candidate.ephemeral || !match(candidate) {
				continue
			}
			if target == nil || candidate.startedAt.After(newestStartedAt) {
				target = candidate
				newestStartedAt = candidate.startedAt
			}
		}
	}

	reservationKey := preferredKey
	if target != nil {
		target.mu.Lock()
		localAgentID := strings.TrimSpace(target.agentID)
		target.mu.Unlock()
		if localAgentID != "" {
			reservationKey = sessionKeyForBinding(localAgentID, conversationBinding{}).String()
		}
	}
	if current, ok := sm.starting[reservationKey]; ok && current != nil {
		return nil, nil, current.done, false
	}

	sm.startGeneration[reservationKey]++
	reservation := &sessionStartReservation{
		key:        reservationKey,
		generation: sm.startGeneration[reservationKey],
		expected:   target,
		done:       make(chan struct{}),
		ownerDone:  make(chan struct{}),
	}
	sm.starting[reservationKey] = reservation
	sm.startOwners[reservation] = struct{}{}
	return target, reservation, reservation.done, true
}

// PublishStart atomically validates the reservation generation and expected
// current session, installs process on sess, and publishes sess as the current
// session. A replacement retires the previous session under the same manager
// lock, so a late launch can never overwrite or outlive a concurrent removal.
func (sm *SessionManager) PublishStart(
	reservation *sessionStartReservation,
	sess *interactiveSession,
	process *agent.Process,
	runtimeStartFingerprint string,
	usesWorkspaceRootMemory bool,
) (*interactiveSession, bool) {
	if sm == nil || reservation == nil || sess == nil || process == nil {
		return nil, false
	}
	sm.mu.Lock()
	defer sm.mu.Unlock()
	if sm.closed {
		return nil, false
	}
	currentReservation, ok := sm.starting[reservation.key]
	if !ok || currentReservation != reservation || currentReservation.generation != reservation.generation {
		return nil, false
	}
	current, exists := sm.sessions[reservation.key]
	if reservation.expected == nil {
		if exists {
			return nil, false
		}
	} else if !exists || current != reservation.expected || reservation.expected.retired.Load() {
		return nil, false
	}
	if sess.retired.Load() {
		return nil, false
	}

	sess.mu.Lock()
	if sess.process != nil && sess.process != process {
		sess.mu.Unlock()
		return nil, false
	}
	sess.terminalReportClaimed = false
	sess.terminalReportDone = nil
	sess.terminalReportWillStop = false
	sess.restartRetirementInFlight = false
	sess.runtimeStartReplacement = false
	sess.process = process
	sess.runtimeStartFingerprint = runtimeStartFingerprint
	sess.processUsesWorkspaceRootMemory = usesWorkspaceRootMemory
	sess.mu.Unlock()

	var previous *interactiveSession
	if current != nil && current != sess {
		previous = current
		previous.retired.Store(true)
	}
	sm.sessions[reservation.key] = sess
	reservation.published = sess
	return previous, true
}

// FinishStart releases exactly one reservation generation. Session
// publication is performed only by PublishStart; an error before publication
// therefore leaves the expected session/current map untouched.
func (sm *SessionManager) FinishStart(reservation *sessionStartReservation) {
	if sm == nil || reservation == nil {
		return
	}
	sm.mu.Lock()
	if current, ok := sm.starting[reservation.key]; ok && current == reservation && current.generation == reservation.generation {
		delete(sm.starting, reservation.key)
		close(reservation.done)
	}
	reservation.finishOnce.Do(func() {
		close(reservation.ownerDone)
	})
	delete(sm.startOwners, reservation)
	sm.mu.Unlock()
}

func (sm *SessionManager) WithSessionByKey(key string, fn func(*interactiveSession)) bool {
	sess, ok := sm.Get(key)
	if !ok || sess == nil {
		return false
	}
	fn(sess)
	return true
}

func (sm *SessionManager) RemoveIfCurrent(key string, sess *interactiveSession) bool {
	if sm == nil || sess == nil {
		return false
	}
	key = strings.TrimSpace(key)
	removed := false
	var process *agent.Process
	sm.mu.Lock()
	if current, ok := sm.sessions[key]; ok && current == sess {
		sess.retired.Store(true)
		sess.mu.Lock()
		process = sess.process
		sess.process = nil
		sess.mu.Unlock()
		delete(sm.sessions, key)
		removed = true
	}
	sm.mu.Unlock()
	if removed {
		revokeSessionMCPToken(sess)
		// If publication won the manager fence immediately before removal,
		// terminate that Process rather than leaving an untracked orphan.
		if process != nil {
			process.Stop()
		}
	}
	return removed
}

func (sm *SessionManager) FirstByPredicate(pred func(*interactiveSession) bool) (*interactiveSession, bool) {
	if sm == nil || pred == nil {
		return nil, false
	}
	sm.mu.RLock()
	defer sm.mu.RUnlock()
	for _, sess := range sm.sessions {
		if sess != nil && pred(sess) {
			return sess, true
		}
	}
	return nil, false
}

func (sm *SessionManager) Each(pred func(*interactiveSession) bool) []*interactiveSession {
	if sm == nil || pred == nil {
		return nil
	}
	sm.mu.RLock()
	defer sm.mu.RUnlock()
	var result []*interactiveSession
	for _, sess := range sm.sessions {
		if sess != nil && pred(sess) {
			result = append(result, sess)
		}
	}
	return result
}

func (sm *SessionManager) FirstByServerAgentID(serverAgentID string) (*interactiveSession, bool) {
	serverAgentID = strings.TrimSpace(serverAgentID)
	if serverAgentID == "" {
		return nil, false
	}
	return sm.FirstByPredicate(func(sess *interactiveSession) bool {
		sess.mu.Lock()
		defer sess.mu.Unlock()
		return strings.TrimSpace(sess.serverAgentID) == serverAgentID
	})
}

func requireSessionKey(key string) (string, error) {
	key = strings.TrimSpace(key)
	if key == "" {
		return "", fmt.Errorf("session key is required")
	}
	return key, nil
}