package runtime

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

	"agent-native/runtime/llm"
	"agent-native/runtime/tools"
)

// ---- N3 WP4：循环 span 形态（名称后缀、关联、usage）----

type capturedSpan struct {
	name    string
	traceID string
	spanID  string
	parent  string
	attrs   map[string]string
}

type spanRecorder struct {
	mu    sync.Mutex
	spans []capturedSpan
}

func (r *spanRecorder) trace() func(string, map[string]string) func() {
	return func(name string, attrs map[string]string) func() {
		return func() {
			// 结束时拷贝（与 telemetry 工厂的结束时过滤同步——usage 是存续期补的）。
			sp := capturedSpan{name: name, attrs: map[string]string{}}
			for k, v := range attrs {
				sp.attrs[k] = v
			}
			r.mu.Lock()
			defer r.mu.Unlock()
			sp.traceID = sp.attrs["_trace_id"]
			sp.spanID = sp.attrs["_span_id"]
			sp.parent = sp.attrs["_parent_span_id"]
			r.spans = append(r.spans, sp)
		}
	}
}

func TestLoopSpansFollowGenAIConvention(t *testing.T) {
	rec := &spanRecorder{}
	reg := tools.NewRegistry(t.TempDir())
	reg.Register(&countingTool{name: "read_file"})
	res, err := Run(context.Background(), Config{
		Model: "claude-x", Tools: reg, MaxTurns: 3, Trace: rec.trace(),
		ProviderName: "anthropic", AgentName: "writer",
		Stream: script(
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
				ToolCalls: []llm.ToolCall{{ID: "c1", Name: "read_file", Arguments: `{}`}}},
			llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd},
		),
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	_ = res
	rec.mu.Lock()
	defer rec.mu.Unlock()
	if len(rec.spans) < 4 {
		t.Fatalf("span 数 = %d（invoke_agent + 2×chat + execute_tool）：%+v", len(rec.spans), rec.spans)
	}
	var root capturedSpan
	byName := map[string][]capturedSpan{}
	for _, sp := range rec.spans {
		switch {
		case strings.HasPrefix(sp.name, "invoke_agent "):
			byName["invoke"] = append(byName["invoke"], sp)
		case strings.HasPrefix(sp.name, "chat "):
			byName["chat"] = append(byName["chat"], sp)
		case strings.HasPrefix(sp.name, "execute_tool "):
			byName["tool"] = append(byName["tool"], sp)
		default:
			t.Fatalf("span 名不合约定：%q", sp.name)
		}
	}
	if len(byName["invoke"]) != 1 || byName["invoke"][0].name != "invoke_agent writer" {
		t.Fatalf("invoke_agent span：%+v", byName["invoke"])
	}
	root = byName["invoke"][0]
	if root.attrs["gen_ai.system"] != "anthropic" || root.attrs["gen_ai.agent.name"] != "writer" ||
		root.attrs["gen_ai.request.model"] != "claude-x" {
		t.Fatalf("根属性：%+v", root.attrs)
	}
	if len(byName["chat"]) != 2 {
		t.Fatalf("chat span 数：%+v", byName["chat"])
	}
	for _, sp := range byName["chat"] {
		if sp.name != "chat claude-x" {
			t.Fatalf("chat span 名：%q", sp.name)
		}
		if sp.traceID != root.traceID || sp.parent != root.spanID {
			t.Fatalf("关联断裂：trace=%q parent=%q 根=%q/%q", sp.traceID, sp.parent, root.traceID, root.spanID)
		}
		if sp.attrs["gen_ai.request.model"] != "claude-x" ||
			sp.attrs["gen_ai.usage.input_tokens"] == "" || sp.attrs["gen_ai.usage.output_tokens"] == "" {
			t.Fatalf("chat 属性/usage：%+v", sp.attrs)
		}
	}
	if len(byName["tool"]) != 1 || byName["tool"][0].name != "execute_tool read_file" {
		t.Fatalf("tool span：%+v", byName["tool"])
	}
	tool := byName["tool"][0]
	if tool.attrs["gen_ai.tool.name"] != "read_file" || tool.traceID != root.traceID || tool.parent != root.spanID {
		t.Fatalf("tool 属性/关联：%+v", tool)
	}
}