"""Pipeline tracing for verbose mode.

When verbose=True is passed to write() or retrieve(), a PipelineTrace
collects intermediate data from each pipeline stage. Zero overhead when
trace is None — each stage simply checks ``if trace:`` before recording.
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from datetime import datetime
from uuid import UUID


@dataclass
class PipelineStep:
    """A single step in a pipeline trace.

    The ``data`` dict contains stage-specific intermediate results.
    Step names and their data vary by pipeline:

    **Write pipeline steps:**

    - ``extraction`` — ``facts``, ``relations``, ``raw_response``
    - ``entity_resolution`` — ``entity_map`` (normalized key -> resolved entity)
    - ``reconciliation`` — ``decisions`` (list of ADD/UPDATE/NOOP/DELETE per fact)
    - ``upsert`` — ``added``, ``updated``, ``deleted``, ``unchanged`` counts

    **Read pipeline steps:**

    - ``retrieval_agent`` — ``plan`` (strategy, entities, similarity_query)
    - ``query_expansion`` — ``primed_entities``, ``expanded_terms``
    - ``semantic_search`` — ``candidates`` count, top scores
    - ``keyword_search`` — ``candidates`` count
    - ``graph_retrieval`` — ``facts`` count, ``edges_traversed``
    - ``reranker`` — ``reranked`` count, score adjustments
    - ``context_compression`` — ``hot_count``, ``warm_count``, ``cold_count``, ``total_tokens``
    """

    name: str
    duration_ms: float
    data: dict


@dataclass
class PipelineTrace:
    """Collects intermediate data from pipeline stages.

    Usage in pipeline functions::

        trace = PipelineTrace()
        t0 = trace.start()
        # ... run extraction ...
        trace.add_step("extraction", {"entities": [...]}, trace.elapsed(t0))
    """

    steps: list[PipelineStep] = field(default_factory=list)

    def start(self) -> float:
        """Return a perf_counter timestamp for measuring step duration."""
        return time.perf_counter()

    def elapsed(self, t0: float) -> float:
        """Return elapsed milliseconds since t0."""
        return round((time.perf_counter() - t0) * 1000, 1)

    def add_step(self, name: str, data: dict, duration_ms: float) -> None:
        """Record a pipeline step with its data and timing."""
        self.steps.append(PipelineStep(name=name, duration_ms=duration_ms, data=data))

    def to_dict(self) -> dict:
        """Serialize to a JSON-compatible dict."""
        return {
            "steps": [
                {
                    "name": step.name,
                    "duration_ms": step.duration_ms,
                    "data": _serialize(step.data),
                }
                for step in self.steps
            ]
        }


def _serialize(obj: object) -> object:
    """Recursively convert non-JSON types to JSON-serializable equivalents."""
    if isinstance(obj, dict):
        return {k: _serialize(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_serialize(item) for item in obj]
    if isinstance(obj, UUID):
        return str(obj)
    if isinstance(obj, datetime):
        return obj.isoformat()
    if hasattr(obj, "model_dump"):
        return obj.model_dump()
    if hasattr(obj, "__dataclass_fields__"):
        from dataclasses import asdict

        return _serialize(asdict(obj))  # type: ignore[call-overload]
    return obj
