import { summarize, summarizeSpan, summarizeTrace } from "./summary.js";
// Event types the viewer skips entirely. `graph` is a one-shot
// schema dump (nodes + edges + start node) emitted at the top of
// each agent run; useful in the JSONL for tooling but not in the
// interactive tree.
const HIDDEN_EVENT_TYPES = new Set(["graph"]);
export function buildForest(events) {
    // traceId → trace root
    const traces = {};
    // span_id → span node (lookup across all traces; span_ids are globally unique per nanoid)
    const spans = {};
    // span_id → its desired parent_span_id (as observed on first sight).
    // Tracked separately so pass 1b can re-resolve parents once every
    // span exists, without polluting the public TreeNode shape.
    const desiredParent = {};
    // Pass 1a: create traces and spans, linking each span to its parent
    // (or trace root) in arrival order. This puts child spans into their
    // parent's `children` array BEFORE any leaf events get appended in
    // pass 2, so a span's child spans are always listed before its leaves.
    for (const evt of events) {
        ensureTrace(traces, evt.trace_id);
        if (evt.span_id) {
            ensureSpan(spans, traces, evt);
            if (!(evt.span_id in desiredParent)) {
                desiredParent[evt.span_id] = evt.parent_span_id ?? null;
            }
        }
    }
    // Pass 1b: re-resolve any span that was attached to the trace root
    // because its parent_span_id had not been seen yet. After pass 1a
    // every span that will ever exist does exist, so we can move the
    // child under its true parent if it shows up now. This makes the
    // tree shape order-independent.
    for (const span of Object.values(spans)) {
        const desiredParentId = desiredParent[span.id];
        if (!desiredParentId)
            continue;
        const trueParent = spans[desiredParentId];
        if (!trueParent || span.parentId === trueParent.id)
            continue;
        // Detach from current parent and re-attach to the true parent.
        const traceRoot = traces[span.traceId];
        traceRoot.children = traceRoot.children.filter((c) => c.id !== span.id);
        span.parentId = trueParent.id;
        trueParent.children.push(span);
    }
    // Pass 2: attach each event as a leaf under its span (or under the
    // trace root if it has no span_id), preserving arrival order. Some
    // event types are noise (e.g. the `graph` schema dump) and are
    // hidden from the viewer entirely.
    let leafCounter = 0;
    for (const evt of events) {
        if (HIDDEN_EVENT_TYPES.has(evt.data.type))
            continue;
        const traceRoot = traces[evt.trace_id];
        const parent = evt.span_id ? spans[evt.span_id] : traceRoot;
        const leaf = {
            id: `evt-${leafCounter++}`,
            traceId: evt.trace_id,
            parentId: parent.id,
            children: [],
            nodeKind: "event",
            label: evt.data.type,
            summary: summarize(evt),
            event: evt,
        };
        parent.children.push(leaf);
    }
    // Pass 3: aggregate metrics on each span.
    for (const span of Object.values(spans)) {
        aggregateMetrics(span);
    }
    // Aggregate trace-level metrics from immediate span children too.
    for (const trace of Object.values(traces)) {
        aggregateMetrics(trace);
        trace.summary = summarizeTrace(trace);
    }
    // Pass 4: sort each node's children in chronological order. Spans
    // sort by their aggregated firstTs (computed in pass 3); leaves by
    // their own timestamp. Children with missing/invalid timestamps
    // keep their original arrival position via a stable sort.
    for (const trace of Object.values(traces)) {
        sortChildrenByTime(trace);
    }
    return Object.values(traces);
}
function sortChildrenByTime(node) {
    // Decorate with the original index so we can fall back to arrival
    // order whenever timestamps tie or are missing (stable sort).
    const decorated = node.children.map((child, idx) => ({
        child,
        ts: nodeSortTs(child),
        idx,
    }));
    decorated.sort((a, b) => {
        if (a.ts !== b.ts)
            return a.ts - b.ts;
        return a.idx - b.idx;
    });
    node.children = decorated.map((d) => d.child);
    for (const child of node.children)
        sortChildrenByTime(child);
}
function nodeSortTs(node) {
    if (node.event) {
        const t = Date.parse(node.event.data.timestamp);
        return Number.isFinite(t) ? t : Number.POSITIVE_INFINITY;
    }
    return node.firstTs ?? Number.POSITIVE_INFINITY;
}
function ensureTrace(traces, traceId) {
    const existing = traces[traceId];
    if (existing)
        return existing;
    const root = {
        id: `trace-${traceId}`,
        traceId,
        parentId: null,
        children: [],
        nodeKind: "trace",
        label: traceId,
        summary: "",
    };
    traces[traceId] = root;
    return root;
}
function ensureSpan(spans, traces, evt) {
    const existing = spans[evt.span_id];
    if (existing)
        return existing;
    const node = {
        id: evt.span_id,
        traceId: evt.trace_id,
        parentId: evt.parent_span_id ?? null,
        children: [],
        nodeKind: "span",
        // Use the introducing event's type as a hint for the span type;
        // refined when we know more (e.g. agentStart → agentRun).
        label: inferSpanLabel(evt),
        summary: "",
    };
    spans[evt.span_id] = node;
    // Attach to parent span (if known) or to the trace root.
    const parent = evt.parent_span_id
        ? (spans[evt.parent_span_id] ?? ensureTrace(traces, evt.trace_id))
        : ensureTrace(traces, evt.trace_id);
    node.parentId = parent.id;
    parent.children.push(node);
    return node;
}
function inferSpanLabel(evt) {
    switch (evt.data.type) {
        case "agentStart":
        case "agentEnd":
            return "agentRun";
        case "enterNode":
            return "nodeExecution";
        case "promptCompletion":
            return "llmCall";
        case "embedCompletion":
            // Embedding spans share the chat-completion shape but get their
            // own label so the viewer can color/filter them separately and
            // cost roll-ups don't conflate the two.
            return "embedding";
        case "toolCall":
            return "toolExecution";
        case "forkStart":
        case "forkEnd":
            return evt.data.mode === "race" ? "race" : "forkAll";
        case "handlerDecision":
            return "handlerChain";
        default:
            // The memory umbrella events (`memoryRemember`, `memoryRecall`,
            // `memoryForget`, `memoryCompaction`) intentionally hit this
            // branch — their event type already matches the SpanType, so the
            // default returns the right label without a per-case mapping.
            return evt.data.type;
    }
}
function aggregateMetrics(node) {
    const leaves = [];
    walk(node, (n) => {
        if (n.event)
            leaves.push(n);
    });
    const sum = (xs) => xs.reduce((a, b) => a + b, 0);
    const tokens = sum(leaves
        .filter((l) => l.event.data.type === "promptCompletion")
        .map((l) => {
        const u = l.event.data.usage ?? {};
        return (u.inputTokens ?? 0) + (u.outputTokens ?? 0);
    }));
    const cost = sum(leaves
        .filter((l) => l.event.data.type === "promptCompletion")
        .map((l) => l.event.data.cost?.totalCost ?? 0));
    const timestamps = leaves
        .map((l) => Date.parse(l.event.data.timestamp))
        .filter(Number.isFinite);
    if (tokens > 0)
        node.tokens = tokens;
    if (cost > 0)
        node.cost = cost;
    if (timestamps.length >= 1) {
        node.firstTs = Math.min(...timestamps);
    }
    // Prefer the authoritative `timeTaken` field on the span's
    // characteristic leaf events. Event timestamps record *emission*
    // time, so a span with a single completion event (like an llmCall
    // with one promptCompletion) would otherwise compute a duration
    // of 0/1ms even when the LLM call took seconds.
    const measured = measuredDuration(node, leaves);
    if (measured !== undefined) {
        node.duration = measured;
    }
    else if (timestamps.length >= 2) {
        node.duration = Math.max(...timestamps) - Math.min(...timestamps);
    }
    if (node.nodeKind === "span") {
        node.summary = summarizeSpan(node);
    }
}
function measuredDuration(node, leaves) {
    if (node.nodeKind !== "span")
        return undefined;
    const matchingLeaves = leaves.filter((l) => l.event.data.type === durationEventType(node.label));
    const times = matchingLeaves
        .map((l) => l.event.data.timeTaken)
        .filter((t) => typeof t === "number");
    if (times.length === 0)
        return undefined;
    return times.reduce((a, b) => a + b, 0);
}
function durationEventType(spanLabel) {
    switch (spanLabel) {
        case "llmCall": return "promptCompletion";
        case "toolExecution": return "toolCall";
        case "agentRun": return "agentEnd";
        case "forkAll":
        case "race": return "forkEnd";
        default: return undefined;
    }
}
function walk(node, visit) {
    visit(node);
    for (const c of node.children)
        walk(c, visit);
}
