package trace

import (
	"fmt"
	"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"
)

// FlowStatusForPhase 把工具执行的细粒度 phase（design/04 §1.1，
// loop 随 tool_execution.end 帧一同记在 tool span 上）映射为流程图
// 节点状态：phase 取值与本包状态枚举（design/19 §1 统一枚举）
// 同构（running/done/error）。未知 phase 返回空串，调用方回落到
// span 终态推导，因此缺 phase 的旧 span 不受影响。
func FlowStatusForPhase(phase string) string {
	switch phase {
	case FlowStatusRunning:
		return FlowStatusRunning
	case FlowStatusDone:
		return FlowStatusDone
	case FlowStatusError:
		return FlowStatusError
	default:
		return ""
	}
}

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

// SeedInput creates the "input" node for the round. It returns the node id so a
// caller can reference it. The label is redacted (design/19 §3), so a message that
// happens to carry a credential cannot publish one.
func (g *RunGraph) SeedInput(text string) string {
	g.Upsert(NodeSpec{
		ID:        "input",
		Label:     RedactText(text),
		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. The label is redacted like SeedInput's.
func (g *RunGraph) SeedOutput(text string) string {
	g.ended = time.Now()
	g.Upsert(NodeSpec{
		ID:        "output",
		Label:     RedactText(text),
		Kind:      "output",
		Status:    FlowStatusDone,
		StartedAt: g.ended,
		EndedAt:   g.ended,
	})
	return "output"
}

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