package orchestration

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

	"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"
)

// parentSession is the flow session the branch tests attribute their sub-agents
// to; it mirrors the web chat session the gateway publishes under.
const parentSession = "web"

// secretArg and secretValue are credential-shaped strings planted in a tool call's
// arguments and in its result. Neither may appear in a published flow graph.
const (
	secretArg   = "sk-arg-0123456789abcdef"
	secretValue = "sk-live-0123456789abcdef"
)

// graphCapture records every flow graph the publisher emits, per session.
type graphCapture struct {
	mu        sync.Mutex
	bySession map[string]*trace.FlowGraph
	roundIDs  map[string][]string
}

func (c *graphCapture) publish(_ string, payload any) {
	g, ok := payload.(*trace.FlowGraph)
	if !ok || g == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.bySession == nil {
		c.bySession = make(map[string]*trace.FlowGraph)
		c.roundIDs = make(map[string][]string)
	}
	c.bySession[g.SessionID] = g
	if ids := c.roundIDs[g.SessionID]; len(ids) == 0 || ids[len(ids)-1] != g.ID {
		c.roundIDs[g.SessionID] = append(ids, g.ID)
	}
}

func (c *graphCapture) graph(t *testing.T, session string) *trace.FlowGraph {
	t.Helper()
	c.mu.Lock()
	defer c.mu.Unlock()
	g := c.bySession[session]
	if g == nil {
		t.Fatalf("no flow graph published for session %q (captured sessions: %v)", session, c.sessionsLocked())
	}
	return g
}

func (c *graphCapture) ids(session string) []string {
	c.mu.Lock()
	defer c.mu.Unlock()
	return append([]string(nil), c.roundIDs[session]...)
}

func (c *graphCapture) has(session string) bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	_, ok := c.bySession[session]
	return ok
}

func (c *graphCapture) sessionsLocked() []string {
	out := make([]string, 0, len(c.bySession))
	for s := range c.bySession {
		out = append(out, s)
	}
	return out
}

// toolThenFinalAdapter behaves like a model that asks for one tool call and then
// answers: it returns the scripted call until the request already carries a tool
// result. Deciding from the request (rather than a call counter) keeps it correct
// when two sub-agents drive the same instance concurrently.
type toolThenFinalAdapter struct {
	toolName string
	args     map[string]any
	final    string
}

func (a *toolThenFinalAdapter) Name() string { return "mock" }

func (a *toolThenFinalAdapter) Chat(_ context.Context, req model.Request, _ model.Options) (model.Response, error) {
	for _, m := range req.Messages {
		if m.Role == model.RoleTool {
			return model.Response{Model: "mock-1", Content: a.final, StopReason: "end_turn"}, nil
		}
	}
	return model.Response{Model: "mock-1", ToolCalls: []model.ToolCall{
		{ID: "call-1", Name: a.toolName, Args: a.args},
	}}, nil
}

func (a *toolThenFinalAdapter) ChatStream(ctx context.Context, req model.Request, on func(model.StreamEvent)) (model.Response, error) {
	// Emit a delta so a test can detect whether the loop forwarded the sub-agent's
	// stream to the caller's UI channel.
	if on != nil {
		on(model.StreamEvent{Kind: model.StreamKindTextDelta, Delta: "内部草稿"})
	}
	return a.Chat(ctx, req, model.Options{})
}

// echoTool is a minimal tool that reports back a scripted result after an optional
// delay, so two branches' steps overlap in time.
type echoTool struct {
	name   string
	result any
	delay  time.Duration
}

func (e *echoTool) Name() string                        { return e.name }
func (e *echoTool) Description() string                 { return "echo the scripted result" }
func (e *echoTool) Parameters() (map[string]any, error) { return nil, nil }

func (e *echoTool) Execute(ctx context.Context, _ any) (any, error) {
	if e.delay > 0 {
		select {
		case <-time.After(e.delay):
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
	return e.result, nil
}

// branchLoop builds a real loop over the given dispatcher with tracing and flow
// publication wired, plus a scheduler driving it — the production shape, minus the
// model provider.
func branchLoop(t *testing.T, flow *loop.FlowPublisher, tracer *trace.Tracer, d *tool.Dispatcher, opts Options) *Scheduler {
	t.Helper()
	ad := &toolThenFinalAdapter{toolName: "echo", args: map[string]any{
		"command": "date",
		"api_key": secretArg,
	}, final: "子代理完成"}
	ag, err := agent.New("be brief", 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, Tracer: tracer, Flow: flow, MaxRounds: 8})
	opts.Loop = lp
	opts.Tracer = tracer
	opts.Flow = flow
	return newScheduler(t, opts)
}

// flowNodeByID returns the node with the given id, failing the test when absent.
func flowNodeByID(t *testing.T, g *trace.FlowGraph, id string) trace.FlowNode {
	t.Helper()
	for _, n := range g.Nodes {
		if n.ID == id {
			return n
		}
	}
	t.Fatalf("no node %q in graph %s (nodes: %v)", id, g.ID, nodeIDs(g))
	return trace.FlowNode{}
}

func nodeIDs(g *trace.FlowGraph) []string {
	out := make([]string, 0, len(g.Nodes))
	for _, n := range g.Nodes {
		out = append(out, n.ID)
	}
	return out
}

// flowKindNodes returns every node of the given kind, in graph insertion order
// (which is the order the spans started, i.e. the lane's step order).
func flowKindNodes(g *trace.FlowGraph, kind string) []trace.FlowNode {
	out := make([]trace.FlowNode, 0, len(g.Nodes))
	for _, n := range g.Nodes {
		if n.Kind == kind {
			out = append(out, n)
		}
	}
	return out
}

// flowLaneNodes returns the nodes inside a lane (a reachable set) in insertion
// order.
func flowLaneNodes(g *trace.FlowGraph, lane map[string]bool) []trace.FlowNode {
	out := make([]trace.FlowNode, 0, len(lane))
	for _, n := range g.Nodes {
		if lane[n.ID] {
			out = append(out, n)
		}
	}
	return out
}

// flowLaneChain returns the branch's steps ordered by insertion time and asserts
// they form one chain: the subagent root feeds the first step and each step feeds
// the next. A cross-linked lane (one step reached from two parents) or a stray node
// breaks the chain assertion.
func flowLaneChain(t *testing.T, g *trace.FlowGraph, sub trace.FlowNode) []trace.FlowNode {
	t.Helper()
	lane := flowReachable(g, sub.ID)
	if len(lane) == 0 {
		t.Fatalf("subagent node %s (label %s) produced no steps", sub.ID, sub.Label)
	}
	chain := flowLaneNodes(g, lane)
	if !hasFlowEdge(g, sub.ID, chain[0].ID) {
		t.Fatalf("edges = %+v, want the lane to start at %s (root %s)", g.Edges, sub.ID, chain[0].ID)
	}
	for i := 1; i < len(chain); i++ {
		if !hasFlowEdge(g, chain[i-1].ID, chain[i].ID) {
			t.Fatalf("lane chain broken between %s and %s (edges: %+v)", chain[i-1].ID, chain[i].ID, g.Edges)
		}
	}
	return chain
}

// chainKinds renders a lane chain as "llm>tool>llm" for failure messages.
func chainKinds(chain []trace.FlowNode) string {
	parts := make([]string, 0, len(chain))
	for _, n := range chain {
		parts = append(parts, n.Kind)
	}
	return strings.Join(parts, ">")
}

// flowSubagentNode returns the branch root whose label names the given sub-agent
// id. A node is keyed by its span id, so the label is what identifies the owner.
func flowSubagentNode(t *testing.T, g *trace.FlowGraph, subID string) trace.FlowNode {
	t.Helper()
	want := "subagent." + subID
	for _, n := range flowKindNodes(g, "subagent") {
		if n.Label == want || n.ID == want {
			return n
		}
	}
	t.Fatalf("no subagent node labelled %q in graph %s (nodes: %v)", want, g.ID, nodeIDs(g))
	return trace.FlowNode{}
}

// flowReachable returns the set of node ids reachable from id, excluding id.
func flowReachable(g *trace.FlowGraph, id string) map[string]bool {
	out := make(map[string]bool)
	queue := []string{id}
	for len(queue) > 0 {
		from := queue[0]
		queue = queue[1:]
		for _, e := range g.Edges {
			if e.From != from || out[e.To] || e.To == id {
				continue
			}
			out[e.To] = true
			queue = append(queue, e.To)
		}
	}
	return out
}

func hasFlowEdge(g *trace.FlowGraph, from, to string) bool {
	for _, e := range g.Edges {
		if e.From == from && e.To == to {
			return true
		}
	}
	return false
}

// flowLogEvent returns the node's log entry of the given type.
func flowLogEvent(t *testing.T, n trace.FlowNode, typ string) trace.Event {
	t.Helper()
	for _, ev := range n.Log {
		if ev.Type == typ {
			return ev
		}
	}
	t.Fatalf("node %s has no %q log entry (log: %+v)", n.ID, typ, n.Log)
	return trace.Event{}
}