"""Prometheus metrics collector for AlbusOS pathway execution.

Listens to the ``PathwayEvent`` stream and updates Prometheus counters
and histograms.  Requires the ``metrics`` optional dependency group::

    pip install albusos[metrics]

Usage::

    from core.observability.metrics import MetricsCollector

    collector = MetricsCollector()
    vm.add_listener(collector)
    # ... run pathway ...

Metrics exposed:

- ``albusos_pathway_duration_seconds`` — Histogram of pathway durations
- ``albusos_node_duration_seconds`` — Histogram of node durations
- ``albusos_tool_duration_seconds`` — Histogram of tool call durations
- ``albusos_tokens_total`` — Counter of tokens consumed
- ``albusos_tool_calls_total`` — Counter of tool invocations
- ``albusos_pathway_runs_total`` — Counter of pathway executions
"""

from __future__ import annotations

import contextlib
import logging
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Callable

    from core.pathways.models.execution import PathwayEvent

logger = logging.getLogger(__name__)

try:
    from prometheus_client import Counter, Histogram

    _HAS_PROMETHEUS = True
except ImportError:
    _HAS_PROMETHEUS = False


def _create_metrics() -> tuple[Histogram, Histogram, Histogram, Counter, Counter, Counter] | None:
    """Create Prometheus metric objects. Returns None if not available."""
    if not _HAS_PROMETHEUS:
        return None
    pathway_duration = Histogram(
        "albusos_pathway_duration_seconds",
        "Pathway execution duration",
        labelnames=["pathway_id", "status"],
    )
    node_duration = Histogram(
        "albusos_node_duration_seconds",
        "Node execution duration",
        labelnames=["pathway_id", "node_id", "node_type", "status"],
    )
    tool_duration = Histogram(
        "albusos_tool_duration_seconds",
        "Tool call duration",
        labelnames=["tool_name", "status"],
    )
    tokens_total = Counter(
        "albusos_tokens_total",
        "Total tokens consumed",
        labelnames=["pathway_id", "model"],
    )
    tool_calls_total = Counter(
        "albusos_tool_calls_total",
        "Total tool invocations",
        labelnames=["tool_name", "status"],
    )
    pathway_runs_total = Counter(
        "albusos_pathway_runs_total",
        "Total pathway executions",
        labelnames=["pathway_id", "status"],
    )
    return (
        pathway_duration,
        node_duration,
        tool_duration,
        tokens_total,
        tool_calls_total,
        pathway_runs_total,
    )


class MetricsCollector:
    """VMEventEmitter listener that updates Prometheus metrics.

    All event handlers are wrapped in ``contextlib.suppress`` so
    metrics failures never break execution.

    Parameters
    ----------
    prefix
        Metric name prefix. Defaults to ``"albusos"``.
    """

    def __init__(self) -> None:
        if not _HAS_PROMETHEUS:
            raise ImportError(
                "prometheus_client is required. Install with: pip install albusos[metrics]"
            )
        metrics = _create_metrics()
        if metrics is None:
            raise ImportError("prometheus_client is required")
        (
            self._pathway_duration,
            self._node_duration,
            self._tool_duration,
            self._tokens_total,
            self._tool_calls_total,
            self._pathway_runs_total,
        ) = metrics

        # Track pathway start times
        self._pathway_starts: dict[str, float] = {}

    # ------------------------------------------------------------------
    # Listener interface
    # ------------------------------------------------------------------

    def __call__(self, event: PathwayEvent) -> None:
        """Dispatch incoming pathway event to the appropriate handler."""
        event_type = getattr(event, "event_type", None)
        handler = _METRICS_DISPATCH.get(event_type)
        if handler:
            with contextlib.suppress(Exception):
                handler(self, event)

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

    def _on_execution_started(self, event: PathwayEvent) -> None:
        data = event.data or {}
        execution_id = str(data.get("execution_id", ""))
        if execution_id:
            import time

            self._pathway_starts[execution_id] = time.monotonic()

    def _on_execution_completed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        pathway_id = event.pathway_id or ""
        execution_id = str(data.get("execution_id", ""))
        duration_ms = data.get("duration_ms")

        if isinstance(duration_ms, (int, float)):
            self._pathway_duration.labels(pathway_id=pathway_id, status="success").observe(
                duration_ms / 1000
            )

        self._pathway_runs_total.labels(pathway_id=pathway_id, status="success").inc()
        self._pathway_starts.pop(execution_id, None)

    def _on_node_completed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        pathway_id = event.pathway_id or ""
        node_id = event.node_id or ""
        node_type = str(data.get("node_type", "unknown"))
        duration_ms = data.get("duration_ms")

        if isinstance(duration_ms, (int, float)):
            self._node_duration.labels(
                pathway_id=pathway_id,
                node_id=node_id,
                node_type=node_type,
                status="success",
            ).observe(duration_ms / 1000)

    def _on_node_failed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        pathway_id = event.pathway_id or ""
        node_id = event.node_id or ""
        node_type = str(data.get("node_type", "unknown"))
        duration_ms = data.get("duration_ms")

        if isinstance(duration_ms, (int, float)):
            self._node_duration.labels(
                pathway_id=pathway_id,
                node_id=node_id,
                node_type=node_type,
                status="failure",
            ).observe(duration_ms / 1000)

    def _on_tool_called(self, event: PathwayEvent) -> None:
        data = event.data or {}
        tool_name = str(data.get("tool_name", "unknown"))
        self._tool_calls_total.labels(tool_name=tool_name, status="started").inc()

    def _on_tool_completed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        tool_name = str(data.get("tool_name", "unknown"))
        duration_ms = data.get("duration_ms")

        self._tool_calls_total.labels(tool_name=tool_name, status="success").inc()
        if isinstance(duration_ms, (int, float)):
            self._tool_duration.labels(tool_name=tool_name, status="success").observe(
                duration_ms / 1000
            )

    def _on_tool_failed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        tool_name = str(data.get("tool_name", "unknown"))
        duration_ms = data.get("duration_ms")

        self._tool_calls_total.labels(tool_name=tool_name, status="failure").inc()
        if isinstance(duration_ms, (int, float)):
            self._tool_duration.labels(tool_name=tool_name, status="failure").observe(
                duration_ms / 1000
            )

    def _on_llm_call_completed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        pathway_id = event.pathway_id or ""
        model = str(data.get("model", "unknown"))
        tokens_in = data.get("tokens_in", 0)
        tokens_out = data.get("tokens_out", 0)

        total = 0
        if isinstance(tokens_in, (int, float)):
            total += int(tokens_in)
        if isinstance(tokens_out, (int, float)):
            total += int(tokens_out)

        if total > 0:
            self._tokens_total.labels(pathway_id=pathway_id, model=model).inc(total)


# Dispatch table
_METRICS_DISPATCH: dict[str, Callable[[MetricsCollector, PathwayEvent], None]] = {
    "execution_started": MetricsCollector._on_execution_started,
    "execution_completed": MetricsCollector._on_execution_completed,
    "node_completed": MetricsCollector._on_node_completed,
    "node_failed": MetricsCollector._on_node_failed,
    "tool_called": MetricsCollector._on_tool_called,
    "tool_completed": MetricsCollector._on_tool_completed,
    "tool_failed": MetricsCollector._on_tool_failed,
    "llm_call_completed": MetricsCollector._on_llm_call_completed,
}


__all__ = [
    "MetricsCollector",
]
