package trace

import (
	"strings"
	"time"
)

// Node kinds the builder emits. They are named constants because the builder
// itself branches on the subagent kind when it resolves a node's lane
// (design/19 §2: parallel sub-agent branches render as separate swimlanes).
const (
	nodeKindLLM         = "llm"
	nodeKindTool        = "tool"
	nodeKindSubagent    = "subagent"
	nodeKindMemory      = "memory"
	nodeKindGuardrail   = "guardrail"
	nodeKindOrchestrate = "orchestration"
)

// maxLaneWalk bounds how far the builder walks a span's parent chain when it
// resolves the lane (nearest enclosing subagent node) of a new node. Span
// parents form a tree, so the bound only guards against malformed input — it
// keeps the walk finite without a magic number in the loop (project rule §10).
const maxLaneWalk = 64

// trunkLane is the lane key of the main trunk: the lane of every span that has no
// sub-agent ancestor. It is a named constant so the map lookups read as intent.
const trunkLane = ""

// Upsert creates or merges a node and wires its incoming edge. The edge source is
// the node's ParentID when that parent is already a node; otherwise the span
// belongs to a lane (subagent branch) or the main trunk and is chained after the
// last node of that same lane — so a lane stays a coherent step chain even while
// several lanes are in flight, and the first node of a lane hangs off its
// subagent root (design/19 §2: 泳道内展示该子代理的步骤链).
func (g *RunGraph) Upsert(spec NodeSpec) {
	if spec.ID == "" {
		return
	}
	g.rememberParent(spec.ID, spec.ParentID)
	if existing := g.nodes[spec.ID]; existing != nil {
		g.merge(existing, spec)
		return
	}
	n := &FlowNode{
		ID:        spec.ID,
		Label:     spec.Label,
		Kind:      spec.Kind,
		Status:    spec.Status,
		StartedAt: spec.StartedAt,
		EndedAt:   spec.EndedAt,
	}
	if spec.Props != nil {
		n.Props = spec.Props
	}
	n.Log = append(n.Log, spec.Log...)
	g.nodes[spec.ID] = n
	g.order = append(g.order, spec.ID)

	lane := g.laneOf(spec.ParentID)
	from := g.laneLast[lane]
	if spec.ParentID != "" {
		if _, ok := g.nodes[spec.ParentID]; ok {
			from = spec.ParentID
		}
	}
	if from == "" {
		from = lane
	}
	if from != "" && from != spec.ID {
		g.addEdge(from, spec.ID)
	}
	// 两类节点是挂在主干上的「标记」而非主干的续点，因此不推进主干游标：
	// subagent 根开出一条泳道（见下），orchestration 记录一次派发/收束
	// （design/19 §1 旁路节点，裁决 11.4.2）—— 它们之后的主干节点仍从原
	// 位置续链，标记本身不会被误当成主干的一步。泳道内的节点（lane != ""）
	// 照常推进本泳道游标。
	if lane != "" || !branchMarker(n.Kind) || g.laneLast[trunkLane] == "" {
		g.laneLast[lane] = spec.ID
	}
}

// branchMarker 报告某类节点是否为「主干上的标记」：子代理泳道根与编排
// 派发节点都挂在主干上却不接替主干链，使主干（input→llm→tool→…→output）
// 与并行分支、派发标记各成一条不会被彼此打断的链。
func branchMarker(kind string) bool {
	return kind == nodeKindSubagent || kind == nodeKindOrchestrate
}

// rememberParent records a span's parent link even when the span is not rendered
// as a node, so laneOf can walk through skipped containers.
func (g *RunGraph) rememberParent(id, parent string) {
	if id == "" {
		return
	}
	g.parents[id] = parent
}

// laneOf returns the id of the nearest enclosing subagent node for a span's
// parent chain, or "" when the chain belongs to the main trunk. Only nodes
// already in the graph can be a lane root, so a lane resolves as soon as its
// first step is folded in.
func (g *RunGraph) laneOf(parentID string) string {
	for id, hops := parentID, 0; id != "" && hops < maxLaneWalk; hops++ {
		if n, ok := g.nodes[id]; ok && n.Kind == nodeKindSubagent {
			return id
		}
		id = g.parents[id]
	}
	return ""
}

// merge updates an existing node in place (a running placeholder being finalized
// by its finished span), keeping label/kind and appending props/log.
func (g *RunGraph) merge(existing *FlowNode, spec NodeSpec) {
	if spec.Label != "" {
		existing.Label = spec.Label
	}
	if spec.Kind != "" {
		existing.Kind = spec.Kind
	}
	if spec.Status != "" {
		existing.Status = spec.Status
	}
	if !spec.EndedAt.IsZero() {
		existing.EndedAt = spec.EndedAt
	}
	if len(spec.Props) > 0 {
		if existing.Props == nil {
			existing.Props = make(map[string]any, len(spec.Props))
		}
		for k, v := range spec.Props {
			existing.Props[k] = v
		}
	}
	existing.Log = append(existing.Log, spec.Log...)
}

// addEdge records a directed edge once.
func (g *RunGraph) addEdge(from, to string) {
	key := from + "\x00" + to
	if g.edges[key] {
		return
	}
	g.edges[key] = true
	g.edgeOrder = append(g.edgeOrder, FlowEdge{From: from, To: to})
}

// AddSpan folds a finished span record into the graph as a node with its final
// status. State-transition spans are the step container and are omitted — but
// their parent link is remembered so lane resolution can walk through them.
// tool 节点的终态优先由其携带的细粒度 phase 驱动（design/04 §1.1 /
// design/19 §1：running/done/error），缺失时回落到 span 终态推导，
// 因此两套来源对同一节点绝不可能得出不同状态。
func (g *RunGraph) AddSpan(rec SpanRecord) {
	kind := kindForSpan(rec.Kind)
	if kind == "" {
		g.rememberParent(rec.ID, rec.ParentID)
		return
	}
	status := statusForFinished(rec.Status)
	if kind == nodeKindTool {
		if ph := FlowStatusForPhase(spanPhase(rec)); ph != "" {
			status = ph
		}
	}
	g.Upsert(NodeSpec{
		ID:        rec.ID,
		Label:     labelForSpan(rec),
		Kind:      kind,
		Status:    status,
		StartedAt: rec.StartedAt,
		EndedAt:   rec.EndedAt,
		Props:     safeSpanProps(rec),
		Log:       redactEvents(rec.Events),
		ParentID:  rec.ParentID,
	})
}

// spanPhase 读取 tool span 上记录的执行 phase；缺失或类型不符时返回
// 空串（调用方据此回落），不区分「没有」与「读不到」。
func spanPhase(rec SpanRecord) string {
	ph, _ := rec.Props["phase"].(string)
	return ph
}

// AddSpanStart registers a running placeholder for a span that just began, so a
// node appears as "running" before it finishes. State-transition spans are
// omitted (their parent link is still remembered for lane resolution), matching
// AddSpan. It reports whether a node was created.
func (g *RunGraph) AddSpanStart(id, name string, kind SpanKind, parent string, at time.Time) bool {
	nodeKind := kindForSpan(kind)
	if nodeKind == "" {
		g.rememberParent(id, parent)
		return false
	}
	g.Upsert(NodeSpec{
		ID:        id,
		Label:     labelForStart(name, kind),
		Kind:      nodeKind,
		Status:    FlowStatusRunning,
		StartedAt: at,
		ParentID:  parent,
	})
	return true
}

// kindForSpan maps a span kind to a flow node kind. state_transition is the step
// container and maps to "" so it is never rendered as a node.
func kindForSpan(k SpanKind) string {
	switch k {
	case SpanKindLLM:
		return nodeKindLLM
	case SpanKindTool:
		return nodeKindTool
	case SpanKindSubagent:
		return nodeKindSubagent
	case SpanKindMemory:
		return nodeKindMemory
	case SpanKindGuardrail:
		return nodeKindGuardrail
	case SpanKindOrchestration:
		return nodeKindOrchestrate
	default:
		return ""
	}
}

// labelForSpan renders a node label from a finished span.
func labelForSpan(rec SpanRecord) string {
	if kind := kindForSpan(rec.Kind); kind == nodeKindTool {
		return strings.TrimPrefix(rec.Name, "tool.")
	}
	if rec.Kind == SpanKindLLM {
		return "思考"
	}
	if rec.Name != "" {
		return rec.Name
	}
	return string(rec.Kind)
}

// labelForStart renders a node label from a just-started span.
func labelForStart(name string, kind SpanKind) string {
	if kind == SpanKindLLM {
		return "思考"
	}
	if kind == SpanKindTool {
		return strings.TrimPrefix(name, "tool.")
	}
	if name != "" {
		return name
	}
	return string(kind)
}

// statusForFinished maps a span terminal status to a node status.
func statusForFinished(s SpanStatus) string {
	switch s {
	case SpanStatusError:
		return FlowStatusError
	case SpanStatusOK:
		return FlowStatusDone
	default:
		return FlowStatusIdle
	}
}