package agent

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

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

// 输出上限续写的行为测试。用的是 scriptedProvider：它能看到每次请求的 system 与
// 历史，所以"prefill 形状对不对"这种关键断言不用真起网关。

func continueAgent(t *testing.T, prov llm.Provider) *Agent {
	t.Helper()
	a := &Agent{
		Provider: prov,
		Tools:    tools.NewRegistry(),
		Session:  &session.Session{Path: "/tmp/none.jsonl"},
		System:   "sys",
		Opt:      Options{Stdout: io.Discard, Stderr: io.Discard},
	}
	return a
}

// setContinueBudget 把续写预算压到 n（用例不该白跑 3 次请求）。
func setContinueBudget(t *testing.T, n int) {
	t.Helper()
	old := outputContinueMax
	outputContinueMax = n
	t.Cleanup(func() { outputContinueMax = old })
}

func assistants(msgs []llm.Message) []llm.Message {
	var out []llm.Message
	for _, m := range msgs {
		if m.Role == llm.RoleAssistant {
			out = append(out, m)
		}
	}
	return out
}

// TestContinueOutputMergesIntoOneMessage 是这条特性的全部意义：半截答案被接上，
// 而且**会话里只有一条 assistant**（两条相邻 assistant 会让回合号与"哪条是最终
// 回复"都多一套特例）。
func TestContinueOutputMergesIntoOneMessage(t *testing.T) {
	setContinueBudget(t, 3)
	var systems []string
	var lastRoles []string
	prov := &scriptedProvider{fn: func(call int, system string, msgs []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		systems = append(systems, system)
		if len(msgs) > 0 {
			lastRoles = append(lastRoles, string(msgs[len(msgs)-1].Role))
		}
		switch call {
		case 1:
			return llm.Message{Role: llm.RoleAssistant, Content: "第一段没写完",
				OutputLimited: true, Usage: &llm.Usage{InputTokens: 10, OutputTokens: 99}}, nil
		default:
			return llm.Message{Role: llm.RoleAssistant, Content: " ——第二段补上了",
				Usage: &llm.Usage{InputTokens: 12, OutputTokens: 7}}, nil
		}
	}}
	a := continueAgent(t, prov)
	if err := a.Run(context.Background(), "讲讲这件事"); err != nil {
		t.Fatalf("续写成功时这一轮该正常结束：%v", err)
	}
	as := assistants(a.Session.Messages)
	if len(as) != 1 {
		t.Fatalf("assistant 消息 %d 条，want 1（必须合并成一条）：%+v", len(as), as)
	}
	if got := as[0].Content; got != "第一段没写完 ——第二段补上了" {
		t.Errorf("合并后的正文 = %q", got)
	}
	if as[0].OutputLimited {
		t.Error("已经写完了，却还标着被切断")
	}
	if as[0].Usage == nil || as[0].Usage.OutputTokens != 106 || as[0].Usage.InputTokens != 22 {
		t.Errorf("usage 没按两次请求累加：%+v", as[0].Usage)
	}
	if a.outputContinuations() != 1 {
		t.Errorf("续写次数记账 = %d，want 1", a.outputContinuations())
	}
	// 续写请求的形状：历史末尾是 assistant（prefill），system 带续写指令。
	if len(lastRoles) != 2 || lastRoles[1] != string(llm.RoleAssistant) {
		t.Errorf("续写请求的历史末尾是 %v，want assistant", lastRoles)
	}
	if !strings.Contains(systems[1], "输出") || strings.Contains(systems[0], "输出被长度") {
		t.Errorf("续写指令只该出现在续写那一轮：sys0=%q sys1 含指令=%v",
			systems[0][:min(20, len(systems[0]))], strings.Contains(systems[1], "输出"))
	}
	// 回合数没有因为续写而多算一步：curStep 只在 step() 里自增。
	if a.curStep != 1 {
		t.Errorf("curStep = %d，want 1（续写不算新的一步）", a.curStep)
	}
}

// TestContinueOutputBudgetExhausted：3 次仍被切断说明是模型/网关在循环，
// 保留已有内容并如实标注——用户得知道为什么停，模型下一轮也得知道接哪。
func TestContinueOutputBudgetExhausted(t *testing.T) {
	setContinueBudget(t, 2)
	prov := &scriptedProvider{fn: func(_ int, _ string, _ []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		return llm.Message{Role: llm.RoleAssistant, Content: "还是没写完", OutputLimited: true}, nil
	}}
	a := continueAgent(t, prov)
	if err := a.Run(context.Background(), "长文"); err != nil {
		t.Fatal(err)
	}
	if prov.calls != 3 {
		t.Errorf("模型被调了 %d 次，want 3（首轮 + 预算 2 次续写）", prov.calls)
	}
	as := assistants(a.Session.Messages)
	if len(as) != 1 {
		t.Fatalf("assistant %d 条，want 1", len(as))
	}
	if !strings.HasPrefix(as[0].Content, "⏸ 回复被输出长度上限切断，已自动续写 2 次") {
		t.Errorf("耗尽时没有如实标注：%q", as[0].Content)
	}
	if !as[0].OutputLimited {
		t.Error("确实没写完，OutputLimited 却被清掉了")
	}
	if strings.Count(as[0].Content, "还是没写完") != 3 {
		t.Errorf("三段输出没都留下：%q", as[0].Content)
	}
}

// TestContinueOutputFailureKeepsPartial：续写请求本身失败（网络/取消/坏流）时，
// 不能把这一轮弄成错误——手里这半截答案比一个 error 有用。
func TestContinueOutputFailureKeepsPartial(t *testing.T) {
	setContinueBudget(t, 3)
	prov := &scriptedProvider{fn: func(call int, _ string, _ []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		if call == 1 {
			return llm.Message{Role: llm.RoleAssistant, Content: "半截但有用", OutputLimited: true}, nil
		}
		return llm.Message{}, fmt.Errorf("openai: HTTP 500")
	}}
	a := continueAgent(t, prov)
	if err := a.Run(context.Background(), "讲"); err != nil {
		t.Fatalf("续写失败不该弄坏这一轮：%v", err)
	}
	as := assistants(a.Session.Messages)
	if len(as) != 1 || !strings.Contains(as[0].Content, "半截但有用") {
		t.Fatalf("半截答案丢了：%+v", as)
	}
	if !strings.Contains(as[0].Content, "回复「继续」") {
		t.Errorf("没告诉用户下一步该说什么：%q", as[0].Content)
	}
}

// TestContinueOutputSkippedWithThinkingSignature：带 thinking 签名时不续写。
// Anthropic 不接受"思考中"的 assistant prefill，而把两段的 thinking 块与签名拼起来
// 会让下一个请求被拒（签名必须与块一一对应）——宁可不续，也不能留下一条坏会话。
func TestContinueOutputSkippedWithThinkingSignature(t *testing.T) {
	setContinueBudget(t, 3)
	prov := &scriptedProvider{fn: func(_ int, _ string, _ []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		return llm.Message{Role: llm.RoleAssistant, Content: "思考着被切断",
			OutputLimited: true, ReasoningSignature: "sig"}, nil
	}}
	a := continueAgent(t, prov)
	if err := a.Run(context.Background(), "讲"); err != nil {
		t.Fatal(err)
	}
	if prov.calls != 1 {
		t.Errorf("带签名还去续写了：调用 %d 次，want 1", prov.calls)
	}
	as := assistants(a.Session.Messages)
	if !strings.Contains(as[0].Content, "回复被输出长度上限切断") {
		t.Errorf("没标注：%q", as[0].Content)
	}
	if strings.Contains(as[0].Content, "已自动续写") {
		t.Errorf("一次都没续，却声称续写过：%q", as[0].Content)
	}
}

// TestNoContinuationWhenNotCut 是反向护栏：正常的最终回复绝不能多打一次请求
// （那等于给每一次对话凭空加一轮往返）。
func TestNoContinuationWhenNotCut(t *testing.T) {
	setContinueBudget(t, 3)
	prov := &scriptedProvider{fn: func(_ int, _ string, _ []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		return llm.Message{Role: llm.RoleAssistant, Content: "一次说完"}, nil
	}}
	a := continueAgent(t, prov)
	if err := a.Run(context.Background(), "讲"); err != nil {
		t.Fatal(err)
	}
	if prov.calls != 1 {
		t.Errorf("调用 %d 次，want 1", prov.calls)
	}
	if a.outputContinuations() != 0 {
		t.Error("没切断却记了续写次数")
	}
	if got := assistants(a.Session.Messages)[0].Content; got != "一次说完" {
		t.Errorf("正文被改动了：%q", got)
	}
}

// TestContinuationCounterResetsPerTurn：预算是"每一轮"的，不是"每个进程"的——
// 上一条消息续写过 3 次，不该让下一条消息一次都续不了。
func TestContinuationCounterResetsPerTurn(t *testing.T) {
	setContinueBudget(t, 1)
	cut := true
	prov := &scriptedProvider{fn: func(_ int, _ string, _ []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		if cut {
			return llm.Message{Role: llm.RoleAssistant, Content: "半", OutputLimited: true}, nil
		}
		return llm.Message{Role: llm.RoleAssistant, Content: "整"}, nil
	}}
	a := continueAgent(t, prov)
	if err := a.Run(context.Background(), "第一轮"); err != nil {
		t.Fatal(err)
	}
	if a.outputContinuations() != 1 {
		t.Fatalf("第一轮续写次数 = %d，want 1", a.outputContinuations())
	}
	cut = false
	if err := a.Run(context.Background(), "第二轮"); err != nil {
		t.Fatal(err)
	}
	if a.outputContinuations() != 0 {
		t.Errorf("第二轮继承了上一轮的续写计数：%d", a.outputContinuations())
	}
}