"""Shared utilities for the review pipeline."""

from __future__ import annotations

from typing import Any


def get_primary_model(trace: dict[str, Any]) -> str | None:
    """Extract the primary model from a trace (first LLM span's model field)."""
    for span in trace.get("spans", []):
        if span.get("kind") == "llm" and span.get("model"):
            return span["model"]
    return None


def build_trace_model_map(traces: list[dict[str, Any]]) -> dict[str, str]:
    """Map trace_id -> primary model for all traces that have one.

    Works with both full traces (extracts from spans) and summaries
    (reads primary_model field directly).
    """
    result = {}
    for t in traces:
        tid = t.get("trace_id", "")
        if not tid:
            continue
        # Summary path: primary_model already extracted
        model = t.get("primary_model")
        # Full trace path: extract from spans
        if not model and "spans" in t:
            model = get_primary_model(t)
        if model:
            result[tid] = model
    return result
