package orchestration

import (
	"context"
	"errors"
	"strings"
	"sync"
	"testing"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/tool"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
)

// requestCaptureAdapter 记录每次模型请求的完整报文（Messages/Tools），
// 使测试能断言「子代理的模型真正看到了什么」——这是 02.4.2 的裁决
// （任务作为一条 user 消息注入）与 02.4.1（工具白名单）的真断言载体：
// 恒答固定文本的 mock 适配器测不出这两类缺陷。
type requestCaptureAdapter struct {
	final string

	mu       sync.Mutex
	requests []model.Request
}

func (a *requestCaptureAdapter) Name() string { return "capture" }

func (a *requestCaptureAdapter) record(req model.Request) {
	a.mu.Lock()
	a.requests = append(a.requests, req)
	a.mu.Unlock()
}

func (a *requestCaptureAdapter) Chat(_ context.Context, req model.Request, _ model.Options) (model.Response, error) {
	a.record(req)
	return model.Response{Model: "mock-1", Content: a.final, StopReason: "end_turn"}, nil
}

func (a *requestCaptureAdapter) ChatStream(_ context.Context, req model.Request, _ func(model.StreamEvent)) (model.Response, error) {
	return a.Chat(context.Background(), req, model.Options{})
}

// request 返回第 i 次捕获的模型请求。
func (a *requestCaptureAdapter) request(i int) model.Request {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.requests[i]
}

func (a *requestCaptureAdapter) count() int {
	a.mu.Lock()
	defer a.mu.Unlock()
	return len(a.requests)
}

// captureLoop 在真实 loop.Loop 之上构建带请求捕获的调度器（生产形状，
// 只是模型换成了捕获适配器）。tools 为 nil 时不接 dispatcher。
func captureLoop(t *testing.T, ad model.Adapter, d *tool.Dispatcher) *Scheduler {
	t.Helper()
	ag, err := agent.New("你是后台子代理，按角色与安全约束行事", d, nil,
		agent.WithModel(ad, nil),
		agent.WithOptions(agent.Options{Model: "mock-1"}))
	if err != nil {
		t.Fatalf("agent.New: %v", err)
	}
	lp := loop.New(loop.Options{Agent: ag, Tools: d, MaxRounds: 8})
	return newScheduler(t, Options{Loop: lp, MaxParallel: 1})
}

// messageTextOf 提取一条消息的纯文本（拼接 text 与 tool_result 分段，
// 后者承载被否决/失败调用的回灌文本）。
func messageTextOf(m model.Message) string {
	var b strings.Builder
	for _, c := range m.Content {
		if c.Type == model.ContentTypeText || c.Type == model.ContentTypeToolResult {
			b.WriteString(c.Text)
		}
	}
	return b.String()
}

// TestSpawnToolInjectsTaskAsUserMessage 是 02.4.2（P0「子代理收不到任务」）
// 的裁决断言：经生产派发工具创建的子代理，其模型请求里必须有一条 user
// 消息同时承载 goal 正文与 attach_info 独立段落；system 只放角色与安全
// 约束，绝不包含任务文本。
func TestSpawnToolInjectsTaskAsUserMessage(t *testing.T) {
	t.Parallel()
	ad := &requestCaptureAdapter{final: "子代理完成"}
	s := captureLoop(t, ad, nil)
	sp := NewSpawnTool(s, WithSpawnWait(3*time.Second))

	raw, err := sp.Execute(context.Background(), map[string]any{
		"goal":        "调研量子纠错现状",
		"attach_info": "只看 2026 年后的论文",
	})
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	got := decodeSpawnResult(t, raw)
	if got.Status != "done" {
		t.Fatalf("Status = %q, want done", got.Status)
	}
	if ad.count() == 0 {
		t.Fatal("the sub-agent never reached a model; nothing to assert the task injection on")
	}
	msgs := ad.request(0).Messages

	var userTexts, systemTexts []string
	for _, m := range msgs {
		switch m.Role {
		case model.RoleUser:
			userTexts = append(userTexts, messageTextOf(m))
		case model.RoleSystem:
			systemTexts = append(systemTexts, messageTextOf(m))
		}
	}
	if len(userTexts) != 1 {
		t.Fatalf("user messages = %d (%v), want exactly one carrying the task", len(userTexts), userTexts)
	}
	task := userTexts[0]
	for _, want := range []string{"[任务目标]", "调研量子纠错现状", "[附加信息]", "只看 2026 年后的论文"} {
		if !strings.Contains(task, want) {
			t.Fatalf("task user message = %q, want it to contain %q", task, want)
		}
	}
	for _, sys := range systemTexts {
		if strings.Contains(sys, "调研量子纠错现状") {
			t.Fatalf("system message carries the task text (%q); the ruling keeps the task out of system", sys)
		}
	}
}

// TestSpawnInjectsDataParagraph 断言 AgentSpec.Data 作为同一条 user 消息内的
// [任务数据] 段落到达模型（库路径调用方同样受裁决覆盖）。
func TestSpawnInjectsDataParagraph(t *testing.T) {
	t.Parallel()
	ad := &requestCaptureAdapter{final: "done"}
	s := captureLoop(t, ad, nil)

	sub, err := s.Spawn(context.Background(), AgentSpec{
		Goal: "处理数据",
		Data: map[string]any{"repo": "cyrene", "stars": 2},
	})
	if err != nil {
		t.Fatalf("Spawn: %v", err)
	}
	if _, werr := s.Wait(context.Background(), []string{sub.ID}); werr != nil {
		t.Fatalf("Wait: %v", werr)
	}
	msgs := ad.request(0).Messages
	var task string
	users := 0
	for _, m := range msgs {
		if m.Role != model.RoleUser {
			continue
		}
		users++
		task = messageTextOf(m)
	}
	if users != 1 {
		t.Fatalf("user messages = %d, want one", users)
	}
	for _, want := range []string{"[任务数据]", `"repo":"cyrene"`} {
		if !strings.Contains(task, want) {
			t.Fatalf("task message = %q, want it to contain %q", task, want)
		}
	}
}

// TestSpawnTimeoutFromToolDefault 断言 WithSpawnTimeout（生产经
// SUBAGENT_TIMEOUT_SEC 填充）覆盖调度器默认时限（02.4.4）。
//
// 「批次在 30ms 前台窗口内没有结算」这一条由门控保证：唯一的并行槽位先被一个
// 托底子代理占住（它的时限是调度器默认的 1h，测试放行之前绝不结束），因此无论
// 机器多忙，被派发的子代理在窗口内都不可能开始跑，更不可能结算。原先这条是靠
// 「30ms 窗口 < 120ms 时限」的墙钟巧合成立的，负载一高就会先等到已结算的 "done"。
func TestSpawnTimeoutFromToolDefault(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 1, PerAgentTimeout: time.Hour})
	gate := make(chan struct{})
	s.runner = &mockRunner{gate: gate} // 阻塞直到测试放行或 ctx 超时

	// 托底者占住唯一槽位；它在测试放行之前不会结束。
	holder, err := s.Spawn(context.Background(), AgentSpec{Goal: "hold the only slot"})
	if err != nil {
		t.Fatalf("Spawn holder: %v", err)
	}
	waitFor(t, func() bool { return statusOf(s, holder.ID) == StatusRunning }, waitTimeout)

	sp := NewSpawnTool(s, WithSpawnWait(30*time.Millisecond), WithSpawnTimeout(120*time.Millisecond))
	raw, err := sp.Execute(context.Background(), map[string]any{"goal": "slow"})
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	got := decodeSpawnResult(t, raw)
	if got.Status != "running" {
		t.Fatalf("Status = %q, want running", got.Status)
	}
	id := got.Agents[0].ID

	// 放行托底者：槽位归还，被派发的子代理这才开始跑并在 120ms 上中止。
	// 该发送与托底者的接收配对（此时还没有别的接收者），因此它同样是确定性的。
	gate <- struct{}{}
	waitFor(t, func() bool { return statusOf(s, id) == StatusAborted }, waitTimeout)
	sub := s.pickOrFail(t, id)
	if !strings.Contains(sub.Summary, "120ms") {
		t.Fatalf("Summary = %q, want the 120ms tool-default deadline named", sub.Summary)
	}
	if !errors.Is(sub.Err, context.DeadlineExceeded) {
		t.Fatalf("Err = %v, want DeadlineExceeded", sub.Err)
	}
}

// TestSpawnTimeoutPerCallOverride 断言按调用覆盖（timeout_seconds）优先于
// 派发默认与调度器默认（design/05 §7 裁决 02.4.4）。
func TestSpawnTimeoutPerCallOverride(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 1, PerAgentTimeout: time.Hour})
	s.runner = &mockRunner{gate: make(chan struct{})}
	sp := NewSpawnTool(s, WithSpawnWait(30*time.Millisecond), WithSpawnTimeout(time.Hour))

	raw, err := sp.Execute(context.Background(), map[string]any{"goal": "slow", "timeout_seconds": float64(1)})
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	id := decodeSpawnResult(t, raw).Agents[0].ID
	waitFor(t, func() bool { return statusOf(s, id) == StatusAborted }, 3*time.Second)
	sub := s.pickOrFail(t, id)
	if !strings.Contains(sub.Summary, "1s") {
		t.Fatalf("Summary = %q, want the 1s per-call deadline named", sub.Summary)
	}
}

// TestSpawnTimeoutSpecOverride 是库路径的按调用覆盖断言：AgentSpec.Timeout
// 优先于 Options.PerAgentTimeout。
func TestSpawnTimeoutSpecOverride(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 1, PerAgentTimeout: time.Hour})
	s.runner = &mockRunner{gate: make(chan struct{})}

	sub, err := s.Spawn(context.Background(), AgentSpec{Goal: "quick", Timeout: 80 * time.Millisecond})
	if err != nil {
		t.Fatalf("Spawn: %v", err)
	}
	waitFor(t, func() bool { return sub.Status() == StatusAborted }, waitTimeout)
	if !strings.Contains(sub.Summary(), "80ms") {
		t.Fatalf("Summary = %q, want the 80ms spec deadline named", sub.Summary())
	}
	if sub.FinishedAt().Sub(sub.StartedAt()) > 5*time.Second {
		t.Fatalf("run took %v; the per-call override never applied", sub.FinishedAt().Sub(sub.StartedAt()))
	}
}

// pickOrFail 按 id 取子代理快照，找不到时失败测试。
func (s *Scheduler) pickOrFail(t *testing.T, id string) SubAgentResult {
	t.Helper()
	for _, a := range s.Agents() {
		if a.ID == id {
			return a
		}
	}
	t.Fatalf("sub-agent %q not registered", id)
	return SubAgentResult{}
}

// TestLateSubagentSpanRejoinsOriginalRound 是 02.4.4「迟到 span 续根」的
// 归属断言（design/05 §9、design/19 §2）：回合（EndRound）先于子代理结束，
// 之后才到达的子代理 span 必须归入原轮（web-1）的泳道——根节点翻转为终态、
// 内部步骤链完整——而不是被丢弃，也不是另起新轮根。
func TestLateSubagentSpanRejoinsOriginalRound(t *testing.T) {
	t.Parallel()
	capture := &graphCapture{}
	flow := loop.NewFlowPublisher(capture.publish, nil)
	tracer := trace.New(nil)

	d, err := tool.New(registry.New())
	if err != nil {
		t.Fatalf("tool.New: %v", err)
	}
	if rerr := d.Register(&echoTool{name: "echo", result: "ok", delay: 150 * time.Millisecond}); rerr != nil {
		t.Fatalf("Register: %v", rerr)
	}
	s := branchLoop(t, flow, tracer, d, Options{MaxParallel: 1})

	flow.StartRound(parentSession, "帮我调研")
	sub, serr := s.Spawn(context.Background(), AgentSpec{Goal: "慢速调研", ParentSession: parentSession})
	if serr != nil {
		t.Fatalf("Spawn: %v", serr)
	}
	// 回合在子代理仍在运行时就结束：其后到达的 span 都是「迟到 span」。
	flow.EndRound(parentSession, "结果如下")
	if _, werr := s.Wait(context.Background(), []string{sub.ID}); werr != nil {
		t.Fatalf("Wait: %v", werr)
	}

	// 不丢、不另起根：整个会话只出版过一轮（web-1），迟到的泳道落在它里面。
	if ids := capture.ids(parentSession); len(ids) != 1 || ids[0] != parentSession+"-1" {
		t.Fatalf("published rounds for %q = %v, want exactly the original round %s-1", parentSession, ids, parentSession)
	}
	if capture.has(sub.ID) {
		t.Fatalf("the late sub-agent opened a graph of its own (%s); it must rejoin the original round", sub.ID)
	}
	g := capture.graph(t, parentSession)
	if g.ID != parentSession+"-1" {
		t.Fatalf("graph id = %q, want the original round's %s-1", g.ID, parentSession)
	}
	// 原轮的收尾仍在，且迟到泳道完整挂靠：根节点终态 + llm>tool>llm 步骤链。
	if flowNodeByID(t, g, "output").Kind != "output" {
		t.Fatal("the closed round lost its output node")
	}
	subNode := flowSubagentNode(t, g, sub.ID)
	if subNode.Status != trace.FlowStatusDone {
		t.Fatalf("late subagent node status = %q, want done (the late finish must flip its root)", subNode.Status)
	}
	if chain := flowLaneChain(t, g, subNode); len(chain) != 3 || chainKinds(chain) != "llm>tool>llm" {
		t.Fatalf("late lane chain = %s, want llm>tool>llm inside the original round", chainKinds(chain))
	}
}

// TestTaskContextIsolationFromCallerMutation 锁定注入消息与调用方切片的隔离：
// Spawn 后修改 spec.Context / spec.ToolWhitelist 不得影响已创建的子代理。
func TestTaskContextIsolationFromCallerMutation(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 1})
	m := &ctxCaptureRunner{}
	s.runner = m

	wl := []string{"alpha"}
	spec := AgentSpec{Goal: "g", ToolWhitelist: wl, Context: []model.Message{model.TextMessage(model.RoleUser, "seed")}}
	sub, err := s.Spawn(context.Background(), spec)
	if err != nil {
		t.Fatalf("Spawn: %v", err)
	}
	wl[0] = "mutated"
	spec.Context[0].Content[0].Text = "mutated"

	if _, werr := s.Wait(context.Background(), []string{sub.ID}); werr != nil {
		t.Fatalf("Wait: %v", werr)
	}
	if got := loop.ToolWhitelistFrom(m.last()); len(got) != 1 || got[0] != "alpha" {
		t.Fatalf("whitelist = %v, want the snapshot [alpha] immune to caller mutation", got)
	}
	if sub.ToolWhitelist[0] != "alpha" {
		t.Fatalf("sub-agent whitelist = %v, want [alpha]", sub.ToolWhitelist)
	}
	if sub.Context[0].Content[0].Text != "seed" {
		t.Fatalf("sub-agent context = %q, want the isolated seed", sub.Context[0].Content[0].Text)
	}
}