"""Circuit breakers for DE outbound calls (Milvus, SeaweedFS, LLM).

Implements SYS-NFR-20: 5 failures (consecutive) -> open for 60s.
State transitions emit a Prometheus Gauge and a warning log.

The Gauge registers against the default ``prometheus_client.REGISTRY`` so the
existing ``/metrics`` endpoint in ``airflow/plugins/api/metrics.py`` scrapes it
without any changes there. Importing this module from the airflow plugin tree
is fine — it has no heavy side effects.
"""

from __future__ import annotations

import logging
from collections.abc import Awaitable, Callable
from typing import Any, TypeVar

import pybreaker
from prometheus_client import REGISTRY, Gauge

log = logging.getLogger(__name__)

T = TypeVar("T")


# ---------------------------------------------------------------------------
# Prometheus gauge — co-located with the listener that updates it
# ---------------------------------------------------------------------------

# 0 = closed (healthy), 1 = half_open (probing), 2 = open (tripped).
_STATE_VALUE: dict[str, int] = {
    "closed": 0,
    "half_open": 1,
    "open": 2,
}


def _gauge(name: str, doc: str, labels: list[str]) -> Gauge:
    existing = REGISTRY._names_to_collectors.get(name)
    if existing is not None:
        return existing
    return Gauge(name, doc, labels)


circuit_breaker_state = _gauge(
    "circuit_breaker_state",
    "DE outbound circuit breaker state (0=closed, 1=half_open, 2=open)",
    ["upstream"],
)


# ---------------------------------------------------------------------------
# Listener
# ---------------------------------------------------------------------------


def _state_label(state: Any) -> str:
    """Return a metric-friendly state name (pybreaker uses 'half-open' with a hyphen)."""
    if state is None:
        return "unknown"
    name = getattr(state, "name", None) or str(state)
    return name.replace("-", "_")


class _MetricsListener(pybreaker.CircuitBreakerListener):
    """Update ``circuit_breaker_state`` gauge and log on every transition."""

    def state_change(self, cb, old_state, new_state):  # noqa: D401
        old_name = _state_label(old_state)
        new_name = _state_label(new_state)
        value = _STATE_VALUE.get(new_name, 0)
        circuit_breaker_state.labels(upstream=cb.name).set(value)
        log.warning(
            "Circuit breaker '%s' transitioned: %s -> %s (fail_max=%s, reset_timeout=%ss)",
            cb.name, old_name, new_name, cb.fail_max, cb.reset_timeout,
        )


_listener = _MetricsListener()


# ---------------------------------------------------------------------------
# Breaker singletons (per upstream)
# ---------------------------------------------------------------------------

# SYS-NFR-20: 5 failures -> open 60s. pybreaker counts consecutive failures by
# default; that's accepted in spec-textiq-394 (sliding-window deferred).
_FAIL_MAX = 5
_RESET_TIMEOUT = 60


def _make_breaker(name: str) -> pybreaker.CircuitBreaker:
    # throw_new_error_on_trip=False: on the trip call (5th consecutive failure)
    # the caller sees the original upstream exception, not CircuitBreakerError.
    # Subsequent short-circuited calls still raise CircuitBreakerError. This
    # matches spec-textiq-394's I/O matrix ("Caller sees raw exception on this call").
    breaker = pybreaker.CircuitBreaker(
        fail_max=_FAIL_MAX,
        reset_timeout=_RESET_TIMEOUT,
        listeners=[_listener],
        name=name,
        throw_new_error_on_trip=False,
    )
    # Seed the gauge with the initial closed state so /metrics shows the label
    # set immediately, before any failure occurs.
    circuit_breaker_state.labels(upstream=name).set(_STATE_VALUE["closed"])
    return breaker


milvus_breaker = _make_breaker("milvus")
s3_breaker = _make_breaker("s3")
llm_breaker = _make_breaker("llm")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def breaker_call(
    breaker: pybreaker.CircuitBreaker,
    fn: Callable[..., T],
    *args: Any,
    **kwargs: Any,
) -> T:
    """Invoke a sync callable through the breaker. Re-raises original exceptions."""
    return breaker.call(fn, *args, **kwargs)


async def async_breaker_call(
    breaker: pybreaker.CircuitBreaker,
    fn: Callable[..., Awaitable[T]],
    *args: Any,
    **kwargs: Any,
) -> T:
    """Invoke an async callable through the breaker.

    Manual reimplementation of ``CircuitBreaker.call`` for coroutines:
    pybreaker 1.4.1's ``call_async`` uses ``@gen.coroutine`` without importing
    tornado's ``gen`` module, so it raises ``NameError`` on first invocation.
    We reproduce ``State.call`` semantics by driving ``before_call``,
    ``_handle_error``, and ``_handle_success`` ourselves around an ``await``.
    """
    with breaker._lock:
        state = breaker.state  # may transition open -> half_open if reset_timeout passed
        state.before_call(fn, *args, **kwargs)  # raises CircuitBreakerError if open
        for listener in breaker.listeners:
            listener.before_call(breaker, fn, *args, **kwargs)

    try:
        result = await fn(*args, **kwargs)
    except BaseException as exc:
        with breaker._lock:
            # _handle_error increments the failure counter, transitions to open
            # at the threshold, and re-raises the original exception.
            breaker.state._handle_error(exc)
        raise  # unreachable: _handle_error reraises, but keeps type-checkers happy
    else:
        with breaker._lock:
            breaker.state._handle_success()
        return result


__all__ = [
    "circuit_breaker_state",
    "milvus_breaker",
    "s3_breaker",
    "llm_breaker",
    "breaker_call",
    "async_breaker_call",
]