"""Tracing helpers shared by every node."""

import time
from functools import wraps
from typing import Callable, Optional


def trace_event(node: str, *, status: str = "ok", duration_ms: float = 0.0,
                summary: str = "", payload: Optional[dict] = None) -> dict:
    """Build a single TraceEvent dict matching the state schema."""
    return {
        "node": node,
        "status": status,
        "duration_ms": round(duration_ms, 2),
        "summary": summary,
        "payload": payload or {},
    }


def traced(node_name: str) -> Callable:
    """
    Decorator: time a node, swallow exceptions into a trace event, and
    guarantee the node always returns a dict containing a `trace` list.

    The wrapped function should return a partial state dict WITHOUT a `trace`
    field — this decorator injects the timing/status event automatically. If
    the wrapped function returns its own `trace` list, those events are kept
    and the timing event is appended.
    """
    def decorator(fn):
        @wraps(fn)
        def wrapper(state, *args, **kwargs):
            t0 = time.perf_counter()
            try:
                result = fn(state, *args, **kwargs) or {}
                dt_ms = (time.perf_counter() - t0) * 1000
                summary = result.pop("_summary", "")
                payload = result.pop("_payload", {})
                event = trace_event(
                    node_name,
                    status="ok",
                    duration_ms=dt_ms,
                    summary=summary,
                    payload=payload,
                )
                existing = result.get("trace", [])
                result["trace"] = existing + [event]
                return result
            except Exception as e:
                dt_ms = (time.perf_counter() - t0) * 1000
                event = trace_event(
                    node_name,
                    status="error",
                    duration_ms=dt_ms,
                    summary=f"{type(e).__name__}: {e}",
                    payload={"error": str(e)},
                )
                return {"trace": [event], "error": f"{node_name}: {e}"}
        return wrapper
    return decorator