package orchestration

import (
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"strings"
	"sync"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/db"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/guardrail"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory"
	"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"
)

// roundExtractor 是确定性的记忆提取器：一轮结束后必然产生一次记忆写入，
// 使「真实跑一轮」里 memory 的写侧也有节点（不实现 asyncExtractor，因此
// 提取留在同步路径上，测试不用等后台）。
type roundExtractor struct{}

func (roundExtractor) Extract(context.Context, []model.Message) ([]memory.MemoryItem, error) {
	return []memory.MemoryItem{{Content: "子代理记住了这次委托"}}, nil
}

// delegateTool 模拟「派活」工具：它把前台会话当作 ParentSession 派一个后台
// 子代理（与生产里的 subagent_spawn 同一个归属口径：会话来自 ctx，绝不从
// 模型输出推断），并记下被派发的子代理 id 供测试等待。
type delegateTool struct {
	sched *Scheduler
	mu    sync.Mutex
	spawn []string
	err   error
}

func (d *delegateTool) Name() string                        { return "delegate" }
func (d *delegateTool) Description() string                 { return "dispatch a background sub-agent" }
func (d *delegateTool) Parameters() (map[string]any, error) { return nil, nil }

func (d *delegateTool) Execute(ctx context.Context, _ any) (any, error) {
	session := loop.SessionFromContext(ctx)
	if session == "" {
		d.setErr(context.Canceled)
		return "no session in context", nil
	}
	// 与生产的 subagent_spawn 同一口径：后台子代理必须活得比前台回合久，
	// 因此派发用一个脱离取消但保留会话值的 context。
	spawnCtx := context.WithoutCancel(ctx)
	sub, err := d.sched.Spawn(spawnCtx, AgentSpec{Goal: "调研一下", ParentSession: session})
	if err != nil {
		d.setErr(err)
		return "dispatch failed", nil
	}
	d.mu.Lock()
	d.spawn = append(d.spawn, sub.ID)
	d.mu.Unlock()
	return "已派发 " + sub.ID, nil
}

func (d *delegateTool) setErr(err error) {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.err == nil {
		d.err = err
	}
}

func (d *delegateTool) spawned() []string {
	d.mu.Lock()
	defer d.mu.Unlock()
	return append([]string(nil), d.spawn...)
}

// roundDiscardLogger 返回一个不写任何地方的 logger，保持测试输出干净。
func roundDiscardLogger() *slog.Logger {
	return slog.New(slog.NewTextHandler(io.Discard, nil))
}

// flowRound 是一轮「前台 + 后台子代理」的真实运行装置：它把两条真实循环
// （前台流式路径、子代理模板）、调度器与派活工具串起来。
type flowRound struct {
	front    *loop.Loop
	flow     *loop.FlowPublisher
	capture  *graphCapture
	sched    *Scheduler
	delegate *delegateTool
	ctx      context.Context
}

// newFlowRound 组装装置：真实 SQLite 内存库上的记忆栈、真实护栏、实时
// tracer/flow，前台走流式路径（生产 chat 路径），子代理模板同样带记忆。
func newFlowRound(t *testing.T) *flowRound {
	t.Helper()
	capture := &graphCapture{}
	flow := loop.NewFlowPublisher(capture.publish, nil)
	tracer := trace.New(nil)

	ctx := context.Background()
	sqlDB, oerr := db.Open(ctx, ":memory:", roundDiscardLogger())
	if oerr != nil {
		t.Fatalf("db.Open: %v", oerr)
	}
	t.Cleanup(func() { _ = sqlDB.Close() })
	memStore, merr := memory.NewStore(ctx, sqlDB, roundDiscardLogger())
	if merr != nil {
		t.Fatalf("memory.NewStore: %v", merr)
	}
	retriever := memory.NewRetriever(memStore)

	front, frontDispatch := newFrontRoundLoop(t, flow, tracer, memStore, retriever)
	subLoop := newSubRoundLoop(t, flow, tracer, memStore, retriever)
	sched := newScheduler(t, Options{MaxParallel: 1, Loop: subLoop, Tracer: tracer, Flow: flow, Memory: memStore})

	delegate := &delegateTool{sched: sched}
	if rerr := frontDispatch.Register(delegate); rerr != nil {
		t.Fatalf("Register delegate: %v", rerr)
	}
	return &flowRound{front: front, flow: flow, capture: capture, sched: sched, delegate: delegate, ctx: ctx}
}

// newFrontRoundLoop 构建前台循环：派活工具 + 流式路径（生产 chat 路径）+ 真实
// 记忆/护栏。它把 dispatcher 一并返回，调用方才能挂上需要调度器的派活工具。
func newFrontRoundLoop(t *testing.T, flow *loop.FlowPublisher, tracer *trace.Tracer, memStore *memory.Store, retriever *memory.Retriever) (*loop.Loop, *tool.Dispatcher) {
	t.Helper()
	dispatch, derr := tool.New(registry.New())
	if derr != nil {
		t.Fatalf("tool.New: %v", derr)
	}
	adapter := &toolThenFinalAdapter{toolName: "delegate", args: map[string]any{"goal": "调研"}, final: "已派发，稍后回填"}
	ag, aerr := agent.New("front", dispatch, nil,
		agent.WithModel(adapter, nil), agent.WithOptions(agent.Options{Model: "mock-1"}))
	if aerr != nil {
		t.Fatalf("agent.New: %v", aerr)
	}
	front := loop.New(loop.Options{
		Agent: ag, Tools: dispatch, Tracer: tracer, Flow: flow, MaxRounds: 8,
		Guards: guardrail.New(), Memory: memStore, Retriever: retriever,
		Extractor: roundExtractor{}, PersonaID: "cyrene",
		OnStream: func(model.StreamEvent) {},
	})
	return front, dispatch
}

// newSubRoundLoop 构建子代理模板：真实循环 + 真实记忆检索，因此泳道里也有
// memory 节点，且它的 llm/tool 步骤会折进派生它的会话的图。
func newSubRoundLoop(t *testing.T, flow *loop.FlowPublisher, tracer *trace.Tracer, memStore *memory.Store, retriever *memory.Retriever) *loop.Loop {
	t.Helper()
	dispatch, serr := tool.New(registry.New())
	if serr != nil {
		t.Fatalf("tool.New: %v", serr)
	}
	if rerr := dispatch.Register(&echoTool{name: "echo", result: map[string]any{"ok": true}}); rerr != nil {
		t.Fatalf("Register: %v", rerr)
	}
	adapter := &toolThenFinalAdapter{toolName: "echo", args: map[string]any{
		"path":    "/tmp/sub-report.txt",
		"api_key": secretArg,
	}, final: "子代理完成"}
	ag, aerr := agent.New("be brief", dispatch, nil,
		agent.WithModel(adapter, nil), agent.WithOptions(agent.Options{Model: "mock-1"}))
	if aerr != nil {
		t.Fatalf("agent.New: %v", aerr)
	}
	return loop.New(loop.Options{
		Agent: ag, Tools: dispatch, Tracer: tracer, Flow: flow, MaxRounds: 8,
		Memory: memStore, Retriever: retriever, Extractor: roundExtractor{}, PersonaID: "cyrene",
	})
}

// run 跑完前台一轮、等子代理收尾，并返回前台会话的最后一份图。
func (r *flowRound) run(t *testing.T) (*trace.FlowGraph, string) {
	t.Helper()
	r.flow.StartRound(parentSession, "帮我派个活")
	res, err := r.front.Run(r.ctx, parentSession, []model.Message{model.TextMessage(model.RoleUser, "帮我派个活")})
	if err != nil {
		t.Fatalf("front Run: %v", err)
	}
	r.flow.EndRound(parentSession, res.Text)
	if len(r.delegate.spawned()) != 1 {
		t.Fatalf("派发次数 = %v，want 1（delegate 错误：%v）", r.delegate.spawned(), r.delegate.err)
	}
	subID := r.delegate.spawned()[0]
	merged, werr := r.sched.Wait(r.ctx, []string{subID})
	if werr != nil {
		t.Fatalf("Wait: %v", werr)
	}
	if merged.Succeeded != 1 {
		t.Fatalf("子代理未成功：%+v（summary=%s）", merged, merged.Summary)
	}
	return r.capture.graph(t, parentSession), subID
}

// TestRealRound_FlowGraphHasAllThreeNodeKinds 是一条端到端的「真实跑一轮」用例
// （裁决 11.4.2）：前台会话带着真实记忆栈、真实护栏与实时 tracer/flow 跑完一轮，
// 期间经工具派发一个后台子代理 —— 同一张图里必须同时出现 memory（记忆召回/写入）、
// guardrail（输入/输出护栏）与 orchestration（派发标记）三类节点，且工具参数里的
// 密钥值绝不出图（裁决 11.4.3 红线）。
func TestRealRound_FlowGraphHasAllThreeNodeKinds(t *testing.T) {
	t.Parallel()
	r := newFlowRound(t)
	g, subID := r.run(t)

	kinds := make(map[string]int)
	for _, n := range g.Nodes {
		kinds[n.Kind]++
	}
	for _, want := range []string{"memory", "guardrail", "orchestration"} {
		if kinds[want] == 0 {
			t.Fatalf("图中缺 %q 节点（实得 %v）", want, kinds)
		}
	}
	assertOrchestrationNode(t, g, subID)
	assertSubagentLane(t, g, subID)
	assertRoundGraphRedaction(t, g)
}

// assertOrchestrationNode 断言派发标记：带被派发的子代理 id、收尾为 done。
func assertOrchestrationNode(t *testing.T, g *trace.FlowGraph, subID string) {
	t.Helper()
	orch := flowKindNodes(g, "orchestration")
	if len(orch) != 1 {
		t.Fatalf("orchestration 节点 = %d 个，want 1", len(orch))
	}
	if got := orch[0].Props["subagent"]; got != subID {
		t.Fatalf("orchestration 节点 subagent = %v，want %s（props：%+v）", got, subID, orch[0].Props)
	}
	if orch[0].Status != trace.FlowStatusDone {
		t.Fatalf("orchestration 节点状态 = %q，want done（props：%+v）", orch[0].Status, orch[0].Props)
	}
}

// assertSubagentLane 断言泳道仍在：子代理根 + 它自己的步骤链（含记忆召回与工具调用）。
func assertSubagentLane(t *testing.T, g *trace.FlowGraph, subID string) {
	t.Helper()
	subNode := flowSubagentNode(t, g, subID)
	lane := flowReachable(g, subNode.ID)
	hasMemoryInLane, hasToolInLane := false, false
	for _, n := range g.Nodes {
		if !lane[n.ID] {
			continue
		}
		if n.Kind == "memory" {
			hasMemoryInLane = true
		}
		if n.Kind == "tool" {
			hasToolInLane = true
		}
	}
	if !hasMemoryInLane || !hasToolInLane {
		t.Fatalf("子代理泳道缺 memory/tool 步骤（memory=%v tool=%v）", hasMemoryInLane, hasToolInLane)
	}
}

// assertRoundGraphRedaction 断言红线：工具参数里的密钥不得出现在发布的图里，
// 而放行的结构类键仍要可读（可视化要能看参数）。
func assertRoundGraphRedaction(t *testing.T, g *trace.FlowGraph) {
	t.Helper()
	blob, jerr := json.Marshal(g)
	if jerr != nil {
		t.Fatalf("marshal graph: %v", jerr)
	}
	if strings.Contains(string(blob), secretArg) {
		t.Fatalf("发布的图泄漏了工具参数：%s", blob)
	}
	if !strings.Contains(string(blob), "/tmp/sub-report.txt") {
		t.Fatalf("放行键 path 的值没出图：%s", blob)
	}
}