package orchestration

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

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

// 本文件钉住**真并行路里派发的子代理落进那一路的图**（切片 FV2；design/15 §2.5.10-1）。
//
// 症状：同一会话真并行时两路各有一张图，而子代理的 span 归属取自 `AgentSpec.ParentSession`
// —— 那是**会话**（总线载荷要用它），不含路 ⇒ 子代理的泳道会落进「会话主干那张」，而那张图
// 在纯并行轮次里根本不存在（span 被静默丢弃），或者落进另一路。这里断言的是：派发工具从
// ctx 取出的**归属桶**（`{会话}#{路}`）一路传到发布器，泳道根只长在**那一路**的图上。

// laneGraphCapture 按 (会话, 路) 记录发布器发出的每一张图（路的键就是帧上的 `turnId`）。
type laneGraphCapture struct {
	mu    sync.Mutex
	lanes map[string]*trace.FlowGraph
}

func laneKey(session, turnID string) string { return session + "|" + turnID }

func (c *laneGraphCapture) publish(_ string, payload any) {
	g, ok := payload.(*trace.FlowGraph)
	if !ok || g == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.lanes == nil {
		c.lanes = make(map[string]*trace.FlowGraph)
	}
	c.lanes[laneKey(g.SessionID, g.TurnID)] = g
}

// lane 返回某一路最新那张图；不存在时判测试失败（并列出收到过哪些路）。
func (c *laneGraphCapture) lane(t *testing.T, session, turnID string) *trace.FlowGraph {
	t.Helper()
	c.mu.Lock()
	defer c.mu.Unlock()
	g := c.lanes[laneKey(session, turnID)]
	if g == nil {
		t.Fatalf("没有收到会话 %q 路 %q 的图（收到过的路：%v）", session, turnID, c.keysLocked())
	}
	return g
}

func (c *laneGraphCapture) keysLocked() []string {
	out := make([]string, 0, len(c.lanes))
	for k := range c.lanes {
		out = append(out, k)
	}
	return out
}

// branchOf 造一个「运行在某一路上」的 ctx（网关经 WithParallelLeg 盖上会话 + 路；
// 子代理自身则经 WithFlowBranch 带着父归属桶）。两条通道的 `FlowOwnerFromContext`
// 取值同源，因此派发工具的取值路径只有这一条。
func branchOf(session, lane string) context.Context {
	return loop.WithFlowBranch(context.Background(), loop.FlowOwnerKey(session, lane))
}

// TestSpawnToolBindsSubAgentToItsLaneGraph 是核心断言：在一路里派发的子代理，其泳道根
// 只出现在**那一路**的图上，且会话主干那张图**不存在**（＝没有被静默丢弃或错落他路）。
func TestSpawnToolBindsSubAgentToItsLaneGraph(t *testing.T) {
	t.Parallel()
	capture := &laneGraphCapture{}
	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"}); rerr != nil {
		t.Fatalf("Register: %v", rerr)
	}
	s := branchLoop(t, flow, tracer, d, Options{MaxParallel: 2})

	// 这一路已经开跑（网关在调用 Run 之前就发过 turn_begin，图的归属桶因此已存在）。
	owner := loop.FlowOwnerKey(parentSession, "turn-a")
	flow.StartRound(owner, "并行的一路")
	sp := NewSpawnTool(s, WithSpawnWait(5*time.Second))
	if _, xerr := sp.Execute(branchOf(parentSession, "turn-a"), map[string]any{"goal": "调研一下"}); xerr != nil {
		t.Fatalf("派发: %v", xerr)
	}
	flow.EndRound(owner, "结果如下")

	g := capture.lane(t, parentSession, "turn-a")
	roots := 0
	for _, n := range g.Nodes {
		if n.Kind == "subagent" {
			roots++
		}
	}
	if roots != 1 {
		t.Fatalf("该路的图里子代理泳道根 = %d 个，want 1（nodes=%v）", roots, nodeIDs(g))
	}
	if strings.HasSuffix(g.ID, "#turn-a") == false {
		t.Fatalf("该路的图 id = %q，want 带 `#turn-a` 后缀", g.ID)
	}
	// 会话主干那张（无路标识）绝不该被这一路的子代理污染：它要么不存在，要么不含泳道根。
	capture.mu.Lock()
	trunk := capture.lanes[laneKey(parentSession, "")]
	capture.mu.Unlock()
	if trunk != nil {
		for _, n := range trunk.Nodes {
			if n.Kind == "subagent" {
				t.Fatalf("子代理泳道根长在了会话主干图上（路桶 %v 才是它该在的地方）", capture.keysLocked())
			}
		}
	}
}