package agent

import (
	"context"
	"errors"
	"fmt"
	"io"
	"strings"
	"sync"
	"testing"
	"unicode/utf8"

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

// scriptedProvider 按调用次序执行脚本式响应，供状态机/压缩/子代理测试复用。
type scriptedProvider struct {
	mu    sync.Mutex
	calls int
	opts  []llm.StreamOptions
	fn    func(call int, system string, msgs []llm.Message, tools []llm.ToolDef) (llm.Message, error)
}

func (p *scriptedProvider) Stream(_ context.Context, system string, msgs []llm.Message, tools []llm.ToolDef, opts llm.StreamOptions, _ func(llm.StreamDelta)) (llm.Message, error) {
	p.mu.Lock()
	p.calls++
	p.opts = append(p.opts, opts)
	c := p.calls
	p.mu.Unlock()
	// 真实 provider 会把动态 system 单独下发；测试回调仍看拼接后的全文，
	// 与改造前 sysNow 的语义一致。
	return p.fn(c, system+opts.SystemDynamic, msgs, tools)
}

// lastOpts 返回最近一次请求的选项（收尾轮靠 DisableTools 识别，而不是靠"tools 为空"）。
func (p *scriptedProvider) lastOpts() llm.StreamOptions {
	p.mu.Lock()
	defer p.mu.Unlock()
	if len(p.opts) == 0 {
		return llm.StreamOptions{}
	}
	return p.opts[len(p.opts)-1]
}

func toolNames(defs []llm.ToolDef) []string {
	names := make([]string, len(defs))
	for i, d := range defs {
		names[i] = d.Name
	}
	return names
}

// TestAgentThinkingPropagation：Agent.Thinking 应通过 StreamOptions 传给 Provider。
func TestAgentThinkingPropagation(t *testing.T) {
	prov := &scriptedProvider{fn: func(_ int, _ string, _ []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		return llm.Message{Role: llm.RoleAssistant, Content: "done"}, nil
	}}
	a := &Agent{
		Provider: prov,
		Tools:    tools.NewRegistry(),
		Session:  &session.Session{},
		System:   "sys",
		Thinking: llm.ThinkingHigh,
		Opt:      Options{Stdout: io.Discard, Stderr: io.Discard},
	}
	if err := a.Run(context.Background(), "hi"); err != nil {
		t.Fatal(err)
	}
	if len(prov.opts) == 0 || prov.opts[0].Thinking != llm.ThinkingHigh {
		t.Fatalf("provider did not receive thinking level: %+v", prov.opts)
	}
}

func hasName(defs []llm.ToolDef, name string) bool {
	for _, d := range defs {
		if d.Name == name {
			return true
		}
	}
	return false
}

// TestEstimateTokens 验证估算随内容增长且对 CJK/ASCII 都有覆盖。
func TestEstimateTokens(t *testing.T) {
	short := estimateTokens([]llm.Message{{Role: llm.RoleUser, Content: "hi"}}, "")
	long := estimateTokens([]llm.Message{{Role: llm.RoleUser, Content: strings.Repeat("a", 3000)}}, "")
	if short <= 0 || long <= short {
		t.Fatalf("estimate not monotonic: short=%d long=%d", short, long)
	}
	cjk := estimateTokens([]llm.Message{{Role: llm.RoleUser, Content: strings.Repeat("好", 99)}}, "")
	if cjk < 30 || cjk > 300 {
		t.Fatalf("cjk estimate out of plausible range: %d", cjk)
	}
}

// TestTruncateToolResults：超长 tool result 在视图层被截断，transcript 不动。
func TestTruncateToolResults(t *testing.T) {
	big := llm.Message{Role: llm.RoleTool, ToolCallID: "x", Content: strings.Repeat("b", 50000)}
	msgs := []llm.Message{
		{Role: llm.RoleUser, Content: "go"},
		{Role: llm.RoleAssistant, Content: "reading"},
		big,
	}
	view := truncateToolResults(msgs, 20000)
	if &view[2] == &msgs[2] {
		t.Fatal("expected a copy, got the same slice")
	}
	if len(view[2].Content) >= 50000 {
		t.Fatal("oversized tool result was not truncated")
	}
	if !strings.Contains(view[2].Content, "[truncated") {
		t.Fatalf("truncation notice missing: %q", view[2].Content[len(view[2].Content)-120:])
	}
	// transcript 不被修改
	if len(msgs[2].Content) != 50000 {
		t.Fatal("original transcript message was mutated")
	}
	// 全部消息都低于阈值时返回原切片（零拷贝路径）
	small := []llm.Message{
		{Role: llm.RoleUser, Content: "go"},
		{Role: llm.RoleTool, ToolCallID: "x", Content: "short result"},
	}
	same := truncateToolResults(small, 20000)
	if &same[0] != &small[0] {
		t.Fatal("expected original slice when nothing to truncate")
	}
}

// TestTruncateToolResultsHeadAndTail：超长结果保留首尾，尾部（报错/汇总所在处）不丢。
func TestTruncateToolResultsHeadAndTail(t *testing.T) {
	content := "HEAD" + strings.Repeat("m", 50000) + "TAIL"
	msgs := []llm.Message{
		{Role: llm.RoleUser, Content: "go"},
		{Role: llm.RoleTool, ToolCallID: "x", Content: content},
	}
	view := truncateToolResults(msgs, 20000)
	got := view[1].Content
	if !strings.HasPrefix(got, "HEAD") {
		t.Fatal("head marker lost")
	}
	if !strings.HasSuffix(got, "TAIL") {
		t.Fatal("tail marker lost (head-only truncation?)")
	}
	if !strings.Contains(got, "[truncated") {
		t.Fatal("truncation notice missing")
	}
	if len(got) > 20000 {
		t.Fatalf("truncated result exceeds max: %d", len(got))
	}
	// 多字节内容截断后仍须是合法 UTF-8（不切出半个字符）。
	cjk := strings.Repeat("中", 30000)
	v := truncateToolResults([]llm.Message{{Role: llm.RoleTool, Content: cjk}}, 1000)
	if !utf8.ValidString(v[0].Content) {
		t.Fatal("invalid utf-8 after truncation")
	}
	if !strings.Contains(v[0].Content, "[truncated") {
		t.Fatal("notice missing for cjk result")
	}
}

// TestAutocompactFlow：超过阈值触发摘要压缩；压缩后的工作集能接上新消息；
// transcript 保持完整。
func TestAutocompactFlow(t *testing.T) {
	var compactedSeen, secondTurnOK bool
	prov := &scriptedProvider{fn: func(call int, system string, msgs []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		switch call {
		case 1:
			if !strings.HasPrefix(system, "You are a context compaction") {
				t.Errorf("call 1: expected compaction system prompt, got %q", system)
			}
			return llm.Message{Role: llm.RoleAssistant, Content: "SUMMARY: user asked X"}, nil
		case 2:
			if len(msgs) != 1 || !strings.Contains(msgs[0].Content, "SUMMARY") {
				t.Errorf("call 2: expected compacted single-message history, got %d msgs", len(msgs))
			} else {
				compactedSeen = true
			}
			return llm.Message{Role: llm.RoleAssistant, Content: "done"}, nil
		case 3:
			// 工作集 = [summary, assistant done, user more]
			if len(msgs) == 3 && strings.Contains(msgs[0].Content, "SUMMARY") && msgs[2].Content == "more" {
				secondTurnOK = true
			} else {
				t.Errorf("call 3: expected [summary, done, new user msg], got %d msgs (last=%q)", len(msgs), msgs[len(msgs)-1].Content)
			}
			return llm.Message{Role: llm.RoleAssistant, Content: "done2"}, nil
		}
		t.Errorf("unexpected call %d", call)
		return llm.Message{}, nil
	}}

	a, sess := newTestAgent(t)
	a.Provider = prov
	// 窗口 1000、阈值 0.5 → 500 tokens：预载 2000 字符（≈682 tokens）触发
	// 压缩；压缩后摘要（≈50 tokens）不会立即再触发。
	a.ContextWindow = 1000
	a.CompactAt = 0.5
	a.MinCompactTokens = 1
	sess.Messages = []llm.Message{
		{Role: llm.RoleUser, Content: strings.Repeat("a", 2000)},
		{Role: llm.RoleAssistant, Content: "partial answer"},
	}

	if err := a.Run(context.Background(), ""); err != nil {
		t.Fatal(err)
	}
	if !compactedSeen {
		t.Fatal("compacted history never reached the model")
	}
	// transcript 未被压缩动过
	if len(sess.Messages) < 3 {
		t.Fatalf("transcript lost messages: %d", len(sess.Messages))
	}
	// 压缩后新消息接在工作集尾部
	if err := a.Run(context.Background(), "more"); err != nil {
		t.Fatal(err)
	}
	if !secondTurnOK {
		t.Fatal("post-compaction new message was not appended to the working set")
	}
}

// TestAutocompactFailureDegradation：摘要调用失败时不致命，退化为不压缩继续。
func TestAutocompactFailureDegradation(t *testing.T) {
	prov := &scriptedProvider{fn: func(call int, system string, _ []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		if strings.HasPrefix(system, "You are a context compaction") {
			return llm.Message{}, context.DeadlineExceeded
		}
		return llm.Message{Role: llm.RoleAssistant, Content: "ok"}, nil
	}}
	a, sess := newTestAgent(t)
	a.Provider = prov
	a.ContextWindow = 100
	a.CompactAt = 0.5
	a.MinCompactTokens = 1
	sess.Messages = []llm.Message{
		{Role: llm.RoleUser, Content: strings.Repeat("a", 1000)},
		{Role: llm.RoleAssistant, Content: "partial"},
	}
	if err := a.Run(context.Background(), ""); err != nil {
		t.Fatalf("run should survive compaction failure: %v", err)
	}
}

// TestSubAgentTool：父 agent 委派子 agent，子 agent 工具表无 agent、
// 结果回填主会话、ephemeral 会话不落盘。
func TestSubAgentTool(t *testing.T) {
	var nestedToolNames []string
	prov := &scriptedProvider{fn: func(call int, _ string, msgs []llm.Message, tools []llm.ToolDef) (llm.Message, error) {
		switch call {
		case 1: // 父 turn 1：发起 agent 工具调用
			if !hasName(tools, "agent") {
				t.Error("call 1: parent toolset should include agent")
			}
			return llm.Message{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{
				{ID: "s1", Name: "agent", Input: `{"prompt":"find the banana"}`},
			}}, nil
		case 2: // 子 agent 唯一一轮：给出最终回复
			nestedToolNames = toolNames(tools)
			if hasName(tools, "agent") {
				t.Error("call 2: nested toolset must not include agent (recursion)")
			}
			return llm.Message{Role: llm.RoleAssistant, Content: "SUB-RESULT banana found"}, nil
		case 3: // 父 turn 2：最终回复
			last := msgs[len(msgs)-1]
			if last.Role != llm.RoleTool || !strings.Contains(last.Content, "SUB-RESULT") {
				t.Errorf("call 3: expected paired tool result, got role=%s content=%q", last.Role, last.Content)
			}
			return llm.Message{Role: llm.RoleAssistant, Content: "done"}, nil
		}
		t.Errorf("unexpected call %d", call)
		return llm.Message{}, nil
	}}

	a, sess := newTestAgent(t)
	a.Provider = prov
	a.MaxDepth = 1 // 允许派发 1 层：顶层可派发、子 agent 不能再派发（历史行为）
	if err := a.Tools.Register(a.SubAgentTool()); err != nil {
		t.Fatal(err)
	}
	if err := a.Run(context.Background(), "find the banana"); err != nil {
		t.Fatal(err)
	}
	if len(nestedToolNames) == 0 {
		t.Fatal("nested agent never ran")
	}
	// ephemeral：主会话 transcript 里不应出现子 agent 的内部消息，
	// 只应有父 user / 父 assistant(toolcall) / tool result / 父最终回复 4 条
	if len(sess.Messages) != 4 {
		t.Fatalf("main transcript polluted: %d messages", len(sess.Messages))
	}
}

// 压缩切点安全不变量：切点永不落在工具回合中间（assistant-with-tool_calls
// 与其 tool results 必须同侧）。
func TestCompactCutSafeBoundary(t *testing.T) {
	mk := func(role llm.Role, calls int) llm.Message {
		m := llm.Message{Role: role, Content: "x"}
		for i := 0; i < calls; i++ {
			m.ToolCalls = append(m.ToolCalls, llm.ToolCall{ID: fmt.Sprintf("c%d", i), Name: "t"})
		}
		return m
	}
	// 20 条：user/assistant 交替夹杂工具回合。
	var msgs []llm.Message
	for i := 0; i < 10; i++ {
		msgs = append(msgs,
			mk(llm.RoleUser, 0),
			mk(llm.RoleAssistant, 2), // 工具回合开始
			mk(llm.RoleTool, 0),
			mk(llm.RoleTool, 0),
			mk(llm.RoleAssistant, 0), // 工具回合结束（最终回复）
		)
	}

	for keep := 1; keep <= 25; keep += 3 {
		cut := compactCut(msgs, keep)
		if cut < 0 || cut > len(msgs) {
			t.Fatalf("keep=%d cut=%d out of range", keep, cut)
		}
		if cut > 0 {
			// 切点前一条若是 assistant，不能带未完成 tool_calls；
			// 切点这条不能是孤儿 tool result。
			if m := msgs[cut]; m.Role == llm.RoleTool {
				t.Fatalf("keep=%d cut=%d lands on orphan tool result", keep, cut)
			}
			if m := msgs[cut-1]; m.Role == llm.RoleAssistant && len(m.ToolCalls) > 0 {
				t.Fatalf("keep=%d cut=%d splits an open tool turn", keep, cut)
			}
		}
	}
}

// TestCompactNowManual：/compact 手动压缩——无视阈值触发摘要、更新工作集缓存、
// transcript 不动；下一次 contextView 基于压缩缓存 + 新消息。
func TestCompactNowManual(t *testing.T) {
	prov := &scriptedProvider{fn: func(call int, system string, msgs []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		if !strings.HasPrefix(system, "You are a context compaction") {
			t.Errorf("expected compaction system prompt, got %q", system)
		}
		return llm.Message{Role: llm.RoleAssistant, Content: "SUMMARY: user asked X"}, nil
	}}
	a, sess := newTestAgent(t)
	a.Provider = prov
	a.MinCompactTokens = 1
	sess.Messages = []llm.Message{
		{Role: llm.RoleUser, Content: strings.Repeat("a", 2000)},
		{Role: llm.RoleAssistant, Content: "partial answer"},
	}

	res, err := a.CompactNow(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	// 少于 compactKeepRecent 条：全量摘要（2 条都进摘要、0 条保留原文）
	if res.Summarized != 2 || res.Kept != 0 {
		t.Fatalf("split wrong: summarized=%d kept=%d", res.Summarized, res.Kept)
	}
	if res.TokensAfter >= res.TokensBefore {
		t.Fatalf("compaction did not shrink view: %d → %d", res.TokensBefore, res.TokensAfter)
	}
	// 工作集缓存已更新，transcript 不动
	if a.compacted == nil || a.compactFrom != len(sess.Messages) {
		t.Fatalf("compaction cache not updated: compacted=%v compactFrom=%d", a.compacted != nil, a.compactFrom)
	}
	if len(sess.Messages) != 2 {
		t.Fatalf("transcript mutated: %d messages", len(sess.Messages))
	}
	// 重复手动压缩：刚压缩过、无新消息时 autocompact 仍会执行（手动即强制），
	// 但 shouldCompact 的防抖不适用——这里只验证不 panic、结果一致。
	if _, err := a.CompactNow(context.Background()); err != nil {
		t.Fatalf("second manual compact failed: %v", err)
	}
}

// TestCompactNowTooSmall：低于 MinCompactTokens 返回 ErrContextTooSmall，
// 不发起摘要调用。
func TestCompactNowTooSmall(t *testing.T) {
	prov := &scriptedProvider{fn: func(call int, _ string, _ []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		t.Errorf("summary call should not happen, got call %d", call)
		return llm.Message{}, nil
	}}
	a, sess := newTestAgent(t)
	a.Provider = prov
	a.MinCompactTokens = 1 << 20
	sess.Messages = []llm.Message{
		{Role: llm.RoleUser, Content: "hi"},
		{Role: llm.RoleAssistant, Content: "hello"},
	}

	res, err := a.CompactNow(context.Background())
	if !errors.Is(err, ErrContextTooSmall) {
		t.Fatalf("want ErrContextTooSmall, got %v", err)
	}
	if res.TokensBefore <= 0 {
		t.Fatalf("TokensBefore not reported: %d", res.TokensBefore)
	}
	if a.compacted != nil {
		t.Fatal("compaction cache must stay empty on refusal")
	}
}