package memory

import (
	"strconv"
	"strings"
	"sync"
	"unicode/utf8"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

// 工作记忆的窗口大小与摘要常量。
const (
	// defaultWorkingWindow 是 NewWorkingMemory 以非正 size
	// 调用时滑动窗口默认保留的消息数。
	defaultWorkingWindow = 20

	// summaryLastTextMax 是朴素摘要中所含末条消息文本
	// 的最大 rune 数，用于保持自摘要紧凑。
	summaryLastTextMax = 80

	// maxWorkingSessions 是工作记忆同时保留的**会话**窗口数上限。
	// 进程内只有一个 WorkingMemory（loop 的每会话滑动窗口），而会话数
	// 随渠道消息无界增长，因此超出上限时淘汰最久未用的会话窗口，
	// 使内存有界（design/03 §0 层①）。
	maxWorkingSessions = 64
)

// window 是一条会话窗口的可变状态：消息滑动窗口 + 摘要。
//
// 它不持锁：所有方法都要求调用方已持有 WorkingMemory.mu
// （读锁或写锁），这正是窗口与摘要共用同一临界区的原因。
type window struct {
	messages []model.Message
	summary  string
}

// add 把 msg 追加到窗口，并在窗口超出 max 时淘汰最旧的消息。
func (w *window) add(msg model.Message, max int) {
	w.messages = append(w.messages, msg)
	w.trim(max)
}

// replace 用 msgs 替换窗口内容，只保留最后 max 条（msgs 是该会话
// 当前对话的完整快照时，窗口即「最近 max 条」的滑窗）。传入的消息
// 逐条深拷贝，因此调用方后续修改自己的切片不会影响窗口。
func (w *window) replace(msgs []model.Message, max int) {
	out := make([]model.Message, 0, len(msgs))
	for _, m := range msgs {
		out = append(out, cloneMessage(m))
	}
	w.messages = out
	w.trim(max)
}

// trim 只保留窗口最后 max 条消息。
func (w *window) trim(max int) {
	if max <= 0 || len(w.messages) <= max {
		return
	}
	overflow := len(w.messages) - max
	copy(w.messages, w.messages[overflow:])
	w.messages = w.messages[:max]
}

// snapshot 返回窗口消息的深拷贝，最旧的在前。返回的切片与每条
// 消息的 content 部分都与内部状态脱离，因此调用方可以随意修改。
func (w *window) snapshot() []model.Message {
	out := make([]model.Message, len(w.messages))
	for i, m := range w.messages {
		out[i] = cloneMessage(m)
	}
	return out
}

// summarize 在已设置摘要时返回存储的摘要，否则返回由当前
// 窗口推导的朴素结构化摘要。它从不修改状态。
func (w *window) summarize() string {
	if w.summary != "" {
		return w.summary
	}
	return naiveSummary(w.messages)
}

// sessionWindow 是一条会话的窗口及其最近使用序号，序号供
// 超出上限时淘汰最久未用的会话（见 maxWorkingSessions）。
type sessionWindow struct {
	win      window
	lastUsed uint64
}

// WorkingMemory 是每会话的短期记忆：当前对话消息的
// 滑动窗口加一份摘要（design §0 layer ①）。它是
// 代理推理所依据的即时上下文。
//
// 窗口有界：超出上限的添加会淘汰最旧的消息。
// Summarize 返回经 SetSummary 设置的摘要，或按需计算的
// 朴素结构化摘要。LLM 生成的摘要推迟到 M4；管线
// 已就位，摘要器可调用 SetSummary。
//
// 两种使用方式并存：
//
//   - 默认窗口（Add/Window/SetSummary/Summarize）：无会话键的
//     单窗口路径，保持该类型最初的语义。
//   - 会话窗口（RecordSession/SessionWindow/SummarizeSession/
//     SetSessionSummary）：**按会话 id 隔离**的窗口，循环在回合
//     落定时写入、在组装上下文时读取。进程内只有一个
//     WorkingMemory 实例（loop.Options.WorkingMemory），因此没有
//     会话键隔离时，群聊/陌生人的窗口会混进主人的工作记忆 ——
//     这是隔离红线（design/03 §5），会话路径即为此而设。
//
// WorkingMemory 并发安全；写路径以写锁串行，
// 读路径在读锁下读取。
type WorkingMemory struct {
	mu       sync.RWMutex
	def      window
	sessions map[string]*sessionWindow
	max      int
	seq      uint64
}

// NewWorkingMemory 构建一个工作记忆，滑动窗口最多
// maxMessages 条消息。非正 maxMessages 选取 defaultWorkingWindow。
func NewWorkingMemory(maxMessages int) *WorkingMemory {
	if maxMessages <= 0 {
		maxMessages = defaultWorkingWindow
	}
	return &WorkingMemory{max: maxMessages}
}

// Add 把 msg 追加到**默认**窗口，并在窗口超出上限时
// 淘汰最旧的消息。它写入的是无会话键的单窗口；按会话隔离的
// 写入走 RecordSession。
func (w *WorkingMemory) Add(msg model.Message) {
	w.mu.Lock()
	defer w.mu.Unlock()

	w.def.add(msg, w.max)
}

// Window 返回默认窗口当前消息的深拷贝，最旧的在前。
// 返回的切片与每条消息的 content 部分都与内部
// 状态脱离，因此调用方可以随意修改。
func (w *WorkingMemory) Window() []model.Message {
	w.mu.RLock()
	defer w.mu.RUnlock()

	return w.def.snapshot()
}

// SetSummary 记录一条显式摘要（例如 M4 中由 LLM 产出）到默认窗口。
func (w *WorkingMemory) SetSummary(s string) {
	w.mu.Lock()
	defer w.mu.Unlock()

	w.def.summary = s
}

// Summarize 在已设置摘要时返回存储的摘要，否则返回由默认窗口
// 推导的朴素结构化摘要。它从不修改状态，因此
// 重复调用返回与调用时窗口一致的值。
func (w *WorkingMemory) Summarize() string {
	w.mu.RLock()
	defer w.mu.RUnlock()

	return w.def.summarize()
}

// RecordSession 把会话 sessionID 的窗口置为 msgs 的最后
// maxMessages 条（msgs 是该会话当前对话的完整快照，因此结果
// 即「最近 N 条」的滑窗）。空会话 id 是空操作 —— 无键的写入无从
// 隔离，绝不落到默认窗口，以免污染无会话键的读点。
//
// 循环在回合落定时以该 turn 的完整对话调用它：会话恢复时，checkpoint
// 还原出的历史随同一个 turn 一起进来，窗口因此被回填，而不是
// 从空白重启（design §1 durable execution）。
func (w *WorkingMemory) RecordSession(sessionID string, msgs []model.Message) {
	if strings.TrimSpace(sessionID) == "" {
		return
	}
	w.mu.Lock()
	defer w.mu.Unlock()

	sw := w.sessionLocked(sessionID)
	sw.win.replace(msgs, w.max)
}

// SessionWindow 返回会话 sessionID 窗口的深拷贝（最旧的在前）。
// 该会话没有窗口时返回空切片。
func (w *WorkingMemory) SessionWindow(sessionID string) []model.Message {
	w.mu.RLock()
	defer w.mu.RUnlock()

	sw, ok := w.sessions[sessionID]
	if !ok {
		return []model.Message{}
	}
	return sw.win.snapshot()
}

// SetSessionSummary 为会话 sessionID 记录一条显式摘要（LLM 摘要器
// 的落点，design §3 §6）。空会话 id 是空操作。
func (w *WorkingMemory) SetSessionSummary(sessionID, s string) {
	if strings.TrimSpace(sessionID) == "" {
		return
	}
	w.mu.Lock()
	defer w.mu.Unlock()

	w.sessionLocked(sessionID).win.summary = s
}

// SummarizeSession 返回会话 sessionID 的摘要：已由 SetSessionSummary
// 显式设置时返回它，否则返回该会话窗口的朴素结构化摘要。
// 该会话没有窗口（或窗口为空）时返回 ""，因此调用方可以据此
// 跳过前导块。它从不修改状态。
func (w *WorkingMemory) SummarizeSession(sessionID string) string {
	w.mu.RLock()
	defer w.mu.RUnlock()

	sw, ok := w.sessions[sessionID]
	if !ok {
		return ""
	}
	return sw.win.summarize()
}

// sessionLocked 返回会话 sessionID 的窗口（必要时新建），并把它
// 标为最近使用；新建会使保留的会话数超出上限时，先淘汰最久
// 未用的那个，使会话窗口数有界。它要求持有 w.mu 的**写锁**。
func (w *WorkingMemory) sessionLocked(sessionID string) *sessionWindow {
	if w.sessions == nil {
		w.sessions = make(map[string]*sessionWindow)
	}
	w.seq++
	sw, ok := w.sessions[sessionID]
	if !ok {
		if len(w.sessions) >= maxWorkingSessions {
			w.evictOldestLocked()
		}
		sw = &sessionWindow{}
		w.sessions[sessionID] = sw
	}
	sw.lastUsed = w.seq
	return sw
}

// evictOldestLocked 删除最久未使用的会话窗口。它要求持有 w.mu 的
// 写锁；空 map 时是空操作。
func (w *WorkingMemory) evictOldestLocked() {
	oldestKey := ""
	var oldest uint64
	for key, sw := range w.sessions {
		if oldestKey == "" || sw.lastUsed < oldest {
			oldestKey, oldest = key, sw.lastUsed
		}
	}
	if oldestKey != "" {
		delete(w.sessions, oldestKey)
	}
}

// cloneMessage 返回 m 的副本，其 content 部分与源脱离。
// model.Content 的每个字段都是值类型，因此复制 parts
// 切片就足够了。
func cloneMessage(m model.Message) model.Message {
	parts := make([]model.Content, len(m.Content))
	copy(parts, m.Content)
	m.Content = parts
	return m
}

// naiveSummary 构造 "{n} 条消息, {role} {count}, ...; 最近: {text}"。
// 它是纯函数：不读取也不修改任何窗口状态。
func naiveSummary(messages []model.Message) string {
	if len(messages) == 0 {
		return ""
	}
	counts := make(map[model.Role]int)
	last := ""
	for _, m := range messages {
		counts[m.Role]++
		last = messageText(m)
	}

	var b strings.Builder
	b.WriteString(strconv.Itoa(len(messages)))
	b.WriteString(" 条消息")
	for _, role := range []model.Role{model.RoleSystem, model.RoleUser, model.RoleAssistant, model.RoleTool} {
		if c := counts[role]; c > 0 {
			b.WriteString(", ")
			b.WriteString(string(role))
			b.WriteByte(' ')
			b.WriteString(strconv.Itoa(c))
		}
	}
	if last != "" {
		b.WriteString("; 最近: ")
		b.WriteString(truncateRunes(last, summaryLastTextMax))
	}
	return b.String()
}

// truncateRunes 把 s 截到最多 n 个 rune。Unicode 安全：它绝不把
// 多字节 rune 切一半。
func truncateRunes(s string, n int) string {
	if utf8.RuneCountInString(s) <= n {
		return s
	}
	runes := []rune(s)
	return string(runes[:n])
}