package trace

import (
	"context"
	"encoding/json"
	"strings"
	"testing"
	"time"
)

// graphNodeKinds 返回图里出现的节点 kind 集合。
func graphNodeKinds(g *FlowGraph) map[string]bool {
	out := make(map[string]bool)
	for _, n := range g.Nodes {
		out[n.Kind] = true
	}
	return out
}

// spanRecordOn 在实时 tracer 上开一个 span、挂上 props、关掉它，并返回完成记录。
// id 由 tracer 生成，因此每个记录天然唯一。
func spanRecordOn(t *testing.T, tr *Tracer, name string, kind SpanKind, props map[string]any) SpanRecord {
	t.Helper()
	_, sp := tr.StartSpan(context.Background(), name, kind)
	for k, v := range props {
		if err := sp.SetProp(k, v); err != nil {
			t.Fatalf("SetProp(%q): %v", k, err)
		}
	}
	sp.Finish(SpanStatusOK)
	return sp.Snapshot()
}

// TestGraphBuilder_MemoryGuardrailOrchestrationNodes 钉住三类旁路节点真的能
// 成为图上的节点（裁决 11.4.2）：memory / guardrail / orchestration 三类 span
// 各有生产者后，kindForSpan 的映射必须把它们原样折进图，而不是当未知 kind 丢掉。
func TestGraphBuilder_MemoryGuardrailOrchestrationNodes(t *testing.T) {
	t.Parallel()
	tr := New(nil)
	g := NewRunGraph("web")
	g.Start(1, time.Now())
	g.SeedInput("帮我看看")

	g.AddSpan(spanRecordOn(t, tr, "memory.recall", SpanKindMemory, map[string]any{"action": "recall", "scope": "owner"}))
	g.AddSpan(spanRecordOn(t, tr, "guardrail.input", SpanKindGuardrail, map[string]any{"action": "guardrail.input", "allow": true}))
	g.AddSpan(spanRecordOn(t, tr, "orchestration.spawn", SpanKindOrchestration, map[string]any{"subagent": "sub-1"}))

	got := graphNodeKinds(g.Snapshot())
	for _, want := range []string{"memory", "guardrail", "orchestration"} {
		if !got[want] {
			t.Fatalf("图中缺 %q 节点（实得 %v）", want, got)
		}
	}
}

// TestGraphBuilder_OrchestrationMarkerKeepsTrunkCursor 钉住编排派发标记的
// 定位：它挂在主干上但不接替主干链 —— 之后的 llm/tool 仍从原位置续链，
// 派发标记与泳道根一样只是主干上的一个标记（否则整条主干会被标记截断）。
func TestGraphBuilder_OrchestrationMarkerKeepsTrunkCursor(t *testing.T) {
	t.Parallel()
	tr := New(nil)
	g := NewRunGraph("web")
	g.Start(1, time.Now())
	g.SeedInput("派活")

	orch := spanRecordOn(t, tr, "orchestration.spawn", SpanKindOrchestration, nil)
	g.AddSpan(orch)
	llm := spanRecordOn(t, tr, "llm.call", SpanKindLLM, nil)
	g.AddSpan(llm)

	graph := g.Snapshot()
	if !hasEdge(graph, "input", orch.ID) {
		t.Fatalf("编排标记未挂在主干 input 上：%+v", graph.Edges)
	}
	if !hasEdge(graph, "input", llm.ID) {
		t.Fatalf("主干被编排标记截断：llm 应从 input 续链，实得 %+v", graph.Edges)
	}
}

// TestSafeSpanProps_ToolArgsSummaryKeepsSecretsOut 是发布面的脱敏用例
// （裁决 11.4.3）：把带密钥的参数摘要喂进 tool span，断言出图后没有密钥值，
// 只留下放行键的摘要与键名清单，且密钥类键在摘要里被替换为 [redacted]。
func TestSafeSpanProps_ToolArgsSummaryKeepsSecretsOut(t *testing.T) {
	t.Parallel()
	const secret = "sk-live-0123456789abcdef"
	tr := New(nil)
	g := NewRunGraph("web")
	g.Start(1, time.Now())
	g.SeedInput("读文件")

	g.AddSpan(spanRecordOn(t, tr, "tool.fs.read", SpanKindTool, map[string]any{
		"args_keys": []string{"api_key", "content", "path"},
		"args_summary": map[string]any{
			"path":    "/tmp/report.txt",
			"api_key": secret,
		},
		"result": secret,
	}))

	blob, err := json.Marshal(g.Snapshot())
	if err != nil {
		t.Fatalf("marshal graph: %v", err)
	}
	out := string(blob)
	if strings.Contains(out, secret) {
		t.Fatalf("出图泄漏了密钥值：%s", out)
	}
	if !strings.Contains(out, "/tmp/report.txt") {
		t.Fatalf("放行键的值没出图（可视化要能看参数）：%s", out)
	}
	if !strings.Contains(out, "args_summary") || !strings.Contains(out, "args_keys") {
		t.Fatalf("摘要键未随节点发布：%s", out)
	}
	if !strings.Contains(out, redactedMarker) {
		t.Fatalf("密钥类键的值应被替换为 %s：%s", redactedMarker, out)
	}
	if strings.Contains(out, "\"content\"") && strings.Contains(out, "args_summary\":{\"content") {
		t.Fatalf("表外键的值不应进入摘要：%s", out)
	}
}

// TestRedactValue_StringKeyedSummaryIsScrubbed 钉住 map[string]string 形态的
// 摘要同样逐键清洗：生产者若改用字符串键值对，密钥仍不能原样出图。
func TestRedactValue_StringKeyedSummaryIsScrubbed(t *testing.T) {
	t.Parallel()
	const secret = "sk-live-0123456789abcdef"
	out, ok := redactValue(map[string]string{"api_key": secret, "path": "/tmp/a"}).(map[string]any)
	if !ok {
		t.Fatal("map[string]string 未被清洗为 map[string]any")
	}
	if out["api_key"] != redactedMarker {
		t.Fatalf("密钥类键 = %v，want %s", out["api_key"], redactedMarker)
	}
	if out["path"] != "/tmp/a" {
		t.Fatalf("放行键值被改坏：%v", out["path"])
	}
}