"""
cortex/llm/circuit_breaker.py — Provider quota/rate-limit circuit breaker.

Prevents LLM provider thrashing during quota exhaustion or sustained 429s.
Tracks per-provider consecutive failures and opens the breaker after a
configurable threshold. Auto-closes after a cooldown period.

Degraded mode (breaker open):
    - Read operations (retrieval) are ALWAYS allowed via non-LLM paths.
    - Write operations (remember/extract) are BLOCKED with a clear error.
      No ambiguous success — callers get an explicit CircuitBreakerOpenError.

Failure classification:
    - quota_exhausted — billing limit hit (hard stop, persistent)
    - rate_limited    — 429 / transient throttle (temporary, back off)
    - transient       — network/5xx (retryable, does NOT trip breaker)
    - model_error     — 400/404 / bad model name (config issue)

Only QUOTA_EXHAUSTED and RATE_LIMITED failures count toward the breaker
threshold. Transient and model errors use the existing retry_with_backoff
mechanism in provider.py and do not contribute to breaker trips.

Dependencies: stdlib only (no vendor SDK imports).
Used by: FallbackLLMProvider (optional wiring), CortexPlugin, tests.
"""
from __future__ import annotations

import enum
import logging
import time
from dataclasses import dataclass
from typing import Optional

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Failure classification
# ---------------------------------------------------------------------------


class FailureType(enum.Enum):
    """Classification of LLM provider failures."""

    QUOTA_EXHAUSTED = "quota_exhausted"
    """Billing/quota limit hit. Hard stop — likely persistent until billing resets."""

    RATE_LIMITED = "rate_limited"
    """HTTP 429 / throttle. Temporary — back off and retry after cooldown."""

    TRANSIENT = "transient"
    """Network error, 5xx, timeout. Retryable; does NOT trip circuit breaker."""

    MODEL_ERROR = "model_error"
    """HTTP 400/404, invalid model name. Config issue; does NOT trip breaker."""


def classify_error(exc: Exception) -> FailureType:
    """Classify a provider exception into a FailureType.

    Inspection order:
        1. HTTP status code (if present on exception as status_code/status/code)
        2. Error message keywords
        3. Exception type (ConnectionError, TimeoutError → TRANSIENT)
        4. Default: TRANSIENT (safe, non-tripping fallback)

    Args:
        exc: Any exception raised by an LLM provider.

    Returns:
        FailureType enum value.
    """
    msg = str(exc).lower()

    # Extract numeric status code from common exception attributes
    raw_code = (
        getattr(exc, "status_code", None)
        or getattr(exc, "status", None)
        or getattr(exc, "code", None)
    )
    code: Optional[int] = None
    if raw_code is not None:
        try:
            code = int(raw_code)
        except (ValueError, TypeError):
            code = None

    # --- 429: could be pure rate-limit or quota exhaustion ---
    is_429 = code == 429 or "429" in msg or "rate limit" in msg or "ratelimit" in msg
    if is_429:
        quota_kws = ("quota", "billing", "exceeded your", "usage limit", "insufficient_quota")
        if any(kw in msg for kw in quota_kws):
            return FailureType.QUOTA_EXHAUSTED
        return FailureType.RATE_LIMITED

    # --- Quota / billing keywords (not necessarily a 429) ---
    quota_kws = ("quota", "billing", "insufficient_quota", "limit exceeded", "usage limit")
    if any(kw in msg for kw in quota_kws):
        return FailureType.QUOTA_EXHAUSTED

    # --- Model config errors (non-retryable, don't trip breaker) ---
    model_kws = ("model not found", "model_not_found", "invalid model", "no such model")
    if code in (400, 404) or any(kw in msg for kw in model_kws):
        return FailureType.MODEL_ERROR

    # --- Transient: network / server errors ---
    if isinstance(exc, (ConnectionError, TimeoutError)):
        return FailureType.TRANSIENT
    if code in (500, 502, 503, 504):
        return FailureType.TRANSIENT
    transient_kws = ("timeout", "connection", "network", "server error", "unavailable")
    if any(kw in msg for kw in transient_kws):
        return FailureType.TRANSIENT

    # Default: treat unknown errors as transient (don't trip breaker)
    return FailureType.TRANSIENT


# ---------------------------------------------------------------------------
# Circuit breaker states
# ---------------------------------------------------------------------------


class BreakerState(enum.Enum):
    """Circuit breaker FSM states."""

    CLOSED = "closed"
    """Normal operation. All calls allowed."""

    OPEN = "open"
    """Tripped. Writes (remember/extract) are blocked. Reads still allowed."""

    HALF_OPEN = "half_open"
    """Cooldown elapsed. One probe allowed to test if provider recovered."""


# ---------------------------------------------------------------------------
# Per-provider state container
# ---------------------------------------------------------------------------


@dataclass
class _ProviderBreaker:
    """Mutable state for one provider's circuit breaker."""

    state: BreakerState = BreakerState.CLOSED
    consecutive_failures: int = 0
    opened_at: Optional[float] = None       # time.monotonic() when opened
    last_failure_type: Optional[FailureType] = None


# ---------------------------------------------------------------------------
# Public exception
# ---------------------------------------------------------------------------


class CircuitBreakerOpenError(Exception):
    """Raised when a write operation is attempted while the breaker is open.

    Always includes a human-readable message explaining:
    - which provider is blocked
    - why (failure type + count)
    - when auto-recovery is expected
    - that reads still work
    """


# ---------------------------------------------------------------------------
# CircuitBreaker
# ---------------------------------------------------------------------------


class CircuitBreaker:
    """
    Per-provider circuit breaker for LLM write operations.

    Tracks consecutive quota/rate-limit failures per provider. Opens the
    breaker after *threshold* consecutive failures and auto-closes after
    *cooldown_sec* seconds.

    Degraded mode (breaker open):
        - ``allow_write(provider)`` → False (writes blocked)
        - ``allow_read(provider)``  → True  (reads always work)
        - ``write_error(provider)`` → clear error message string
        - ``assert_write_allowed(provider)`` → raises CircuitBreakerOpenError

    Only QUOTA_EXHAUSTED and RATE_LIMITED failures trip the breaker.
    TRANSIENT and MODEL_ERROR failures are ignored by the breaker (they
    are handled by retry_with_backoff in provider.py).

    Thread safety: Not thread-safe. Designed for single-event-loop async use.
    If shared across threads, add external locking.

    Args:
        enabled: Master on/off toggle. When False, all checks are no-ops.
        threshold: Consecutive failures before opening (default 5).
        cooldown_sec: Seconds before auto-transitioning OPEN → HALF_OPEN (default 300).
    """

    def __init__(
        self,
        *,
        enabled: bool = True,
        threshold: int = 5,
        cooldown_sec: int = 300,
    ) -> None:
        self.enabled = enabled
        self.threshold = threshold
        self.cooldown_sec = cooldown_sec
        self._breakers: dict[str, _ProviderBreaker] = {}

    # ------------------------------------------------------------------ #
    # Internal helpers
    # ------------------------------------------------------------------ #

    def _get_breaker(self, provider: str) -> _ProviderBreaker:
        if provider not in self._breakers:
            self._breakers[provider] = _ProviderBreaker()
        return self._breakers[provider]

    def _check_auto_recovery(self, breaker: _ProviderBreaker) -> None:
        """Transition OPEN → HALF_OPEN if cooldown has elapsed."""
        if breaker.state == BreakerState.OPEN and breaker.opened_at is not None:
            elapsed = time.monotonic() - breaker.opened_at
            if elapsed >= self.cooldown_sec:
                logger.info(
                    "CircuitBreaker: cooldown elapsed (%.1fs >= %ds), "
                    "transitioning OPEN → HALF_OPEN",
                    elapsed,
                    self.cooldown_sec,
                )
                breaker.state = BreakerState.HALF_OPEN

    # ------------------------------------------------------------------ #
    # State mutation
    # ------------------------------------------------------------------ #

    def record_failure(self, provider: str, exc: Exception) -> None:
        """Record a provider failure. May open the breaker.

        Only QUOTA_EXHAUSTED and RATE_LIMITED failures count toward the
        threshold. TRANSIENT and MODEL_ERROR failures are ignored.

        Args:
            provider: Provider name (e.g., "openai", "anthropic").
            exc: The exception that was raised by the provider.
        """
        if not self.enabled:
            return

        breaker = self._get_breaker(provider)
        failure_type = classify_error(exc)
        breaker.last_failure_type = failure_type

        # Only quota/rate-limit failures trip the breaker
        trip_types = {FailureType.QUOTA_EXHAUSTED, FailureType.RATE_LIMITED}
        if failure_type not in trip_types:
            logger.debug(
                "CircuitBreaker[%s]: ignoring %s failure (not a trip type)",
                provider,
                failure_type.value,
            )
            return

        breaker.consecutive_failures += 1
        logger.debug(
            "CircuitBreaker[%s]: failure %d/%d (type=%s)",
            provider,
            breaker.consecutive_failures,
            self.threshold,
            failure_type.value,
        )

        if breaker.state in (BreakerState.CLOSED, BreakerState.HALF_OPEN):
            if breaker.consecutive_failures >= self.threshold:
                breaker.state = BreakerState.OPEN
                breaker.opened_at = time.monotonic()
                logger.warning(
                    "CircuitBreaker[%s]: OPEN after %d consecutive %s failures. "
                    "Writes blocked. Cooldown: %ds.",
                    provider,
                    breaker.consecutive_failures,
                    failure_type.value,
                    self.cooldown_sec,
                )

    def record_success(self, provider: str) -> None:
        """Record a successful provider call. Resets failure count.

        Always resets consecutive_failures and transitions any state → CLOSED.
        Call this after every successful LLM response.

        Args:
            provider: Provider name.
        """
        if not self.enabled:
            return
        breaker = self._get_breaker(provider)
        if breaker.state == BreakerState.HALF_OPEN:
            logger.info(
                "CircuitBreaker[%s]: success in HALF_OPEN → CLOSED (recovered)",
                provider,
            )
        elif breaker.state == BreakerState.OPEN:
            logger.info(
                "CircuitBreaker[%s]: success in OPEN → CLOSED (recovered without half-open)",
                provider,
            )
        breaker.state = BreakerState.CLOSED
        breaker.consecutive_failures = 0
        breaker.opened_at = None

    # ------------------------------------------------------------------ #
    # State queries
    # ------------------------------------------------------------------ #

    def is_open(self, provider: str) -> bool:
        """Return True if the breaker is OPEN (writes should be blocked).

        Automatically checks for cooldown expiry and transitions
        OPEN → HALF_OPEN if cooldown has elapsed.

        HALF_OPEN is treated as NOT open — it allows one probe write
        to test recovery.

        Args:
            provider: Provider name.

        Returns:
            True only when state is OPEN (not HALF_OPEN, not CLOSED).
        """
        if not self.enabled:
            return False
        breaker = self._get_breaker(provider)
        self._check_auto_recovery(breaker)
        return breaker.state == BreakerState.OPEN

    def get_state(self, provider: str) -> BreakerState:
        """Return the current BreakerState for *provider*.

        Triggers auto-recovery check (OPEN → HALF_OPEN) as a side effect.

        Args:
            provider: Provider name.

        Returns:
            Current BreakerState enum value.
        """
        breaker = self._get_breaker(provider)
        self._check_auto_recovery(breaker)
        return breaker.state

    def allow_write(self, provider: str) -> bool:
        """Return True if write operations (remember/extract) are allowed.

        Writes are blocked when the breaker is OPEN. Allowed when CLOSED
        or HALF_OPEN (probe attempt).

        Args:
            provider: Provider name.

        Returns:
            True when writes are permitted.
        """
        return not self.is_open(provider)

    def allow_read(self, provider: str) -> bool:  # noqa: ARG002
        """Return True if read operations (retrieval) are allowed.

        Reads are ALWAYS allowed regardless of breaker state — they use
        non-LLM paths (vector search, BM25) in degraded mode.

        Args:
            provider: Provider name (unused; reads are always permitted).

        Returns:
            Always True.
        """
        return True

    def assert_write_allowed(self, provider: str) -> None:
        """Raise CircuitBreakerOpenError if writes are blocked.

        Use this at the top of write operations (remember/extract) for an
        explicit, unambiguous failure rather than an ambiguous partial result.

        Args:
            provider: Provider name.

        Raises:
            CircuitBreakerOpenError: If the breaker is open for *provider*.
        """
        if not self.allow_write(provider):
            raise CircuitBreakerOpenError(self.write_error(provider))

    def write_error(self, provider: str) -> str:
        """Return a human-readable error message for a blocked write.

        Includes provider name, failure cause, remaining cooldown, and
        confirmation that reads still work.

        Args:
            provider: Provider name.

        Returns:
            Descriptive error string (no exception raised).
        """
        breaker = self._get_breaker(provider)
        failure_info = (
            f" (last failure: {breaker.last_failure_type.value})"
            if breaker.last_failure_type
            else ""
        )
        elapsed = (
            time.monotonic() - breaker.opened_at
            if breaker.opened_at is not None
            else 0.0
        )
        remaining = max(0.0, self.cooldown_sec - elapsed)
        return (
            f"LLM provider '{provider}' circuit breaker is OPEN{failure_info}. "
            f"Write operations (remember/extract) are blocked to prevent thrashing. "
            f"Read operations (retrieval) still work via non-LLM paths. "
            f"Auto-recovery in {remaining:.0f}s "
            f"({breaker.consecutive_failures} consecutive quota/rate-limit failures)."
        )

    # ------------------------------------------------------------------ #
    # Maintenance / observability
    # ------------------------------------------------------------------ #

    def reset(self, provider: str) -> None:
        """Force-reset a breaker to CLOSED state.

        Use for manual operator recovery or test teardown. Does NOT
        require cooldown to have elapsed.

        Args:
            provider: Provider name to reset.
        """
        if provider in self._breakers:
            del self._breakers[provider]
            logger.info("CircuitBreaker[%s]: manually reset to CLOSED", provider)

    def reset_all(self) -> None:
        """Reset all tracked provider breakers."""
        self._breakers.clear()
        logger.info("CircuitBreaker: all breakers reset")

    def status(self) -> dict[str, dict]:
        """Return a status snapshot for all tracked providers.

        Triggers auto-recovery checks as a side effect.

        Returns:
            Dict keyed by provider name with state, failure count,
            last_failure_type, and opened_at timestamp.
        """
        result: dict[str, dict] = {}
        for name, breaker in list(self._breakers.items()):
            self._check_auto_recovery(breaker)
            result[name] = {
                "state": breaker.state.value,
                "consecutive_failures": breaker.consecutive_failures,
                "last_failure_type": (
                    breaker.last_failure_type.value
                    if breaker.last_failure_type
                    else None
                ),
                "opened_at": breaker.opened_at,
            }
        return result
