package loop

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

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

// Compaction tuning constants. These bound the compactor so it stays
// deterministic and offline: the summary is capped, the recent tail is a share
// of the budget, and very short conversations are left alone even when they
// exceed the budget (design §1.3 上下文溢出压缩).
const (
	// DefaultMaxContextTokens is the context budget a caller can assign to
	// Options.MaxContextTokens to enable overflow compaction (design §1.3).
	DefaultMaxContextTokens = 32768

	// keepRecentFraction is the share of the budget reserved for the recent tail
	// during compaction; the remainder is left for the summary message and prompt.
	keepRecentFraction = 0.5

	// minKeepTokens is the floor for the recent-tail budget so a tiny
	// MaxContextTokens does not force every turn into the summary.
	minKeepTokens = 128

	// summaryMarker tags a compaction summary system message so the compactor can
	// recognize an existing summary prefix and preserve (merge into) it rather
	// than re-summarize it on a later compaction.
	summaryMarker = "[前文摘要 · history summary]"

	// maxSummaryChars caps the compaction summary so repeated compactions cannot
	// inflate the summary without bound (design §1.3 摘要容量上限).
	maxSummaryChars = 4000

	// minCompactMessages is the smallest conversation the compactor will touch; a
	// shorter one is left alone even when it is over budget.
	minCompactMessages = 3

	// messageOverheadTokens is the per-message token overhead charged for the role
	// and separator tokens the provider adds around each message.
	messageOverheadTokens = 4

	// imageTokens is the rough token charge for a single image part.
	imageTokens = 1000

	// charsPerToken is the character-to-token divisor used by the offline estimate.
	charsPerToken = 4

	// summarySampleChars bounds the per-message text sampled into a naive summary.
	summarySampleChars = 120
)

// Compactor compacts an over-budget conversation by summarizing older turns into
// a single leading summary message while preserving the recent tail — including
// the trailing tool results — unmodified (design §1.3, aligned with OpenClaw's
// context-engine). It is injected through Options.Compactor so a caller can swap
// in an LLM-backed summarizer; the loop ships a structural default that never
// calls a model.
type Compactor interface {
	// Compact returns messages with older turns compressed into a leading summary
	// when the estimated tokens exceed budget. It returns the input unchanged when
	// the messages already fit within budget or there is nothing safe to
	// summarize. ctx is the abort signal: a canceled context stops compaction and
	// returns a wrapped context.Canceled so the caller can bail out.
	Compact(ctx context.Context, messages []model.Message, budget int) ([]model.Message, error)
}

// defaultCompactor is the built-in structural compactor. It is stateless and
// safe for concurrent use.
type defaultCompactor struct{}

// Compact implements Compactor.
func (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)
}

// compactMessages is the shared compaction implementation. It never mutates the
// input slices; a compacted result is always freshly allocated.
func compactMessages(ctx context.Context, messages []model.Message, budget int) ([]model.Message, error) {
	if budget <= 0 || len(messages) < minCompactMessages {
		return messages, nil
	}
	if estimateContextTokens(messages) <= budget {
		return messages, nil
	}

	prefixEnd := summaryPrefixEnd(messages)
	keepTokens := max(int(float64(budget)*keepRecentFraction), minKeepTokens)
	keepStart := findKeepStart(messages, prefixEnd, keepTokens)
	if keepStart <= prefixEnd {
		// Everything is being kept; there is no safe prefix to summarize.
		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 returns the count of leading compaction-summary messages,
// which are never re-summarized. It is 0 (no summary) or 1 (a single summary
// prefix), since compaction always emits exactly one leading summary.
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 returns the index of the first message to keep given that the
// recent tail should fit within keepTokens. It always keeps whole turns: a turn
// starts at a user message and runs to the next user message. It scores turns
// from the back, so the most recent turns survive and a trailing tool result is
// never sliced apart. When no user turn exists (unusual), it keeps nothing extra
// and returns from, so nothing is summarized.
func findKeepStart(messages []model.Message, from, keepTokens 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]
		}
		if i == len(turns)-1 || acc+estimateContextTokens(messages[turns[i]:to]) <= keepTokens {
			acc += estimateContextTokens(messages[turns[i]:to])
			keptStart = turns[i]
		} else {
			break
		}
	}
	return keptStart
}

// estimateContextTokens returns a conservative token estimate for messages. It is
// a heuristic (chars ÷ charsPerToken with a per-message and per-image overhead)
// so it stays dependency-free and offline, but it is monotonic enough to drive
// the compact decision.
func estimateContextTokens(messages []model.Message) 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
			}
		}
	}
	return total
}

// estimateStringTokens is the character-level token estimate for one text string,
// weighting wide (CJK) runes double and dividing by charsPerToken. An empty string
// costs nothing; a non-empty one costs at least one 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 renders a compact structural summary of the given messages: a role
// distribution plus the most recent user and assistant text. It is the default
// (offline, no-model) summarizer; a caller that wants a semantic digest injects
// its own Compactor instead.
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 truncates s to at most max runes, appending a "..." marker when it
// had to cut. It is rune-safe so it never slices a multi-byte rune in half.
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 shortens s to at most n runes, Unicode-safe.
func truncateRunes(s string, n int) string {
	runes := []rune(s)
	if len(runes) <= n {
		return s
	}
	return string(runes[:n])
}