package trace

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"
)

// Flow graph wire types (DevDocs/design/19 §1). The JSON field names are
// camelCase so the frontend's FlowGraph type and parseGraph read them directly.
// Node status strings mirror the frontend AgentStatus keys so parseGraph and
// FlowGraphView pass them through unchanged.

// Node status labels kept as named constants so no bare string appears in the
// mapping code (project rule §10).
const (
	// FlowStatusIdle is a node that has not started.
	FlowStatusIdle = "idle"
	// FlowStatusRunning is a node whose span is in flight.
	FlowStatusRunning = "running"
	// FlowStatusDone is a successfully completed node.
	FlowStatusDone = "done"
	// FlowStatusError is a node whose span finished with an error.
	FlowStatusError = "error"
)

// maxSummaryRunes bounds how many runes of a user string / tool result are
// rendered into a node label or prop summary, so sensitive or long content is
// collapsed to a safe, dashboard-friendly prefix.
const maxSummaryRunes = 80

// 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 = ""

// FlowNode is one node in a flow graph. Status/SpanKind mirror the frontier.
type FlowNode struct {
	ID        string         `json:"id"`
	Label     string         `json:"label"`
	Kind      string         `json:"kind,omitempty"`
	Status    string         `json:"status,omitempty"`
	StartedAt time.Time      `json:"startedAt,omitempty"`
	EndedAt   time.Time      `json:"endedAt,omitempty"`
	Props     map[string]any `json:"props,omitempty"`
	Log       []Event        `json:"log,omitempty"`
}

// FlowEdge is a directed edge between two nodes.
type FlowEdge struct {
	From string `json:"from"`
	To   string `json:"to"`
}

// FlowGraph is a live, single-round flow graph published as ws.graph_update.
type FlowGraph struct {
	ID        string     `json:"id"`
	Nodes     []FlowNode `json:"nodes"`
	Edges     []FlowEdge `json:"edges"`
	SessionID string     `json:"sessionId,omitempty"`
	StartedAt time.Time  `json:"startedAt,omitempty"`
	EndedAt   time.Time  `json:"endedAt,omitempty"`
}

// NodeSpec is the input to RunGraph.Upsert: the full description of a node plus
// the parent id used to wire its incoming edge.
type NodeSpec struct {
	ID        string
	Label     string
	Kind      string
	Status    string
	StartedAt time.Time
	EndedAt   time.Time
	Props     map[string]any
	Log       []Event
	ParentID  string
}

// RunGraph accumulates the spans of one round into a FlowGraph. It is a pure
// structure: it knows nothing about the bus, sessions or the loop, so it is
// unit-testable in isolation. A caller serializes access (see loop.FlowPublisher).
type RunGraph struct {
	session   string
	round     int
	nodes     map[string]*FlowNode
	order     []string
	edges     map[string]bool
	edgeOrder []FlowEdge
	// parents remembers every span's parent link, including spans that are not
	// rendered as nodes (a state_transition step container), so a node's lane can
	// be resolved by walking through the skipped containers.
	parents map[string]string
	// laneLast is the per-lane cursor used to chain sequential steps: keyed by the
	// enclosing subagent node id, and by "" for the main trunk. One cursor per lane
	// (instead of one per graph) is what keeps two parallel sub-agent lanes from
	// cross-linking when their spans interleave.
	laneLast map[string]string
	started  time.Time
	ended    time.Time
}

// NewRunGraph allocates an empty graph for one session round.
func NewRunGraph(session string) *RunGraph {
	return &RunGraph{
		session:  session,
		nodes:    make(map[string]*FlowNode),
		edges:    make(map[string]bool),
		parents:  make(map[string]string),
		laneLast: make(map[string]string),
	}
}

// Start marks the round began at now and resets the node/edge sets so a reused
// graph starts fresh for a new round.
func (g *RunGraph) Start(round int, now time.Time) {
	g.round = round
	g.started = now
	g.ended = time.Time{}
	g.nodes = make(map[string]*FlowNode)
	g.order = nil
	g.edges = make(map[string]bool)
	g.edgeOrder = nil
	g.parents = make(map[string]string)
	g.laneLast = make(map[string]string)
}

// 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)
	}
	// A sub-agent root is drawn on the trunk it hangs off, but it does not become
	// the trunk's continuation point: the main chain resumes where the branch left
	// it (the spawn node). Without this, every node after a spawn would look like a
	// step of the last branch (the frontend derives lanes from the edges). The
	// first trunk node still advances the cursor, so a graph built without an input
	// seed cannot orphan its later trunk nodes.
	if lane != "" || n.Kind != nodeKindSubagent || g.laneLast[trunkLane] == "" {
		g.laneLast[lane] = spec.ID
	}
}

// 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})
}

// SeedInput creates the "input" node for the round. It returns the node id so a
// caller can reference it.
func (g *RunGraph) SeedInput(text string) string {
	g.Upsert(NodeSpec{
		ID:        "input",
		Label:     truncate(text, maxSummaryRunes),
		Kind:      "input",
		Status:    FlowStatusDone,
		StartedAt: g.started,
	})
	return "input"
}

// SeedOutput creates the "output" node for the round, wired to the last node so
// the trunk ends at the answer.
func (g *RunGraph) SeedOutput(text string) string {
	g.ended = time.Now()
	g.Upsert(NodeSpec{
		ID:        "output",
		Label:     truncate(text, maxSummaryRunes),
		Kind:      "output",
		Status:    FlowStatusDone,
		StartedAt: g.ended,
		EndedAt:   g.ended,
	})
	return "output"
}

// 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.
func (g *RunGraph) AddSpan(rec SpanRecord) {
	kind := kindForSpan(rec.Kind)
	if kind == "" {
		g.rememberParent(rec.ID, rec.ParentID)
		return
	}
	g.Upsert(NodeSpec{
		ID:        rec.ID,
		Label:     labelForSpan(rec),
		Kind:      kind,
		Status:    statusForFinished(rec.Status),
		StartedAt: rec.StartedAt,
		EndedAt:   rec.EndedAt,
		Props:     safeSpanProps(rec),
		Log:       rec.Events,
		ParentID:  rec.ParentID,
	})
}

// 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
}

// Snapshot returns a copy of the graph safe to encode and publish.
func (g *RunGraph) Snapshot() *FlowGraph {
	nodes := make([]FlowNode, 0, len(g.order))
	for _, id := range g.order {
		nodes = append(nodes, *g.nodes[id])
	}
	return &FlowGraph{
		ID:        fmt.Sprintf("%s-%d", g.session, g.round),
		Nodes:     nodes,
		Edges:     append([]FlowEdge(nil), g.edgeOrder...),
		SessionID: g.session,
		StartedAt: g.started,
		EndedAt:   g.ended,
	}
}

// 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)
}

// safeSpanProps carries only a whitelist of span props to the frontend, plus a
// desensitized tool result summary. It never forwards raw arguments or secrets,
// and truncates string text so long content cannot leak a full sensitive payload.
func safeSpanProps(rec SpanRecord) map[string]any {
	var props map[string]any
	for _, k := range safePropKeys {
		v, ok := rec.Props[k]
		if !ok {
			continue
		}
		if props == nil {
			props = make(map[string]any)
		}
		props[k] = redactValue(v)
	}
	if rec.Kind == SpanKindTool {
		if v, ok := rec.Props["result"]; ok {
			if props == nil {
				props = make(map[string]any)
			}
			props["result_summary"] = truncate(summarize(v), maxSummaryRunes)
		}
	}
	return props
}

// redactValue shortens string props so a long error/step value cannot leak
// sensitive content through the dashboard; non-string values pass through.
func redactValue(v any) any {
	if s, ok := v.(string); ok {
		return truncate(s, maxSummaryRunes)
	}
	return v
}

// safePropKeys are the span prop keys forwarded to a node's props.
var safePropKeys = []string{"stop_reason", "step", "error", "error_kind", "tool"}

// 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
	}
}

// summarize renders an arbitrary value as a short text string for a prop summary.
func summarize(v any) string {
	switch x := v.(type) {
	case nil:
		return ""
	case string:
		return x
	default:
		b, err := json.Marshal(x)
		if err != nil {
			return fmt.Sprintf("%v", x)
		}
		return string(b)
	}
}

// truncate cuts s to at most max runes, appending an ellipsis when it overflows.
func truncate(s string, max int) string {
	if max <= 0 {
		return s
	}
	r := []rune(s)
	if len(r) <= max {
		return s
	}
	return string(r[:max]) + "…"
}