package dispatcher

import (
	"bufio"
	"encoding/json"
	"log/slog"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

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

// -----------------------------------------------------------------------
// Tuning constants
// -----------------------------------------------------------------------

const (
	// maxHotBytes caps the in-memory cache per agent (4 MB).
	maxHotBytes = 4 << 20

	// maxTotalHotBytes caps the combined in-memory cache across all agents.
	maxTotalHotBytes = 32 << 20

	// trajectoryReplayLimit caps the initial cold+hot replay buffered for a
	// single SSE subscriber.
	trajectoryReplayLimit = 512

	trajectoryLiveBuffer = 256

	// trajectoryTTL controls how long cold (disk) entries are kept.
	trajectoryTTL = 24 * time.Hour

	// trajectoryGCInterval is the background sweep frequency for both
	// memory and disk.
	trajectoryGCInterval = 5 * time.Minute

	// coldSubDir is the directory under the log root where trajectory
	// NDJSON files live.
	coldSubDir = "trajectory"
)

// -----------------------------------------------------------------------
// Per-agent in-memory state
// -----------------------------------------------------------------------

type agentCache struct {
	items     []types.TrajectoryLogItem
	sizeBytes int // approximate serialised size of items
}

// -----------------------------------------------------------------------
// trajectoryHub
// -----------------------------------------------------------------------

// trajectoryHub maintains per-agent recent trajectory history (hot in
// memory, cold on disk) and fans out live events to SSE subscribers.
type trajectoryHub struct {
	mu          sync.Mutex
	hot         map[string]*agentCache                               // agentID -> hot ring
	hotBytes    int                                                  // approximate combined size of hot entries
	subscribers map[string]map[chan types.TrajectoryLogItem]struct{} // agentID -> subscriber set
	allSubs     map[chan types.TrajectoryLogItem]struct{}            // all-agent live subscriber set
	stopGC      chan struct{}
	gcDone      chan struct{}
	stopOnce    sync.Once
}

func newTrajectoryHub() *trajectoryHub {
	h := &trajectoryHub{
		hot:         make(map[string]*agentCache),
		subscribers: make(map[string]map[chan types.TrajectoryLogItem]struct{}),
		allSubs:     make(map[chan types.TrajectoryLogItem]struct{}),
		stopGC:      make(chan struct{}),
		gcDone:      make(chan struct{}),
	}
	go h.gcLoop()
	return h
}

// BeginStop concurrently and idempotently signals the background GC goroutine
// and returns its exact completion owner. Bounded dispatcher shutdown can wait
// on this channel as one task under its shared deadline.
func (h *trajectoryHub) BeginStop() <-chan struct{} {
	if h == nil {
		done := make(chan struct{})
		close(done)
		return done
	}
	h.stopOnce.Do(func() {
		close(h.stopGC)
	})
	return h.gcDone
}

// Stop is the strong local boundary: after it returns the GC goroutine can no
// longer mutate hot state or trajectory cold files.
func (h *trajectoryHub) Stop() {
	<-h.BeginStop()
}

// -----------------------------------------------------------------------
// Cold storage helpers
// -----------------------------------------------------------------------

// coldDir returns the on-disk directory for trajectory logs.
func coldDir() string {
	cfg := profile.Default().Config()
	logDir := cfg.Log.Dir
	if logDir == "" {
		logDir = filepath.Dir(cfg.LogFile())
	}
	return filepath.Join(logDir, coldSubDir)
}

// coldPath returns the NDJSON file path for one agent.
func coldPath(agentID string) string {
	return filepath.Join(coldDir(), agentID+".ndjson")
}

// appendCold appends items to the agent's NDJSON file on disk.
func appendCold(agentID string, items []types.TrajectoryLogItem) {
	dir := coldDir()
	if err := os.MkdirAll(dir, 0755); err != nil {
		slog.Warn("trajectory: create cold dir", "error", err)
		return
	}
	f, err := os.OpenFile(coldPath(agentID), os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
	if err != nil {
		slog.Warn("trajectory: open cold file", "agent", agentID, "error", err)
		return
	}
	defer f.Close()

	enc := json.NewEncoder(f)
	for i := range items {
		if err := enc.Encode(&items[i]); err != nil {
			slog.Warn("trajectory: encode cold item", "agent", agentID, "error", err)
		}
	}
}

// readCold returns items from the agent's cold file whose timestamp is at or
// after cutoff (unix ms). When limit is positive, only the latest limit items
// are retained in memory; when limit is zero, no cold items are returned; when
// limit is negative, all matching items are returned. Returns nil on any error.
func readCold(agentID string, cutoff int64, limit int) []types.TrajectoryLogItem {
	if limit == 0 {
		return nil
	}
	f, err := os.Open(coldPath(agentID))
	if err != nil {
		return nil
	}
	defer f.Close()

	var out []types.TrajectoryLogItem
	sc := bufio.NewScanner(f)
	sc.Buffer(make([]byte, 0, 256*1024), 1<<20)
	for sc.Scan() {
		var item types.TrajectoryLogItem
		if err := json.Unmarshal(sc.Bytes(), &item); err != nil {
			continue
		}
		if item.Timestamp >= cutoff {
			out = append(out, item)
			if limit > 0 && len(out) > limit {
				copy(out, out[1:])
				out = out[:limit]
			}
		}
	}
	return out
}

// -----------------------------------------------------------------------
// Broadcast
// -----------------------------------------------------------------------

// Broadcast appends entries to the hot cache (spilling old entries to disk
// if the 4 MB cap is exceeded) and fans out to live subscribers.
func (h *trajectoryHub) Broadcast(agentID string, entries []types.TrajectoryEntry) {
	h.BroadcastAt(agentID, entries, time.Now().UnixMilli())
}

func (h *trajectoryHub) BroadcastAt(agentID string, entries []types.TrajectoryEntry, timestamp int64) {
	if len(entries) == 0 {
		return
	}

	if timestamp <= 0 {
		timestamp = time.Now().UnixMilli()
	}

	h.mu.Lock()
	defer h.mu.Unlock()

	items := make([]types.TrajectoryLogItem, 0, len(entries))
	addedBytes := 0
	for _, e := range entries {
		item := types.TrajectoryLogItem{
			AgentID:   agentID,
			Entry:     e,
			Timestamp: timestamp,
		}
		items = append(items, item)
		addedBytes += estimateItemBytes(item)
	}

	ac := h.hot[agentID]
	if ac == nil {
		ac = &agentCache{}
		h.hot[agentID] = ac
	}

	ac.items = append(ac.items, items...)
	ac.sizeBytes += addedBytes
	h.hotBytes += addedBytes

	// If over 4 MB, spill the oldest entries to cold storage.
	if ac.sizeBytes > maxHotBytes {
		h.spillToCold(agentID, ac)
	}
	h.enforceTotalHotBytes()

	// Fan-out to per-agent subscribers.
	subs := h.subscribers[agentID]
	for ch := range subs {
		for _, item := range items {
			select {
			case ch <- item:
			default:
			}
		}
	}

	// Fan-out to all-agent subscribers.
	for ch := range h.allSubs {
		for _, item := range items {
			select {
			case ch <- item:
			default:
			}
		}
	}
}

// spillToCold moves the oldest hot entries to disk until sizeBytes is
// within the per-agent limit. The caller must hold h.mu.
func (h *trajectoryHub) spillToCold(agentID string, ac *agentCache) {
	var spilled []types.TrajectoryLogItem
	for ac.sizeBytes > maxHotBytes && len(ac.items) > 1 {
		item := ac.items[0]
		ac.items = ac.items[1:]
		size := estimateItemBytes(item)
		ac.sizeBytes -= size
		h.hotBytes -= size
		spilled = append(spilled, item)
	}
	if ac.sizeBytes < 0 {
		ac.sizeBytes = 0
	}
	if h.hotBytes < 0 {
		h.hotBytes = 0
	}
	if len(spilled) > 0 {
		appendCold(agentID, spilled)
	}
}

// enforceTotalHotBytes spills the oldest hot entries across all agents
// to enforce the global memory budget. The caller must hold h.mu.
func (h *trajectoryHub) enforceTotalHotBytes() {
	spilledByAgent := make(map[string][]types.TrajectoryLogItem)
	for h.hotBytes > maxTotalHotBytes {
		agentID, ac := h.oldestHotAgent()
		if ac == nil {
			return
		}
		item := ac.items[0]
		ac.items = ac.items[1:]
		size := estimateItemBytes(item)
		ac.sizeBytes -= size
		h.hotBytes -= size
		spilledByAgent[agentID] = append(spilledByAgent[agentID], item)
		if ac.sizeBytes < 0 {
			ac.sizeBytes = 0
		}
		if len(ac.items) == 0 && len(h.subscribers[agentID]) == 0 {
			delete(h.hot, agentID)
		}
	}
	if h.hotBytes < 0 {
		h.hotBytes = 0
	}
	for agentID, items := range spilledByAgent {
		appendCold(agentID, items)
	}
}

func (h *trajectoryHub) oldestHotAgent() (string, *agentCache) {
	oldestAgentID := ""
	var oldestCache *agentCache
	var oldestTimestamp int64
	for agentID, ac := range h.hot {
		if ac == nil || len(ac.items) == 0 {
			if len(h.subscribers[agentID]) == 0 {
				delete(h.hot, agentID)
			}
			continue
		}
		ts := ac.items[0].Timestamp
		if oldestCache == nil || ts < oldestTimestamp {
			oldestAgentID = agentID
			oldestCache = ac
			oldestTimestamp = ts
		}
	}
	return oldestAgentID, oldestCache
}

// -----------------------------------------------------------------------
// Subscribe / Unsubscribe
// -----------------------------------------------------------------------

// Subscribe returns a channel pre-filled with recent history (cold + hot,
// within TTL) for the given agent, then receives live events.
func (h *trajectoryHub) Subscribe(agentID string) <-chan types.TrajectoryLogItem {
	h.mu.Lock()
	defer h.mu.Unlock()

	cutoff := time.Now().Add(-trajectoryTTL).UnixMilli()

	// Gather hot entries within TTL.
	var hot []types.TrajectoryLogItem
	if ac := h.hot[agentID]; ac != nil {
		for _, item := range ac.items {
			if item.Timestamp >= cutoff {
				hot = append(hot, item)
			}
		}
	}

	items := compactTrajectoryLogItems(append(readCold(agentID, cutoff, -1), hot...))
	items = latestTrajectoryItems(items, trajectoryReplayLimit)

	ch := make(chan types.TrajectoryLogItem, len(items)+trajectoryLiveBuffer)
	for _, item := range items {
		ch <- item
	}

	if h.subscribers[agentID] == nil {
		h.subscribers[agentID] = make(map[chan types.TrajectoryLogItem]struct{})
	}
	h.subscribers[agentID][ch] = struct{}{}

	return ch
}

// SubscribeLive returns a channel that receives only new trajectory events.
func (h *trajectoryHub) SubscribeLive(agentID string) <-chan types.TrajectoryLogItem {
	h.mu.Lock()
	defer h.mu.Unlock()

	ch := make(chan types.TrajectoryLogItem, trajectoryLiveBuffer)
	if h.subscribers[agentID] == nil {
		h.subscribers[agentID] = make(map[chan types.TrajectoryLogItem]struct{})
	}
	h.subscribers[agentID][ch] = struct{}{}
	return ch
}

// SubscribeAllLive returns a channel that receives new trajectory events for all agents.
func (h *trajectoryHub) SubscribeAllLive() <-chan types.TrajectoryLogItem {
	h.mu.Lock()
	defer h.mu.Unlock()

	ch := make(chan types.TrajectoryLogItem, trajectoryLiveBuffer)
	h.allSubs[ch] = struct{}{}
	return ch
}

// Unsubscribe removes and closes a subscriber channel.
func (h *trajectoryHub) Unsubscribe(agentID string, ch <-chan types.TrajectoryLogItem) {
	h.mu.Lock()
	defer h.mu.Unlock()

	subs := h.subscribers[agentID]
	for sub := range subs {
		if sub == ch {
			delete(subs, sub)
			close(sub)
			return
		}
	}
}

// UnsubscribeAll removes and closes an all-agent subscriber channel.
func (h *trajectoryHub) UnsubscribeAll(ch <-chan types.TrajectoryLogItem) {
	h.mu.Lock()
	defer h.mu.Unlock()

	for sub := range h.allSubs {
		if sub == ch {
			delete(h.allSubs, sub)
			close(sub)
			return
		}
	}
}

// Recent returns the latest n trajectory items for the given agent from the
// hot cache. If fewer than n items are available it returns all of them.
// The returned slice is ordered oldest-first (same order as the hot ring).
func (h *trajectoryHub) Recent(agentID string, n int) []types.TrajectoryLogItem {
	h.mu.Lock()
	defer h.mu.Unlock()

	ac := h.hot[agentID]
	if ac == nil || len(ac.items) == 0 {
		return nil
	}

	items := latestTrajectoryItems(compactTrajectoryLogItems(ac.items), n)

	out := make([]types.TrajectoryLogItem, len(items))
	copy(out, items)
	return out
}

// History returns recent trajectory items ordered newest-first. If before is
// positive, only entries older than that timestamp are returned.
func (h *trajectoryHub) History(agentID string, limit int, before int64) []types.TrajectoryLogItem {
	if limit <= 0 {
		return nil
	}

	cutoff := time.Now().Add(-trajectoryTTL).UnixMilli()

	h.mu.Lock()
	var hot []types.TrajectoryLogItem
	if ac := h.hot[agentID]; ac != nil {
		hot = append(hot, ac.items...)
	}
	h.mu.Unlock()

	items := append(readCold(agentID, cutoff, -1), hot...)
	filtered := make([]types.TrajectoryLogItem, 0, len(items))
	for _, item := range items {
		if item.Timestamp < cutoff {
			continue
		}
		if before > 0 && item.Timestamp >= before {
			continue
		}
		filtered = append(filtered, item)
	}

	compacted := compactTrajectoryLogItems(filtered)
	sort.SliceStable(compacted, func(i, j int) bool {
		return compacted[i].Timestamp > compacted[j].Timestamp
	})
	if len(compacted) > limit {
		compacted = compacted[:limit]
	}

	out := make([]types.TrajectoryLogItem, len(compacted))
	copy(out, compacted)
	return out
}

func mergeTrajectoryEntries(previous types.TrajectoryEntry, next types.TrajectoryEntry) (types.TrajectoryEntry, bool) {
	if previous.Kind == consts.TrajectoryKindContextUsage &&
		next.Kind == consts.TrajectoryKindContextUsage &&
		previous.SessionID == next.SessionID {
		return next, true
	}

	if sameSubagentActivityItem(previous, next) {
		merged := previous
		if strings.TrimSpace(next.Activity) != "" {
			merged.Activity = next.Activity
		}
		if strings.TrimSpace(next.Detail) != "" {
			merged.Detail = next.Detail
		}
		merged.Metadata = mergeTrajectoryMetadata(previous.Metadata, next.Metadata)
		return merged, true
	}

	if previous.Kind == consts.ParsedEventKindText && next.Kind == consts.ParsedEventKindText {
		merged := previous
		merged.Text += next.Text
		return merged, true
	}

	if previous.Kind == consts.ParsedEventKindThinking && next.Kind == consts.ParsedEventKindThinking {
		if strings.TrimSpace(previous.SectionID) != strings.TrimSpace(next.SectionID) {
			return types.TrajectoryEntry{}, false
		}
		merged := previous
		merged.Text = mergeThinkingText(previous.Text, next.Text)
		if strings.TrimSpace(next.Title) != "" {
			merged.Title = next.Title
		}
		return merged, true
	}

	if previous.Kind == consts.ParsedEventKindToolOutput &&
		next.Kind == consts.ParsedEventKindToolOutput &&
		previous.ToolName == next.ToolName &&
		previous.ToolCallID == next.ToolCallID {
		merged := previous
		merged.Text = mergeToolOutputText(previous.Text, next.Text)
		if len(merged.Metadata) == 0 && len(next.Metadata) > 0 {
			merged.Metadata = next.Metadata
		}
		return merged, true
	}

	return types.TrajectoryEntry{}, false
}

func mergeTrajectoryMetadata(previous, next map[string]any) map[string]any {
	if len(previous) == 0 && len(next) == 0 {
		return nil
	}
	merged := make(map[string]any, len(previous)+len(next))
	for key, value := range previous {
		merged[key] = value
	}
	for key, value := range next {
		merged[key] = value
	}
	return merged
}

func subagentActivityItemID(entry types.TrajectoryEntry) (string, bool) {
	if entry.Kind != consts.TrajectoryKindStatus || entry.Metadata == nil {
		return "", false
	}
	runtimeItemType, ok := entry.Metadata["runtimeItemType"].(string)
	if !ok || strings.TrimSpace(runtimeItemType) != "subAgentActivity" {
		return "", false
	}
	itemID, _ := entry.Metadata["itemId"].(string)
	return strings.TrimSpace(itemID), true
}

func sameSubagentActivityItem(previous types.TrajectoryEntry, next types.TrajectoryEntry) bool {
	previousItemID, previousIsSubagentActivity := subagentActivityItemID(previous)
	nextItemID, nextIsSubagentActivity := subagentActivityItemID(next)
	return previousIsSubagentActivity && nextIsSubagentActivity &&
		previousItemID != "" && previousItemID == nextItemID
}

func mergeThinkingText(previous, next string) string {
	return previous + next
}

func mergeToolOutputText(previous, next string) string {
	if next == "" {
		return previous
	}
	if previous == "" {
		return next
	}
	if previous == next || strings.HasPrefix(previous, next) {
		return previous
	}
	if strings.HasPrefix(next, previous) {
		return next
	}
	return previous + next
}

func compactTrajectoryLogItems(items []types.TrajectoryLogItem) []types.TrajectoryLogItem {
	if len(items) == 0 {
		return nil
	}

	out := make([]types.TrajectoryLogItem, 0, len(items))
	for _, item := range items {
		previousIndex := len(out) - 1
		if previousIndex >= 0 {
			previous := out[previousIndex]
			if previous.AgentID == item.AgentID {
				if merged, ok := mergeTrajectoryEntries(previous.Entry, item.Entry); ok {
					timestamp := previous.Timestamp
					if merged.Kind == consts.TrajectoryKindContextUsage {
						timestamp = item.Timestamp
					}
					out[previousIndex] = types.TrajectoryLogItem{
						AgentID:   previous.AgentID,
						Entry:     merged,
						Timestamp: timestamp,
					}
					continue
				}

				_, previousIsSubagentActivity := subagentActivityItemID(previous.Entry)
				_, nextIsSubagentActivity := subagentActivityItemID(item.Entry)
				if !previousIsSubagentActivity && !nextIsSubagentActivity &&
					previous.Entry.Kind == consts.TrajectoryKindStatus &&
					item.Entry.Kind == consts.TrajectoryKindStatus &&
					previous.Entry.Activity == item.Entry.Activity &&
					previous.Entry.Detail == item.Entry.Detail {
					continue
				}
			}
		}

		out = append(out, item)
	}

	return out
}

// LatestByAgent returns the latest trajectory item for each requested agent.
func (h *trajectoryHub) LatestByAgent(agentIDs []string) map[string]types.TrajectoryLogItem {
	out := make(map[string]types.TrajectoryLogItem, len(agentIDs))
	for _, agentID := range agentIDs {
		items := h.History(agentID, 1, 0)
		if len(items) > 0 {
			out[agentID] = items[0]
		}
	}
	return out
}

// ClearAgent removes all hot history and closes all subscribers.
// Cold files are left for the GC to clean up by TTL.
func (h *trajectoryHub) ClearAgent(agentID string) {
	h.mu.Lock()
	defer h.mu.Unlock()

	if ac := h.hot[agentID]; ac != nil {
		h.hotBytes -= ac.sizeBytes
		if h.hotBytes < 0 {
			h.hotBytes = 0
		}
	}
	delete(h.hot, agentID)
	subs := h.subscribers[agentID]
	for ch := range subs {
		close(ch)
	}
	delete(h.subscribers, agentID)
}

// -----------------------------------------------------------------------
// Background GC
// -----------------------------------------------------------------------

func (h *trajectoryHub) gcLoop() {
	defer close(h.gcDone)
	ticker := time.NewTicker(trajectoryGCInterval)
	defer ticker.Stop()
	for {
		select {
		case <-h.stopGC:
			return
		case <-ticker.C:
			h.evictExpiredHot()
			evictExpiredCold()
		}
	}
}

// evictExpiredHot drops in-memory entries older than the TTL.
func (h *trajectoryHub) evictExpiredHot() {
	cutoff := time.Now().Add(-trajectoryTTL).UnixMilli()
	h.mu.Lock()
	defer h.mu.Unlock()

	for agentID, ac := range h.hot {
		trimmed, removed := trimBefore(ac.items, cutoff)
		ac.items = trimmed
		ac.sizeBytes -= removed
		h.hotBytes -= removed
		if ac.sizeBytes < 0 {
			ac.sizeBytes = 0
		}
		if h.hotBytes < 0 {
			h.hotBytes = 0
		}
		if len(ac.items) == 0 && len(h.subscribers[agentID]) == 0 {
			delete(h.hot, agentID)
		}
	}
}

// evictExpiredCold rewrites each cold NDJSON file keeping only entries
// within the TTL. Files that become empty are deleted.
func evictExpiredCold() {
	dir := coldDir()
	entries, err := os.ReadDir(dir)
	if err != nil {
		return
	}
	cutoff := time.Now().Add(-trajectoryTTL).UnixMilli()

	for _, de := range entries {
		if de.IsDir() {
			continue
		}
		name := de.Name()
		if filepath.Ext(name) != ".ndjson" {
			continue
		}
		agentID := name[:len(name)-len(".ndjson")]
		kept := readCold(agentID, cutoff, -1)
		path := coldPath(agentID)
		if len(kept) == 0 {
			if err := os.Remove(path); err != nil {
				slog.Warn("trajectory: remove expired cold file", "path", path, "error", err)
			}
			continue
		}
		// Rewrite in-place via temp file.
		tmp := path + ".tmp"
		f, err := os.Create(tmp)
		if err != nil {
			slog.Warn("trajectory: create temp cold file", "path", tmp, "error", err)
			continue
		}
		enc := json.NewEncoder(f)
		for i := range kept {
			_ = enc.Encode(&kept[i])
		}
		f.Close()
		if err := os.Rename(tmp, path); err != nil {
			slog.Warn("trajectory: rename temp cold file", "error", err)
			os.Remove(tmp)
		}
	}
}

// -----------------------------------------------------------------------
// Helpers
// -----------------------------------------------------------------------

// estimateItemBytes returns a rough byte size of a serialised TrajectoryLogItem.
func estimateItemBytes(item types.TrajectoryLogItem) int {
	// Fixed overhead (JSON keys, agentID, timestamp, punctuation): ~80 bytes.
	const overhead = 80
	e := item.Entry
	return overhead + len(e.Kind) + len(e.Text) + len(e.ToolName) + len(e.ToolInput) + len(e.Activity) + len(e.Detail)
}

// trimBefore returns the suffix of items whose Timestamp >= cutoff and the
// estimated byte size of the removed prefix.
func trimBefore(items []types.TrajectoryLogItem, cutoff int64) ([]types.TrajectoryLogItem, int) {
	removed := 0
	for i, item := range items {
		if item.Timestamp >= cutoff {
			return items[i:], removed
		}
		removed += estimateItemBytes(item)
	}
	return nil, removed
}

func latestTrajectoryItems(items []types.TrajectoryLogItem, limit int) []types.TrajectoryLogItem {
	if limit <= 0 || len(items) <= limit {
		return items
	}
	return items[len(items)-limit:]
}