package agent

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

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

// 中断时保留半截回复的测试。scriptedProvider 不发 delta（它直接返回整条消息），
// 所以这里要一个会"先流几个字再断"的 provider——真实事故就是这个形状。

type cuttingProvider struct {
	text      string // 断之前流出的正文
	reasoning string
	err       error
	calls     int
}

func (p *cuttingProvider) Stream(_ context.Context, _ string, _ []llm.Message, _ []llm.ToolDef,
	_ llm.StreamOptions, onDelta func(llm.StreamDelta)) (llm.Message, error) {
	p.calls++
	if p.reasoning != "" && onDelta != nil {
		onDelta(llm.StreamDelta{Reasoning: p.reasoning})
	}
	if p.text != "" && onDelta != nil {
		onDelta(llm.StreamDelta{Text: p.text})
	}
	// provider 出错时返回零值 Message：半截内容不在返回值里（这正是它丢字的根因）。
	return llm.Message{}, p.err
}

func cutAgent(t *testing.T, prov llm.Provider) (*Agent, *session.Session) {
	t.Helper()
	a, sess := newTestAgent(t)
	a.Provider = prov
	a.Tools = tools.NewRegistry()
	a.Opt.Stdout = io.Discard
	a.Opt.Stderr = io.Discard
	return a, sess
}

// TestInterruptKeepsStreamedText：屏幕上都打出字了，会话里却一个字都没有——
// 刷新/重进就消失，模型下一轮也不知道自己说到哪。
func TestInterruptKeepsStreamedText(t *testing.T) {
	noWait(t)
	prov := &cuttingProvider{text: "已经分析到第三步", err: errors.New("openai: 流被掐断")}
	a, sess := cutAgent(t, prov)
	err := a.Run(context.Background(), "分析这个仓库")
	if err == nil {
		t.Fatal("这一轮该以失败结束")
	}
	var as []llm.Message
	for _, m := range sess.Messages {
		if m.Role == llm.RoleAssistant {
			as = append(as, m)
		}
	}
	if len(as) != 1 {
		t.Fatalf("assistant 消息 %d 条，want 1：%+v", len(as), as)
	}
	if as[0].Content != "已经分析到第三步" {
		t.Errorf("半截内容没原样保留：%q", as[0].Content)
	}
	// 半截就是半截，得让消费方看得出来：webui 的重试判定靠它区分"这条没写完"和
	// "这一轮已经答完"，漏了标记就会把可续跑的回合判成"没有可重试的消息"。
	if !as[0].Interrupted {
		t.Error("没打 interrupted 标记")
	}
	// Step 必须盖上（否则 trace 里看不出这是第几步）；耗时用度不苛求正数——
	// 假 provider 是瞬时返回的，真跑起来才有毫秒。
	if as[0].Step != 1 || as[0].DurationMs < 0 || as[0].StartedAt == 0 {
		t.Errorf("trace 没盖上：%+v", as[0])
	}
	// 回合账仍然记的是失败/停止，不是"说完了"。
	if got := sess.TurnEnds[len(sess.TurnEnds)-1].Reason; got == "completed" {
		t.Error("失败被记成了 completed")
	}
}

// TestInterruptKeepsReasoning：思考文本也一样保留（前端展开的"思考过程"不该空掉）。
func TestInterruptKeepsReasoning(t *testing.T) {
	noWait(t)
	prov := &cuttingProvider{reasoning: "先看看入口在哪", err: errors.New("stream stalled")}
	a, sess := cutAgent(t, prov)
	_ = a.Run(context.Background(), "讲")
	if len(sess.Messages) < 2 || sess.Messages[len(sess.Messages)-1].Reasoning != "先看看入口在哪" {
		t.Fatalf("思考文本没保留：%+v", sess.Messages)
	}
}

// TestInterruptNothingStreamed：连接一开就断时**不能**留一条空 assistant 消息——
// 它会被"没有 tool_calls ⇒ 说完了"读成一次正常的空回答，正是我们在治的那个病。
func TestInterruptNothingStreamed(t *testing.T) {
	noWait(t)
	prov := &cuttingProvider{err: errors.New("connection reset by peer")}
	a, sess := cutAgent(t, prov)
	if err := a.Run(context.Background(), "讲"); err == nil {
		t.Fatal("该失败")
	}
	for _, m := range sess.Messages {
		if m.Role == llm.RoleAssistant {
			t.Fatalf("什么都没流出来，却落了一条 assistant：%+v", m)
		}
	}
}

// TestInterruptNeverPersistsPartialToolCalls：半截的工具参数是没写完的 JSON。
// 落了就会有一条永远配不上对的 tool_use，Anthropic 会拒绝整条会话（见
// synthMissingToolResults 的注释）。这里断言它压根不会被写进会话。
func TestInterruptNeverPersistsPartialToolCalls(t *testing.T) {
	noWait(t)
	prov := &cuttingProvider{text: `我想调用 write，参数 {"path":"a.go",`, err: errors.New("aborted")}
	a, sess := cutAgent(t, prov)
	_ = a.Run(context.Background(), "写个文件")
	last := sess.Messages[len(sess.Messages)-1]
	if len(last.ToolCalls) != 0 {
		t.Errorf("半截 tool_calls 被落盘了：%+v", last.ToolCalls)
	}
	if !strings.Contains(last.Content, "write") {
		t.Error("正文没按原样保留")
	}
}

// TestCancelKeepsPartialIsIdempotentAcrossRetries：空回复重试环里每次 call() 都会
// 清空累积器，落盘的只能是**最后一次**尝试流出的内容，不能把两轮的字拼在一起。
func TestCancelKeepsPartialIsIdempotentAcrossRetries(t *testing.T) {
	noWait(t)
	prov := &retryThenCutProvider{}
	a, sess := cutAgent(t, prov)
	_ = a.Run(context.Background(), "讲")
	var texts []string
	for _, m := range sess.Messages {
		if m.Role == llm.RoleAssistant {
			texts = append(texts, m.Content)
		}
	}
	if len(texts) != 1 || texts[0] != "第二次才流出的字" {
		t.Errorf("落盘内容不对（重试拼接或丢失）：%+v", texts)
	}
}

// TestAbandonStreamFiresBeforeRetry：被丢弃的那次尝试流出的字，宿主需要知道——
// 它已经打到屏幕上了，重跑一次就是两遍同样的话。通知只给**正文**的字节数
// （思考走另一条通道，服务端没有它的字节账）。
func TestAbandonStreamFiresBeforeRetry(t *testing.T) {
	noWait(t)
	type abandon struct {
		reason string
		bytes  int
	}
	var got []abandon
	prov := &cuttingProvider{text: "第一遍说出的话", err: llm.ErrTruncatedStream}
	a, sess := cutAgent(t, prov)
	a.Opt.OnAbandon = func(reason string, n int) { got = append(got, abandon{reason, n}) }
	if err := a.Run(context.Background(), "讲"); err == nil {
		t.Fatal("这一轮该以失败结束")
	}
	// 预算现在按时间算（不再是一个固定次数），所以这里只钉"每次重试都作废一次，
	// 而且不止两次"——旧的 2 次上限正是"网关抖一下就中断整轮"的根源。
	if len(got) < 3 {
		t.Fatalf("作废通知只有 %d 次，重试预算没有生效：%+v", len(got), got)
	}
	for i, ab := range got {
		if ab.bytes != len("第一遍说出的话") {
			t.Errorf("第 %d 次通知带的字节数 = %d, want %d", i+1, ab.bytes, len("第一遍说出的话"))
		}
		if ab.reason == "" {
			t.Errorf("第 %d 次通知没带原因", i+1)
		}
	}
	// 会话里始终只有一份：作废只影响屏幕，落盘走的是最后一次尝试攒下的内容。
	var kept int
	for _, m := range sess.Messages {
		if m.Role == llm.RoleAssistant {
			kept++
			if strings.Count(m.Content, "第一遍说出的话") != 1 {
				t.Errorf("半截内容被拼成了多份：%q", m.Content)
			}
		}
	}
	if kept != 1 {
		t.Fatalf("assistant 消息 %d 条，want 1", kept)
	}
}

// TestAbandonStreamSilentWhenNothingStreamed：什么都没流出来时不发空通知——
// 宿主白擦一次屏幕，什么都没换回来。
func TestAbandonStreamSilentWhenNothingStreamed(t *testing.T) {
	noWait(t)
	called := 0
	a, _ := cutAgent(t, &cuttingProvider{err: errors.New("connection reset by peer")})
	a.Opt.OnAbandon = func(string, int) { called++ }
	_ = a.Run(context.Background(), "讲")
	if called != 0 {
		t.Fatalf("什么都没流出，却通知了 %d 次", called)
	}
}

// retryThenCutProvider：第一次报空补全（触发重试环），第二次流出几个字后断掉。
type retryThenCutProvider struct{ n int }

func (p *retryThenCutProvider) Stream(_ context.Context, _ string, _ []llm.Message, _ []llm.ToolDef,
	_ llm.StreamOptions, onDelta func(llm.StreamDelta)) (llm.Message, error) {
	p.n++
	if p.n == 1 {
		return llm.Message{}, llm.ErrEmptyCompletion
	}
	if onDelta != nil {
		onDelta(llm.StreamDelta{Text: "第二次才流出的字"})
	}
	return llm.Message{}, errors.New("stream cut")
}