package loop

import (
	"context"
	"fmt"
	"strings"
	"unicode/utf8"

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

// 压缩调优常量。它们界定 compactor，使其保持确定性与
// 离线：摘要被限额、近期尾部占预算的一个比例，
// 且极短的对话即使超出预算也原样不动
// （design §1.3 上下文溢出压缩）。
const (
	// DefaultMaxContextTokens 是调用方可赋给
	// Options.MaxContextTokens 以启用溢出压缩的上下文预算（design §1.3）。
	DefaultMaxContextTokens = 32768

	// keepRecentFraction 是压缩时为近期尾部保留的
	// 预算比例；余下部分留给摘要消息与提示词。
	keepRecentFraction = 0.5

	// minKeepTokens 是近期尾部预算的下限，使极小的
	// MaxContextTokens 不会把每个 turn 都逼进摘要。
	minKeepTokens = 128

	// summaryMarker 标记压缩摘要系统消息，使 compactor 能
	// 识别已存在的摘要前缀并保留（并入）它，
	// 而不是在后续压缩中重新摘要。
	summaryMarker = "[前文摘要 · history summary]"

	// maxSummaryChars 限制压缩摘要，使反复压缩无法
	// 无界地膨胀摘要（design §1.3 摘要容量上限）。
	maxSummaryChars = 4000

	// minCompactMessages 是 compactor 会触碰的最小对话；
	// 更短的即使超出预算也原样不动。
	minCompactMessages = 3

	// messageOverheadTokens 是每条消息的 token 开销，对应 provider
	// 在每条消息周围添加的角色与分隔符 token。
	messageOverheadTokens = 4

	// imageTokens 是单个图像部件的粗略 token 计数。
	imageTokens = 1000

	// charsPerToken 是离线估算使用的字符到 token 除数。附件注入预算
	// （DocumentInjectCharsFor）与压缩估算共用它，是本仓库唯一的字符 ↔ token
	// 换算口径。
	charsPerToken = 4

	// CharsPerToken 是 charsPerToken 的**导出别名**：组合根里按同一口径换算的
	// 地方（cmd/cyrene 的输入护栏上限推导：上限 = 推导预算 token 数 × 本常量）
	// 必须复用这一处定义，而不是另写一个字面量 4 —— 护栏上限与附件注入预算
	// 曾因两边各写一套换算（1 char = 1 token vs 4 char = 1 token）而对不上，
	// 表现是「2 份各自撞满注入预算的附件就被整回合拦死」（切片 G2b）。
	CharsPerToken = charsPerToken

	// summarySampleChars 限制被抽样进朴素摘要的每条消息文本长度。
	summarySampleChars = 120
)

// Compactor 通过把更早的轮次摘要成单条前导摘要消息来压缩
// 超预算的对话，同时原样保留近期尾部 —— 包括
// 尾部的工具结果（design §1.3，与 OpenClaw 的
// context-engine 对齐）。它经 Options.Compactor 注入，使调用方能
// 换成 LLM 支持的摘要器；循环自带一个绝不调用
// 模型的结构化默认实现。
type Compactor interface {
	// Compact 在估算 token 超过 budget 时返回把更早轮次压缩为
	// 一条前导摘要的 messages。消息本来就在预算内，或
	// 没有可安全摘要的内容时，它原样返回输入。
	// ctx 是中止信号：被取消的 context 停止压缩并
	// 返回包装过的 context.Canceled，使调用方能退出。
	Compact(ctx context.Context, messages []model.Message, budget int) ([]model.Message, error)
}

// defaultCompactor 是内置的结构化 compactor。它无状态（documentInjectChars 只在
// 构造时给定），且并发安全。
type defaultCompactor struct {
	// documentInjectChars 是单份附件注入上下文的字符预算，用于估算文档部件的
	// token 代价（documentPartTokens）。0 表示用 defaultDocumentInjectChars。
	documentInjectChars int
}

// Compact 实现 Compactor。
func (c defaultCompactor) Compact(ctx context.Context, messages []model.Message, budget int) ([]model.Message, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return nil, fmt.Errorf("loop: compact: %w", err)
	}
	return compactMessages(ctx, messages, budget, c.documentInjectChars)
}

// compactMessages 是共享的压缩实现。它绝不改写
// 输入切片；压缩结果始终是新分配的。
func compactMessages(ctx context.Context, messages []model.Message, budget, docInjectChars int) ([]model.Message, error) {
	if budget <= 0 || len(messages) < minCompactMessages {
		return messages, nil
	}
	if estimateContextTokensWith(messages, docInjectChars) <= budget {
		return messages, nil
	}

	prefixEnd := summaryPrefixEnd(messages)
	keepTokens := max(int(float64(budget)*keepRecentFraction), minKeepTokens)
	keepStart := findKeepStart(messages, prefixEnd, keepTokens, docInjectChars)
	if keepStart <= prefixEnd {
		// 全部内容都被保留；没有可安全摘要的前缀。
		return messages, nil
	}

	if err := ctx.Err(); err != nil {
		return nil, fmt.Errorf("loop: compact: %w", err)
	}

	existing := ""
	if prefixEnd > 0 {
		existing = strings.TrimPrefix(messageText(messages[0]), summaryMarker+"\n\n")
	}
	fresh := buildSummary(messages[prefixEnd:keepStart])
	if fresh == "" && existing == "" {
		return messages, nil
	}
	full := fresh
	if existing != "" {
		full = existing + "\n\n" + fresh
	}

	summary := model.TextMessage(model.RoleSystem, summaryMarker+"\n\n"+capSummary(full, maxSummaryChars))
	out := make([]model.Message, 0, len(messages)-(keepStart-prefixEnd)+1)
	out = append(out, summary)
	out = append(out, messages[keepStart:]...)
	return out, nil
}

// summaryPrefixEnd 返回前导压缩摘要消息的数量，
// 它们绝不再次被摘要。取值为 0（无摘要）或 1（单条摘要
// 前缀），因为压缩始终只产出一条前导摘要。
func summaryPrefixEnd(messages []model.Message) int {
	if len(messages) == 0 || messages[0].Role != model.RoleSystem {
		return 0
	}
	if strings.HasPrefix(messageText(messages[0]), summaryMarker) {
		return 1
	}
	return 0
}

// findKeepStart 在近期尾部应装入 keepTokens 的前提下，
// 返回第一条要保留的消息索引。它始终保留完整的 turn：一个 turn
// 从用户消息开始，直到下一条用户消息。它从后往前
// 给 turn 计分，因此最近的 turn 得以存活，尾部的工具结果
// 绝不被切开。不存在用户 turn 时（少见），它不额外保留任何东西
// 并从 from 返回，因此什么也不被摘要。
func findKeepStart(messages []model.Message, from, keepTokens, docInjectChars int) int {
	n := len(messages)
	turns := make([]int, 0, 8)
	for i := from; i < n; i++ {
		if messages[i].Role == model.RoleUser {
			turns = append(turns, i)
		}
	}
	if len(turns) == 0 {
		return from
	}

	acc := 0
	keptStart := turns[len(turns)-1]
	for i := len(turns) - 1; i >= 0; i-- {
		to := n
		if i+1 < len(turns) {
			to = turns[i+1]
		}
		size := estimateContextTokensWith(messages[turns[i]:to], docInjectChars)
		if i == len(turns)-1 || acc+size <= keepTokens {
			acc += size
			keptStart = turns[i]
		} else {
			break
		}
	}
	return keptStart
}

// estimateContextTokens 返回 messages 的保守 token 估算，文档部件按**缺省**注入
// 上限计入（defaultDocumentInjectChars）。需要按本回合真实注入预算估算的调用方
// （compactor）请用 estimateContextTokensWith。
func estimateContextTokens(messages []model.Message) int {
	return estimateContextTokensWith(messages, defaultDocumentInjectChars)
}

// estimateContextTokensWith 返回 messages 的保守 token 估算。它是
// 启发式（字符 ÷ charsPerToken，加上每条消息、每个图像与每个文档的开销），
// 因此无依赖且离线，但单调性足够驱动压缩决策。
// docInjectChars 是单份文档注入上下文的字符上限，决定文档部件在估算里最多
// 算多少（0 = 缺省值）。
func estimateContextTokensWith(messages []model.Message, docInjectChars int) int {
	total := 0
	for _, m := range messages {
		total += messageOverheadTokens
		for _, part := range m.Content {
			switch part.Type {
			case model.ContentTypeText:
				total += estimateStringTokens(part.Text)
			case model.ContentTypeImage:
				total += imageTokens
			case model.ContentTypeDocument:
				total += documentPartTokens(part, docInjectChars)
			}
		}
	}
	return total
}

// documentPartTokens 估算一个文档部件在装配后可能占用的 token。窗口里只有
// 引用（design/24 §2.4），正文要等装配期才注入，因此这里按「原文件字节数
// （UTF-8 文本下 ≈ 字符数），上限为注入预算」估算。
//
// 关键是文档部件**不是免费的**：只认 ContentTypeImage 的分支会让纯文档回合
// 被估成 0 token，在压缩里等同于空内容（design/24 §9.5 移交的静默吞掉）。
// 大文件（尤其在预算处截断注入、或压根读不到）会被这个估计偏高地计入，方向
// 是「少放内容进上下文」——保守收紧，而不是溢出模型窗口。
func documentPartTokens(part model.Content, docInjectChars int) int {
	if docInjectChars <= 0 {
		docInjectChars = defaultDocumentInjectChars
	}
	chars := part.DocumentSize
	if chars > int64(docInjectChars) {
		chars = int64(docInjectChars)
	}
	if chars < 0 {
		chars = 0
	}
	t := int(chars) / charsPerToken
	if t < documentMinTokens {
		return documentMinTokens
	}
	return t
}

// estimateStringTokens 是单个文本字符串的字符级 token 估算，
// 把宽字符（CJK）rune 计为双倍权重再除以 charsPerToken。空字符串
// 不花费任何 token；非空至少花费一个 token。
func estimateStringTokens(s string) int {
	if s == "" {
		return 0
	}
	runes := utf8.RuneCountInString(s)
	wide := 0
	for _, r := range s {
		if r >= 0x2E80 {
			wide++
		}
	}
	t := (runes + wide) / charsPerToken
	if t < 1 {
		t = 1
	}
	return t
}

// buildSummary 渲染给定 messages 的紧凑结构化摘要：角色
// 分布，加上最近的用户与 assistant 文本。它是默认的
// （离线、不调模型）摘要器；想要语义化摘要的调用方
// 改为注入自己的 Compactor。
func buildSummary(messages []model.Message) string {
	if len(messages) == 0 {
		return ""
	}
	counts := map[model.Role]int{}
	lastUser, lastAssistant := "", ""
	for _, m := range messages {
		counts[m.Role]++
		text := messageText(m)
		if text == "" {
			continue
		}
		switch m.Role {
		case model.RoleUser:
			lastUser = text
		case model.RoleAssistant:
			lastAssistant = text
		}
	}

	var b strings.Builder
	fmt.Fprintf(&b, "已压缩 %d 条历史消息（", len(messages))
	first := true
	for _, role := range []model.Role{model.RoleSystem, model.RoleUser, model.RoleAssistant, model.RoleTool} {
		if c := counts[role]; c > 0 {
			if !first {
				b.WriteString("，")
			}
			first = false
			fmt.Fprintf(&b, "%s %d", role, c)
		}
	}
	b.WriteString("）")
	if lastUser != "" {
		fmt.Fprintf(&b, "。最近用户: %s", truncateRunes(lastUser, summarySampleChars))
	}
	if lastAssistant != "" {
		fmt.Fprintf(&b, "。最近助手: %s", truncateRunes(lastAssistant, summarySampleChars))
	}
	return b.String()
}

// capSummary 把 s 截断到至多 max 个 rune，需要裁剪时追加 "..." 标记。
// 它是 rune 安全的，因此绝不把多字节 rune 切成两半。
func capSummary(s string, max int) string {
	runes := []rune(s)
	if len(runes) <= max {
		return s
	}
	if max < 3 {
		return string(runes[:max])
	}
	return string(runes[:max-3]) + "..."
}

// truncateRunes 把 s 缩短到至多 n 个 rune，Unicode 安全。
func truncateRunes(s string, n int) string {
	out, _ := truncateRunesAt(s, n)
	return out
}

// truncateRunesAt 与 truncateRunes 同义，另外报告是否发生了截断。
// 「是否截断」是清单行里续读游标的唯一判据（截断才给游标），因此它必须来自
// 截断动作本身，而不是「长度是否等于上限」这类再推断。
func truncateRunesAt(s string, n int) (string, bool) {
	if n <= 0 {
		return "", s != ""
	}
	runes := []rune(s)
	if len(runes) <= n {
		return s, false
	}
	return string(runes[:n]), true
}