package memory

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

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

// Working memory sizing and summarization constants.
const (
	// defaultWorkingWindow is the default number of messages retained in the
	// sliding window when NewWorkingMemory is called with a non-positive size.
	defaultWorkingWindow = 20

	// summaryLastTextMax is the maximum rune length of the trailing message text
	// included in a naive summary, keeping self-summaries compact.
	summaryLastTextMax = 80
)

// WorkingMemory is the per-session short-term memory: a sliding window of the
// current conversation's messages plus a summary (design §0 layer ①). It is the
// immediate context an agent reasons over.
//
// The window is bounded: adding beyond the maximum evicts the oldest message.
// Summarize returns the summary set via SetSummary, or a naive structural
// summary computed on demand. LLM-generated summaries are deferred to M4; the
// plumbing is in place so a summarizer can call SetSummary.
//
// WorkingMemory is safe for concurrent use; Add/SetSummary serialize with a
// write lock and Window/Summarize read under a read lock.
type WorkingMemory struct {
	mu      sync.RWMutex
	window  []model.Message
	max     int
	summary string
}

// NewWorkingMemory builds a working memory with a sliding window of up to
// maxMessages messages. A non-positive maxMessages selects defaultWorkingWindow.
func NewWorkingMemory(maxMessages int) *WorkingMemory {
	if maxMessages <= 0 {
		maxMessages = defaultWorkingWindow
	}
	return &WorkingMemory{max: maxMessages}
}

// Add appends msg to the window and evicts the oldest message when the window
// exceeds its limit.
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 returns a deep copy of the current message window, oldest first. The
// returned slice and each message's content parts are detached from internal
// state, so callers may mutate them freely.
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 returns a copy of m whose content parts are detached from the
// source. Every field of model.Content is a value type, so copying the parts
// slice is sufficient.
func cloneMessage(m model.Message) model.Message {
	parts := make([]model.Content, len(m.Content))
	copy(parts, m.Content)
	m.Content = parts
	return m
}

// SetSummary records an explicit summary (e.g. produced by an LLM in M4).
func (w *WorkingMemory) SetSummary(s string) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.summary = s
}

// Summarize returns the stored summary when one was set, otherwise a naive
// structural summary derived from the current window. It never mutates state, so
// repeated calls return a value consistent with the window at call time.
func (w *WorkingMemory) Summarize() string {
	w.mu.RLock()
	defer w.mu.RUnlock()

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

// naiveSummaryLocked builds "{n} 条消息, {role} {count}, ...; 最近: {text}". It
// requires w.mu (read or write) to be held.
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 shortens s to at most n runes. Unicode-safe: it never slices a
// multi-byte rune in half.
func truncateRunes(s string, n int) string {
	if utf8.RuneCountInString(s) <= n {
		return s
	}
	runes := []rune(s)
	return string(runes[:n])
}