"""
Provider Circuit Breaker for the kvict Proxy.

Adapted from packages/caae/circuit_breaker.py. The original CAAE circuit breaker
monitors PCIe queue depth and switches between LWKCP (cost-aware) and LRU
(conservative) eviction policies.

The proxy circuit breaker monitors provider API latency/errors and switches
between "normal" (forward to provider) and "degraded" (serve stale cache)
modes. Same state machine: CLOSED -> OPEN -> HALF_OPEN -> CLOSED.

  CAAE analog:
    PCIe queue saturated  ->  circuit trips  ->  fall back to LRU (recompute)
  Proxy analog:
    Provider API degraded ->  circuit trips  ->  fall back to stale cache

This protects users from provider outages: if OpenAI is slow or returning
errors, the proxy serves cached responses instead of hanging.
"""

from __future__ import annotations

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

from proxy.config import ProxyConfig, get_config

logger = logging.getLogger(__name__)


class CircuitState(Enum):
    """Circuit breaker states (mirrors caae.circuit_breaker.CircuitState)."""

    CLOSED = "closed"  # Normal operation, forward to provider
    OPEN = "open"  # Provider degraded, serve stale cache
    HALF_OPEN = "half_open"  # Testing if provider has recovered


@dataclass
class ProviderCircuitStats:
    """Statistics for the provider circuit breaker."""

    total_checks: int = 0
    trips: int = 0
    current_state: CircuitState = CircuitState.CLOSED
    last_trip_time: Optional[float] = None
    last_recovery_time: Optional[float] = None
    stale_responses_served: int = 0


class ProviderCircuitBreaker:
    """
    Circuit breaker for LLM provider APIs.

    Monitors provider response latency and error rates. When the provider
    is degraded, the circuit "trips" and the proxy serves stale cached
    responses rather than timing out or returning errors.

    States (same as CAAE):
    - CLOSED: Normal operation, forward all requests to provider
    - OPEN: Provider degraded, serve stale cache when available
    - HALF_OPEN: Testing recovery by sending a probe request

    Adapted from packages/caae/circuit_breaker.py with the same
    threshold/cooldown/recovery logic.
    """

    def __init__(self, config: ProxyConfig | None = None):
        self._config = config or get_config()
        self._threshold_ms = self._config.circuit_breaker_threshold_ms
        self._cooldown_ms = self._config.circuit_breaker_cooldown_ms
        self._failure_threshold = self._config.circuit_breaker_failure_threshold
        self._recovery_threshold = self._config.circuit_breaker_recovery_threshold

        self._state = CircuitState.CLOSED
        self._last_trip_time: float = 0.0
        self._consecutive_successes = 0
        self._consecutive_failures = 0
        self._stats = ProviderCircuitStats()

    @property
    def state(self) -> CircuitState:
        return self._state

    @property
    def is_open(self) -> bool:
        """True if circuit is tripped (provider considered degraded)."""
        return self._state != CircuitState.CLOSED

    def record_success(self, latency_ms: float) -> None:
        """Record a successful provider response.

        If latency exceeds threshold, treat as a "slow success" (partial failure).
        """
        self._stats.total_checks += 1

        if latency_ms > self._threshold_ms:
            # Slow response — treat as failure
            self._record_failure()
            return

        # Fast success
        if self._state == CircuitState.CLOSED:
            self._consecutive_failures = 0
            self._consecutive_successes += 1

        elif self._state == CircuitState.HALF_OPEN:
            self._consecutive_successes += 1
            self._consecutive_failures = 0
            if self._consecutive_successes >= self._recovery_threshold:
                self._recover()

        elif self._state == CircuitState.OPEN:
            # Check cooldown
            self._check_cooldown()

    def record_failure(self) -> None:
        """Record a provider error (timeout, 5xx, connection refused, etc.)."""
        self._stats.total_checks += 1
        self._record_failure()

    def should_allow_request(self) -> bool:
        """Check whether to forward a request to the provider or serve stale cache.

        Returns:
            True if the request should be forwarded to the provider.
            False if the circuit is open and stale cache should be used.
        """
        if self._state == CircuitState.CLOSED:
            return True

        if self._state == CircuitState.OPEN:
            # Check if cooldown has elapsed
            elapsed_ms = (time.time() - self._last_trip_time) * 1000
            if elapsed_ms >= self._cooldown_ms:
                self._state = CircuitState.HALF_OPEN
                self._consecutive_successes = 0
                logger.info("Circuit breaker entering HALF_OPEN (testing recovery)")
                return True  # Allow one probe request
            return False  # Still in cooldown, serve stale

        # HALF_OPEN: allow the probe request
        return True

    def record_stale_served(self) -> None:
        """Record that a stale cached response was served due to circuit being open."""
        self._stats.stale_responses_served += 1

    def get_stats(self) -> ProviderCircuitStats:
        self._stats.current_state = self._state
        return self._stats

    # ---- internal ----

    def _record_failure(self) -> None:
        self._consecutive_failures += 1
        self._consecutive_successes = 0

        if self._state == CircuitState.CLOSED:
            if self._consecutive_failures >= self._failure_threshold:
                self._trip()

        elif self._state == CircuitState.HALF_OPEN:
            # Failed during recovery — re-trip
            self._trip()

        elif self._state == CircuitState.OPEN:
            pass  # Already open, nothing to do

    def _trip(self) -> None:
        self._state = CircuitState.OPEN
        self._last_trip_time = time.time()
        self._stats.trips += 1
        self._stats.current_state = CircuitState.OPEN
        self._stats.last_trip_time = self._last_trip_time
        logger.warning(
            "Circuit breaker TRIPPED (state=OPEN, trips=%d). "
            "Serving stale cache until provider recovers.",
            self._stats.trips,
        )

    def _recover(self) -> None:
        self._state = CircuitState.CLOSED
        self._consecutive_failures = 0
        self._stats.current_state = CircuitState.CLOSED
        self._stats.last_recovery_time = time.time()
        logger.info("Circuit breaker RECOVERED (state=CLOSED)")

    def _check_cooldown(self) -> None:
        elapsed_ms = (time.time() - self._last_trip_time) * 1000
        if elapsed_ms >= self._cooldown_ms:
            self._state = CircuitState.HALF_OPEN
            self._consecutive_successes = 0

    def force_open(self) -> None:
        """Force circuit open (for testing/emergency)."""
        self._trip()

    def force_close(self) -> None:
        """Force circuit closed (for testing/recovery)."""
        self._recover()
