"""Execution metrics middleware.

This middleware is *pure instrumentation*:
- It does not change execution semantics.
- It emits counters and timings to a sink.

The sink is intentionally backend-agnostic; higher layers decide whether to use an
in-memory sink (tests/CLI) or integrate with a real metrics system.
"""

from __future__ import annotations

from collections import Counter, defaultdict
from typing import Any, Generic, Protocol, TypeVar, runtime_checkable

from ..execution import BaseExecRequest, Executor
from .base import _call_timed_maybe_async

TReq = TypeVar("TReq", bound=BaseExecRequest)
TOut = TypeVar("TOut")


@runtime_checkable
class MetricsSink(Protocol):
    """Protocol-agnostic metrics sink for middleware."""

    def observe_ms(self, *, operation: str, ms: float, kind: str, meta: dict[str, Any] | None = None) -> None: ...
    def inc(self, *, operation: str, kind: str, meta: dict[str, Any] | None = None) -> None: ...


class NoOpMetricsSink:
    """Default sink used when metrics are disabled."""

    def observe_ms(self, *, operation: str, ms: float, kind: str, meta: dict[str, Any] | None = None) -> None:
        return None

    def inc(self, *, operation: str, kind: str, meta: dict[str, Any] | None = None) -> None:
        return None


class InMemoryMetricsSink:
    """Simple in-memory metrics sink for tests/debugging.

    Keys are stored as `<kind>:<operation>` to match the existing metrics printing style.
    """

    def __init__(self) -> None:
        self.counters: Counter[str] = Counter()
        self.timings_ms: dict[str, list[float]] = defaultdict(list)

    def observe_ms(self, *, operation: str, ms: float, kind: str, meta: dict[str, Any] | None = None) -> None:  # noqa: ARG002
        key = f"{kind}:{operation}"
        self.timings_ms[key].append(float(ms))

    def inc(self, *, operation: str, kind: str, meta: dict[str, Any] | None = None) -> None:  # noqa: ARG002
        key = f"{kind}:{operation}"
        self.counters[key] += 1


class MetricsMiddleware(Generic[TReq, TOut]):
    """Record attempt/success/error counters and timings.

    Emitted kinds follow a convention:
    - `<req.kind>.attempt`
    - `<req.kind>.success` / `<req.kind>.success_ms`
    - `<req.kind>.error` / `<req.kind>.error_ms`
    """

    def __init__(self, *, metrics: MetricsSink | None = None) -> None:
        self._metrics = metrics or NoOpMetricsSink()

    @property
    def metrics(self) -> MetricsSink:
        return self._metrics

    def wrap(self, nxt: Executor[TReq, TOut]) -> Executor[TReq, TOut]:
        m = self._metrics

        class _Wrapped:
            def execute(self, req: TReq) -> TOut:
                m.inc(operation=req.operation, kind=f"{req.kind}.attempt", meta=req.meta)

                def _on_success(ms: float) -> None:
                    m.observe_ms(operation=req.operation, ms=ms, kind=f"{req.kind}.success_ms", meta=req.meta)
                    m.inc(operation=req.operation, kind=f"{req.kind}.success", meta=req.meta)

                def _on_error(ms: float, _e: Exception) -> None:
                    m.observe_ms(operation=req.operation, ms=ms, kind=f"{req.kind}.error_ms", meta=req.meta)
                    m.inc(operation=req.operation, kind=f"{req.kind}.error", meta=req.meta)

                return _call_timed_maybe_async(
                    lambda: nxt.execute(req),
                    on_success=_on_success,
                    on_error=_on_error,
                )

        return _Wrapped()

