"""Execution trace collection and visualization.

Provides:

- **Trace models** — ``ExecutionTrace``, ``TraceSpan``, ``ToolCallSpan``, ``LLMCallSpan``
- **TraceCollector** — a permanent ``VMEventEmitter`` listener that builds
  ``ExecutionTrace`` objects from pathway events.  Partitions traces by
  ``execution_id`` so a single collector handles multiple runs.
- **trace_to_mermaid** — renders an ``ExecutionTrace`` as a Mermaid gantt chart.

The collector is attached once to the ``PathwayVM`` at construction and stays
attached for the lifetime of the VM.  Access traces via ``collector.last_trace``
or ``collector.get_trace(execution_id)``.
"""

from __future__ import annotations

from collections import OrderedDict
from typing import TYPE_CHECKING, Literal

from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from collections.abc import Callable

    from core.pathways.models.execution import PathwayEvent

# =============================================================================
# TRACE MODELS
# =============================================================================


class ToolCallSpan(BaseModel):
    """A single tool invocation within a node."""

    tool_name: str
    call_id: str | None = None
    duration_ms: float | None = None
    success: bool = True
    error: str | None = None


class LLMCallSpan(BaseModel):
    """A single LLM call within a node."""

    model: str
    tokens_in: int = 0
    tokens_out: int = 0
    duration_ms: float | None = None


class DecisionPoint(BaseModel):
    """A single LLM decision within an iterative tool-calling loop.

    Captures why the LLM chose specific tools at each round.
    """

    round: int
    reasoning: str | None = None
    tools_chosen: list[str] = Field(default_factory=list)
    tools_available: list[str] = Field(default_factory=list)


class TraceSpan(BaseModel):
    """Execution span for a single node."""

    node_id: str
    node_type: str
    status: Literal["success", "failure", "skipped"] = "success"
    duration_ms: float | None = None
    tool_calls: list[ToolCallSpan] = Field(default_factory=list)
    llm_calls: list[LLMCallSpan] = Field(default_factory=list)
    decisions: list[DecisionPoint] = Field(default_factory=list)


class ExecutionTrace(BaseModel):
    """Full execution trace for a pathway run.

    Collect via ``TraceCollector``, then inspect as a dict
    (``trace.model_dump()``) or render with ``trace_to_mermaid()``.
    """

    execution_id: str = ""
    pathway_id: str = ""
    total_duration_ms: float | None = None
    total_tokens: int = 0
    spans: list[TraceSpan] = Field(default_factory=list)


# =============================================================================
# PER-EXECUTION STATE (internal)
# =============================================================================


class _ExecutionState:
    """Mutable accumulator for a single execution's trace data."""

    __slots__ = (
        "current_node",
        "execution_id",
        "pathway_id",
        "pending_llm_calls",
        "pending_tools",
        "span_order",
        "spans",
        "total_duration_ms",
        "total_tokens",
    )

    def __init__(self, execution_id: str, pathway_id: str) -> None:
        self.execution_id = execution_id
        self.pathway_id = pathway_id
        self.total_duration_ms: float | None = None
        self.spans: dict[str, TraceSpan] = {}
        self.span_order: list[str] = []
        self.pending_tools: dict[str, tuple[ToolCallSpan, str]] = {}
        self.pending_llm_calls: dict[str, tuple[LLMCallSpan, str]] = {}
        self.current_node: str | None = None
        self.total_tokens: int = 0

    def build(self) -> ExecutionTrace:
        ordered_spans = [self.spans[nid] for nid in self.span_order if nid in self.spans]
        return ExecutionTrace(
            execution_id=self.execution_id,
            pathway_id=self.pathway_id,
            total_duration_ms=self.total_duration_ms,
            total_tokens=self.total_tokens,
            spans=ordered_spans,
        )


# =============================================================================
# TRACE COLLECTOR (permanent VMEventEmitter listener)
# =============================================================================

# Maximum number of completed traces to retain in memory.
_MAX_RETAINED_TRACES = 64


class TraceCollector:
    """Pathway event listener that accumulates ``ExecutionTrace`` objects.

    Designed to be attached **once** to a ``PathwayVM`` and remain attached
    for the VM's lifetime.  Each execution is partitioned by its
    ``execution_id``, so concurrent or sequential runs are tracked
    independently.

    Access traces after execution::

        vm.last_trace  # most recent completed trace
        vm.get_trace(exec_id)  # look up by execution ID
    """

    def __init__(self) -> None:
        # In-flight executions keyed by execution_id
        self._active: dict[str, _ExecutionState] = {}

        # Completed traces (bounded LRU)
        self._completed: OrderedDict[str, ExecutionTrace] = OrderedDict()

    # -- Public API -----------------------------------------------------------

    @property
    def last_trace(self) -> ExecutionTrace | None:
        """The most recently completed ``ExecutionTrace``, or ``None``."""
        if not self._completed:
            return None
        # Last inserted item
        return next(reversed(self._completed.values()))

    def get_trace(self, execution_id: str) -> ExecutionTrace | None:
        """Look up a completed trace by *execution_id*."""
        return self._completed.get(execution_id)

    # -- Listener interface (callable accepting PathwayEvent) -----------------

    def __call__(self, event: PathwayEvent) -> None:
        """Dispatch incoming pathway event to the appropriate handler."""
        handler = _TRACE_DISPATCH.get(event.event_type)
        if handler:
            handler(self, event)

    # -- Internal: resolve execution state ------------------------------------

    def _state_for(self, event: PathwayEvent) -> _ExecutionState | None:
        """Resolve the ``_ExecutionState`` for an event."""
        data = event.data or {}
        exec_id = str(data.get("execution_id", ""))
        if not exec_id:
            return None
        return self._active.get(exec_id)

    # -- Per-event handlers ---------------------------------------------------

    def _on_execution_started(self, event: PathwayEvent) -> None:
        data = event.data or {}
        exec_id = str(data.get("execution_id", ""))
        if not exec_id:
            return
        state = _ExecutionState(exec_id, event.pathway_id or "")
        self._active[exec_id] = state

    def _on_execution_completed(self, event: PathwayEvent) -> None:
        state = self._state_for(event)
        if state is None:
            return
        data = event.data or {}
        duration_ms = data.get("duration_ms")
        if isinstance(duration_ms, (int, float)):
            state.total_duration_ms = float(duration_ms)
        # Finalize and move to completed
        trace = state.build()
        self._completed[state.execution_id] = trace
        self._active.pop(state.execution_id, None)
        # Evict oldest if over limit
        while len(self._completed) > _MAX_RETAINED_TRACES:
            self._completed.popitem(last=False)

    def _on_node_started(self, event: PathwayEvent) -> None:
        state = self._state_for(event)
        if state is None:
            return
        node_id = event.node_id or ""
        data = event.data or {}
        node_type = str(data.get("node_type", "unknown"))
        span = TraceSpan(node_id=node_id, node_type=node_type)
        state.spans[node_id] = span
        if node_id not in state.span_order:
            state.span_order.append(node_id)
        state.current_node = node_id

    def _on_node_completed(self, event: PathwayEvent) -> None:
        state = self._state_for(event)
        if state is None:
            return
        node_id = event.node_id or ""
        data = event.data or {}
        span = state.spans.get(node_id)
        if span:
            span.status = "success"
            duration_ms = data.get("duration_ms")
            if isinstance(duration_ms, (int, float)):
                span.duration_ms = float(duration_ms)

    def _on_node_failed(self, event: PathwayEvent) -> None:
        state = self._state_for(event)
        if state is None:
            return
        node_id = event.node_id or ""
        data = event.data or {}
        span = state.spans.get(node_id)
        if span:
            span.status = "failure"
            duration_ms = data.get("duration_ms")
            if isinstance(duration_ms, (int, float)):
                span.duration_ms = float(duration_ms)

    def _on_tool_called(self, event: PathwayEvent) -> None:
        state = self._state_for(event)
        if state is None:
            return
        data = event.data or {}
        tool_name = str(data.get("tool_name", "unknown"))
        call_id = str(data.get("call_id", ""))
        node_id = event.node_id or state.current_node or ""
        tool_span = ToolCallSpan(tool_name=tool_name, call_id=call_id or None)
        span = state.spans.get(node_id)
        if span:
            span.tool_calls.append(tool_span)
        if call_id:
            state.pending_tools[call_id] = (tool_span, node_id)

    def _on_tool_completed(self, event: PathwayEvent) -> None:
        state = self._state_for(event)
        if state is None:
            return
        data = event.data or {}
        call_id = str(data.get("call_id", ""))
        if call_id and call_id in state.pending_tools:
            tool_span, _node_id = state.pending_tools.pop(call_id)
            duration_ms = data.get("duration_ms")
            if isinstance(duration_ms, (int, float)):
                tool_span.duration_ms = float(duration_ms)
            tool_span.success = True

    def _on_tool_failed(self, event: PathwayEvent) -> None:
        state = self._state_for(event)
        if state is None:
            return
        data = event.data or {}
        call_id = str(data.get("call_id", ""))
        if call_id and call_id in state.pending_tools:
            tool_span, _node_id = state.pending_tools.pop(call_id)
            duration_ms = data.get("duration_ms")
            if isinstance(duration_ms, (int, float)):
                tool_span.duration_ms = float(duration_ms)
            tool_span.success = False
            tool_span.error = event.error or str(data.get("error", ""))

    def _on_llm_call_started(self, event: PathwayEvent) -> None:
        state = self._state_for(event)
        if state is None:
            return
        data = event.data or {}
        node_id = event.node_id or state.current_node or ""
        model = str(data.get("model", "unknown"))
        call_id = str(data.get("call_id", ""))
        llm_span = LLMCallSpan(model=model)
        span = state.spans.get(node_id)
        if span:
            span.llm_calls.append(llm_span)
        if call_id:
            state.pending_llm_calls[call_id] = (llm_span, node_id)

    def _on_llm_call_completed(self, event: PathwayEvent) -> None:
        state = self._state_for(event)
        if state is None:
            return
        data = event.data or {}
        call_id = str(data.get("call_id", ""))
        tokens_in = data.get("tokens_in", 0)
        tokens_out = data.get("tokens_out", 0)
        duration_ms = data.get("duration_ms")
        model = str(data.get("model", "unknown"))

        if call_id and call_id in state.pending_llm_calls:
            llm_span, _node_id = state.pending_llm_calls.pop(call_id)
        else:
            node_id = event.node_id or state.current_node or ""
            llm_span = LLMCallSpan(model=model)
            span = state.spans.get(node_id)
            if span:
                span.llm_calls.append(llm_span)

        if isinstance(tokens_in, (int, float)):
            llm_span.tokens_in = int(tokens_in)
        if isinstance(tokens_out, (int, float)):
            llm_span.tokens_out = int(tokens_out)
        if isinstance(duration_ms, (int, float)):
            llm_span.duration_ms = float(duration_ms)

        state.total_tokens += llm_span.tokens_in + llm_span.tokens_out

    def _on_llm_decision(self, event: PathwayEvent) -> None:
        state = self._state_for(event)
        if state is None:
            return
        data = event.data or {}
        node_id = event.node_id or state.current_node or ""
        span = state.spans.get(node_id)
        if span:
            decision = DecisionPoint(
                round=int(data.get("round", 0)),
                reasoning=str(data.get("reasoning", "")) or None,
                tools_chosen=list(data.get("tools_chosen", [])),
                tools_available=list(data.get("tools_available", [])),
            )
            span.decisions.append(decision)


# Dispatch table: event_type string → handler method
_TRACE_DISPATCH: dict[str, Callable[[TraceCollector, PathwayEvent], None]] = {
    "execution_started": TraceCollector._on_execution_started,
    "execution_completed": TraceCollector._on_execution_completed,
    "node_started": TraceCollector._on_node_started,
    "node_completed": TraceCollector._on_node_completed,
    "node_failed": TraceCollector._on_node_failed,
    "tool_called": TraceCollector._on_tool_called,
    "tool_completed": TraceCollector._on_tool_completed,
    "tool_failed": TraceCollector._on_tool_failed,
    "llm_call_started": TraceCollector._on_llm_call_started,
    "llm_call_completed": TraceCollector._on_llm_call_completed,
    "llm_decision": TraceCollector._on_llm_decision,
}


# =============================================================================
# MERMAID GANTT RENDERER
# =============================================================================


def _format_gantt_duration(ms: float | None) -> str:
    """Format a duration for display in Mermaid gantt task labels."""
    if ms is None:
        return "?"
    if ms < 1:
        return "<1ms"
    if ms < 1000:
        return f"{ms:.0f}ms"
    if ms < 60_000:
        return f"{ms / 1000:.1f}s"
    return f"{ms / 60_000:.1f}m"


def _sanitize_gantt_label(text: str) -> str:
    """Remove characters that break Mermaid gantt task labels."""
    return text.replace(":", "-").replace(";", "-").replace(",", " ")


def trace_to_mermaid(trace: ExecutionTrace) -> str:
    """Render an ``ExecutionTrace`` as a Mermaid gantt chart.

    Each node becomes a task bar; tool calls are shown as sub-tasks
    within their parent node's section.

    Parameters
    ----------
    trace
        The execution trace to render.

    Returns
    -------
    str
        A complete Mermaid gantt chart definition.
    """
    lines: list[str] = [
        "gantt",
        "    title Execution Trace",
        "    dateFormat x",
        "    axisFormat %Lms",
    ]

    # Lay spans out sequentially based on durations (wall-clock offsets
    # are not available in the event data).
    offset_ms = 0

    for span in trace.spans:
        dur = span.duration_ms or 0
        dur_label = _format_gantt_duration(span.duration_ms)
        status_tag = "done," if span.status == "success" else "crit,"
        label = _sanitize_gantt_label(f"{span.node_id} ({span.node_type}) {dur_label}")

        lines.append(f"    section {span.node_id}")
        lines.append(f"    {label} :{status_tag} {int(offset_ms)}, {int(offset_ms + dur)}")

        # Sub-bars for tool calls
        sub_offset = offset_ms
        for tc in span.tool_calls:
            tc_dur = tc.duration_ms or 0
            tc_tag = "done," if tc.success else "crit,"
            tc_label = _sanitize_gantt_label(
                f"{tc.tool_name} {_format_gantt_duration(tc.duration_ms)}"
            )
            lines.append(f"    {tc_label} :{tc_tag} {int(sub_offset)}, {int(sub_offset + tc_dur)}")
            sub_offset += tc_dur

        # Sub-bars for LLM calls
        for lc in span.llm_calls:
            lc_dur = lc.duration_ms or 0
            tok_label = f"{lc.tokens_in}+{lc.tokens_out}tok"
            lc_label = _sanitize_gantt_label(
                f"{lc.model} ({tok_label}) {_format_gantt_duration(lc.duration_ms)}"
            )
            lines.append(f"    {lc_label} :done, {int(sub_offset)}, {int(sub_offset + lc_dur)}")
            sub_offset += lc_dur

        offset_ms += dur

    return "\n".join(lines) + "\n"


__all__ = [
    "DecisionPoint",
    "ExecutionTrace",
    "LLMCallSpan",
    "ToolCallSpan",
    "TraceCollector",
    "TraceSpan",
    "trace_to_mermaid",
]
