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
)

// WorkingMemory 是每会话的短期记忆：当前对话消息的
// 滑动窗口加一份摘要（design §0 layer ①）。它是
// 代理推理所依据的即时上下文。
//
// 窗口有界：超出上限的添加会淘汰最旧的消息。
// Summarize 返回经 SetSummary 设置的摘要，或按需计算的
// 朴素结构化摘要。LLM 生成的摘要推迟到 M4；管线
// 已就位，摘要器可调用 SetSummary。
//
// WorkingMemory 并发安全；Add/SetSummary 以写锁串行，
// Window/Summarize 在读锁下读取。
type WorkingMemory struct {
	mu      sync.RWMutex
	window  []model.Message
	max     int
	summary string
}

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

// Add 把 msg 追加到窗口，并在窗口超出上限时
// 淘汰最旧的消息。
func (w *WorkingMemory) Add(msg model.Message) {
	w.mu.Lock()
	defer w.mu.Unlock()

	w.window = append(w.window, msg)
	if len(w.window) > w.max {
		overflow := len(w.window) - w.max
		copy(w.window, w.window[overflow:])
		w.window = w.window[:w.max]
	}
}

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

	out := make([]model.Message, len(w.window))
	for i, m := range w.window {
		out[i] = cloneMessage(m)
	}
	return out
}

// 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
}

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

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

	if w.summary != "" {
		return w.summary
	}
	return w.naiveSummaryLocked()
}

// naiveSummaryLocked 构造 "{n} 条消息, {role} {count}, ...; 最近: {text}"。
// 它要求持有 w.mu（读锁或写锁）。
func (w *WorkingMemory) naiveSummaryLocked() string {
	if len(w.window) == 0 {
		return ""
	}
	counts := make(map[model.Role]int)
	last := ""
	for _, m := range w.window {
		counts[m.Role]++
		last = messageText(m)
	}

	var b strings.Builder
	b.WriteString(strconv.Itoa(len(w.window)))
	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])
}