package agent

import (
	"context"
	"errors"
	"fmt"
	"io"
	"regexp"
	"strings"
	"testing"

	"git.gopher.ink/mirrors2/AutoZero/llm"
	"git.gopher.ink/mirrors2/AutoZero/session"
)

// TestCheckpointUpdatedInPlace：第二次压缩要把旧 checkpoint + 新消息交给更新
// prompt，而不是把整段历史（含旧摘要）重新摘要一遍；sidecar 同步更新。
func TestCheckpointUpdatedInPlace(t *testing.T) {
	var inputs []string
	var summaryCalls int
	prov := &scriptedProvider{fn: func(_ int, system string, msgs []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		if strings.HasPrefix(system, "You are a context compaction") {
			summaryCalls++
			inputs = append(inputs, msgs[0].Content)
			return llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("CHECKPOINT-V%d", summaryCalls)}, nil
		}
		return llm.Message{Role: llm.RoleAssistant, Content: "ok"}, nil
	}}
	a, sess := newTestAgent(t)
	a.Provider = prov
	a.MinCompactTokens = 1
	a.Opt.Stderr = io.Discard

	if err := sess.Append(llm.Message{Role: llm.RoleUser, Content: "first task"}); err != nil {
		t.Fatal(err)
	}
	if err := sess.Append(llm.Message{Role: llm.RoleAssistant, Content: "working"}); err != nil {
		t.Fatal(err)
	}
	if _, err := a.CompactNow(context.Background()); err != nil {
		t.Fatal(err)
	}
	if a.checkpoint != "CHECKPOINT-V1" || a.checkpointCovered != 2 {
		t.Fatalf("first checkpoint = %q covered=%d", a.checkpoint, a.checkpointCovered)
	}
	if len(a.compacted) != 1 || !strings.Contains(a.compacted[0].Content, "CHECKPOINT-V1") {
		t.Fatalf("working set = %+v", a.compacted)
	}

	if err := sess.Append(llm.Message{Role: llm.RoleUser, Content: "more context"}); err != nil {
		t.Fatal(err)
	}
	if err := sess.Append(llm.Message{Role: llm.RoleAssistant, Content: "answer"}); err != nil {
		t.Fatal(err)
	}
	if _, err := a.CompactNow(context.Background()); err != nil {
		t.Fatal(err)
	}
	if a.checkpoint != "CHECKPOINT-V2" {
		t.Fatalf("second checkpoint = %q", a.checkpoint)
	}
	if a.checkpointCovered != 4 {
		t.Fatalf("covered = %d, want 4", a.checkpointCovered)
	}
	if len(inputs) != 2 {
		t.Fatalf("summary calls = %d", len(inputs))
	}
	second := inputs[1]
	if !strings.Contains(second, "Existing checkpoint") || !strings.Contains(second, "CHECKPOINT-V1") {
		t.Fatalf("update prompt did not carry the prior checkpoint: %q", second)
	}
	if !strings.Contains(second, "more context") {
		t.Fatalf("update prompt missing newer history: %q", second)
	}
	if strings.Contains(second, "first task") {
		t.Fatalf("update prompt re-summarized already-covered history: %q", second)
	}
	cp, ok := session.LoadCheckpoint(a.sessionID())
	if !ok || cp.Text != "CHECKPOINT-V2" || cp.Covered != 4 {
		t.Fatalf("sidecar = %+v ok=%v", cp, ok)
	}
}

// TestLoadCheckpointRestoresAndValidates：新回合能从 sidecar 直接接上
// checkpoint + 未覆盖消息；transcript 被改写后指纹失效，退回完整历史。
func TestLoadCheckpointRestoresAndValidates(t *testing.T) {
	prov := &scriptedProvider{fn: func(_ int, system string, _ []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		if strings.HasPrefix(system, "You are a context compaction") {
			return llm.Message{Role: llm.RoleAssistant, Content: "CHECKPOINT-X"}, nil
		}
		return llm.Message{Role: llm.RoleAssistant, Content: "ok"}, nil
	}}
	a, sess := newTestAgent(t)
	a.Provider = prov
	a.MinCompactTokens = 1
	a.Opt.Stderr = io.Discard
	for _, content := range []string{"u1", "a1", "u2", "a2"} {
		role := llm.RoleUser
		if strings.HasPrefix(content, "a") {
			role = llm.RoleAssistant
		}
		if err := sess.Append(llm.Message{Role: role, Content: content}); err != nil {
			t.Fatal(err)
		}
	}
	if _, err := a.CompactNow(context.Background()); err != nil {
		t.Fatal(err)
	}

	b := &Agent{Provider: prov, Tools: a.Tools, Session: sess, System: "test",
		Opt: Options{Stdout: io.Discard, Stderr: io.Discard}}
	if !b.LoadCheckpoint() {
		t.Fatal("valid checkpoint not restored")
	}
	if b.checkpoint != "CHECKPOINT-X" || b.compactFrom != 4 {
		t.Fatalf("restored state: checkpoint=%q from=%d", b.checkpoint, b.compactFrom)
	}
	if err := sess.Append(llm.Message{Role: llm.RoleUser, Content: "next"}); err != nil {
		t.Fatal(err)
	}
	b.ContextWindow = 1 << 30
	b.CompactAt = 1
	b.MinCompactTokens = 0
	view, err := b.contextView(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if len(view) != 2 || !strings.Contains(view[0].Content, "CHECKPOINT-X") || view[1].Content != "next" {
		t.Fatalf("restored view = %+v", view)
	}

	// 改写已覆盖的历史：指纹失效，不能继续用旧 checkpoint。
	sess.Messages[3].Content = "tampered"
	c := &Agent{Provider: prov, Tools: a.Tools, Session: sess, System: "test",
		Opt: Options{Stdout: io.Discard, Stderr: io.Discard}}
	if c.LoadCheckpoint() {
		t.Fatal("stale checkpoint must be rejected after transcript changes")
	}
}

// TestObservedTokensPrefersActualUsage：压缩判据用网关回报的真实 prompt tokens，
// 不再被 bytes/3 的估算单独带偏；请求之后新增的消息仍按估算补上。
func TestObservedTokensPrefersActualUsage(t *testing.T) {
	a := &Agent{
		System:           strings.Repeat("s", 300),
		ContextWindow:    10000,
		CompactAt:        0.5,
		MinCompactTokens: 1,
		lastInputTokens:  0,
	}
	view := []llm.Message{
		{Role: llm.RoleUser, Content: "short"},
		{Role: llm.RoleAssistant, Content: "ok"},
	}
	a.noteRequestUsage(&llm.Usage{InputTokens: 6000}, view)
	if !a.shouldCompact(view) {
		t.Fatal("actual 6000 tokens should trigger compaction despite a low estimate")
	}

	big := []llm.Message{
		{Role: llm.RoleUser, Content: strings.Repeat("a", 30000)},
		{Role: llm.RoleAssistant, Content: "ok"},
	}
	a.noteRequestUsage(&llm.Usage{InputTokens: 1000}, big)
	if a.shouldCompact(big) {
		t.Fatal("actual 1000 tokens should not be overridden by the estimate")
	}

	grown := append(append([]llm.Message{}, big...),
		llm.Message{Role: llm.RoleUser, Content: strings.Repeat("b", 15000)})
	if !a.shouldCompact(grown) {
		t.Fatal("messages added after the measured request must count toward the trigger")
	}
}

// TestForceCompactTargetHint：收缩阶梯要把目标 token 数写进摘要提示词。
func TestForceCompactTargetHint(t *testing.T) {
	var inputs []string
	prov := &scriptedProvider{fn: func(_ int, system string, msgs []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		if strings.HasPrefix(system, "You are a context compaction") {
			inputs = append(inputs, msgs[0].Content)
			return llm.Message{Role: llm.RoleAssistant, Content: "CP"}, nil
		}
		return llm.Message{Role: llm.RoleAssistant, Content: "ok"}, nil
	}}
	a, sess := newTestAgent(t)
	a.Provider = prov
	a.MinCompactTokens = 1
	a.Opt.Stderr = io.Discard
	for _, content := range []string{"u1", "a1"} {
		if err := sess.Append(llm.Message{Role: llm.RoleUser, Content: content}); err != nil {
			t.Fatal(err)
		}
	}
	view := append([]llm.Message(nil), sess.Messages...)
	a.noteRequestUsage(&llm.Usage{InputTokens: 10000}, view)
	if err := a.forceCompactForSize(context.Background(), view, 0.5); err != nil {
		t.Fatal(err)
	}
	if len(inputs) != 1 || !strings.Contains(inputs[0], "under about 5000 tokens") {
		t.Fatalf("target hint missing: %v", inputs)
	}
}

// TestOversizedShrinkLadder：请求过大被拒后按 0.7 / 0.5 / 0.35 逐级收缩，
// 每次重发都基于刚压出来的新视图；压到网关肯收为止。
func TestOversizedShrinkLadder(t *testing.T) {
	realCalls, summaryCalls := 0, 0
	var targets []int
	prov := &scriptedProvider{fn: func(_ int, system string, msgs []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		if strings.HasPrefix(system, "You are a context compaction") {
			summaryCalls++
			m := regexp.MustCompile(`under about (\d+) tokens`).FindStringSubmatch(msgs[0].Content)
			if m == nil {
				t.Fatalf("summary input missing target: %q", msgs[0].Content)
			}
			target := atoiTest(m[1])
			targets = append(targets, target)
			return llm.Message{Role: llm.RoleAssistant, Content: "CP:" + strings.Repeat("x", target*3)}, nil
		}
		realCalls++
		if realCalls <= 3 {
			return llm.Message{}, errors.New("HTTP 400: context length too large")
		}
		return llm.Message{Role: llm.RoleAssistant, Content: "done"}, nil
	}}
	a, sess := newTestAgent(t)
	a.Provider = prov
	a.Opt.Stderr = io.Discard
	a.ContextWindow = 1 << 30
	a.CompactAt = 1
	a.MinCompactTokens = 1
	sess.Messages = []llm.Message{
		{Role: llm.RoleUser, Content: strings.Repeat("a", 30000)},
		{Role: llm.RoleAssistant, Content: strings.Repeat("b", 30000)},
	}

	if err := a.Run(context.Background(), ""); err != nil {
		t.Fatalf("run: %v (real=%d summary=%d targets=%v)", err, realCalls, summaryCalls, targets)
	}
	if realCalls != 4 || summaryCalls != 3 {
		t.Fatalf("real=%d summary=%d targets=%v", realCalls, summaryCalls, targets)
	}
	if !(targets[0] > targets[1] && targets[1] > targets[2]) {
		t.Fatalf("targets must shrink: %v", targets)
	}
}

func atoiTest(s string) int {
	n := 0
	for _, r := range s {
		if r < '0' || r > '9' {
			break
		}
		n = n*10 + int(r-'0')
	}
	return n
}