"""Circuit breaker pattern for provider resilience.

Wraps LLM providers with circuit breaker state machine:
CLOSED (normal) -> OPEN (failing) -> HALF_OPEN (probing) -> CLOSED.

Prevents cascading failures by skipping known-broken providers
until a configurable recovery window has elapsed. Also includes
exponential backoff with jitter for transient failure recovery.

Stdlib-only (ADR-021).
"""

from __future__ import annotations

import enum
import random
import time
from dataclasses import dataclass

from orxaq.orchestration.providers.base import LLMProvider, LLMResponse


class CircuitState(enum.Enum):
    """Circuit breaker states."""

    CLOSED = "closed"
    OPEN = "open"
    HALF_OPEN = "half_open"


@dataclass
class CircuitConfig:
    """Configuration for a circuit breaker."""

    failure_threshold: int = 5
    recovery_timeout_s: float = 30.0
    half_open_max_calls: int = 1
    success_threshold: int = 2

    def to_dict(self) -> dict:
        return {
            "failure_threshold": self.failure_threshold,
            "recovery_timeout_s": self.recovery_timeout_s,
            "half_open_max_calls": self.half_open_max_calls,
            "success_threshold": self.success_threshold,
        }


@dataclass
class CircuitStats:
    """Runtime statistics for a circuit breaker."""

    state: CircuitState = CircuitState.CLOSED
    consecutive_failures: int = 0
    consecutive_successes: int = 0
    total_calls: int = 0
    total_failures: int = 0
    total_successes: int = 0
    total_rejected: int = 0
    last_failure_time: float = 0.0
    last_success_time: float = 0.0
    opened_at: float = 0.0
    half_open_calls: int = 0

    def to_dict(self) -> dict:
        return {
            "state": self.state.value,
            "consecutive_failures": self.consecutive_failures,
            "consecutive_successes": self.consecutive_successes,
            "total_calls": self.total_calls,
            "total_failures": self.total_failures,
            "total_successes": self.total_successes,
            "total_rejected": self.total_rejected,
            "last_failure_time": self.last_failure_time,
            "last_success_time": self.last_success_time,
            "opened_at": self.opened_at,
        }


class CircuitBreaker:
    """Circuit breaker for a single provider.

    State transitions:
      CLOSED -> OPEN:      when consecutive failures >= threshold
      OPEN -> HALF_OPEN:   after recovery_timeout_s elapsed
      HALF_OPEN -> CLOSED: after success_threshold consecutive successes
      HALF_OPEN -> OPEN:   on any failure
    """

    def __init__(self, config: CircuitConfig | None = None) -> None:
        self._config = config or CircuitConfig()
        self._stats = CircuitStats()

    @property
    def state(self) -> CircuitState:
        # Auto-transition from OPEN to HALF_OPEN on timeout
        if self._stats.state == CircuitState.OPEN:
            elapsed = time.monotonic() - self._stats.opened_at
            if elapsed >= self._config.recovery_timeout_s:
                self._stats.state = CircuitState.HALF_OPEN
                self._stats.half_open_calls = 0
                self._stats.consecutive_successes = 0
        return self._stats.state

    @property
    def stats(self) -> CircuitStats:
        # Trigger state refresh
        _ = self.state
        return self._stats

    def allow_request(self) -> bool:
        """Check if a request is allowed through the circuit."""
        current = self.state
        if current == CircuitState.CLOSED:
            return True
        if current == CircuitState.OPEN:
            return False
        # HALF_OPEN: allow limited calls
        return self._stats.half_open_calls < self._config.half_open_max_calls

    def record_success(self) -> None:
        """Record a successful call."""
        self._stats.total_calls += 1
        self._stats.total_successes += 1
        self._stats.consecutive_successes += 1
        self._stats.consecutive_failures = 0
        self._stats.last_success_time = time.monotonic()

        if self._stats.state == CircuitState.HALF_OPEN:
            self._stats.half_open_calls += 1
            if self._stats.consecutive_successes >= self._config.success_threshold:
                self._stats.state = CircuitState.CLOSED

    def record_failure(self) -> None:
        """Record a failed call."""
        self._stats.total_calls += 1
        self._stats.total_failures += 1
        self._stats.consecutive_failures += 1
        self._stats.consecutive_successes = 0
        self._stats.last_failure_time = time.monotonic()

        if self._stats.state == CircuitState.HALF_OPEN:
            # Immediately re-open
            self._stats.state = CircuitState.OPEN
            self._stats.opened_at = time.monotonic()
        elif self._stats.consecutive_failures >= self._config.failure_threshold:
            self._stats.state = CircuitState.OPEN
            self._stats.opened_at = time.monotonic()

    def record_rejected(self) -> None:
        """Record a rejected call (circuit open)."""
        self._stats.total_rejected += 1

    def reset(self) -> None:
        """Reset the circuit breaker to initial state."""
        self._stats = CircuitStats()


@dataclass
class BackoffConfig:
    """Configuration for exponential backoff with jitter."""

    base_delay_s: float = 1.0
    max_delay_s: float = 30.0
    multiplier: float = 2.0
    jitter_factor: float = 0.5
    max_retries: int = 3

    def to_dict(self) -> dict:
        return {
            "base_delay_s": self.base_delay_s,
            "max_delay_s": self.max_delay_s,
            "multiplier": self.multiplier,
            "jitter_factor": self.jitter_factor,
            "max_retries": self.max_retries,
        }


def compute_backoff(
    attempt: int,
    config: BackoffConfig | None = None,
) -> float:
    """Compute delay for attempt N with exponential backoff + jitter."""
    c = config or BackoffConfig()
    if attempt <= 0:
        return 0.0
    delay = c.base_delay_s * (c.multiplier ** (attempt - 1))
    delay = min(delay, c.max_delay_s)
    jitter = delay * c.jitter_factor * random.random()  # noqa: S311
    return delay + jitter


class ResilientProvider:
    """Wraps an LLMProvider with circuit breaker + backoff.

    Usage::

        provider = VLLMProvider(...)
        resilient = ResilientProvider(
            provider,
            circuit_config=CircuitConfig(failure_threshold=3),
            backoff_config=BackoffConfig(max_retries=2),
        )
        response = resilient.generate("Hello")
    """

    def __init__(
        self,
        provider: LLMProvider,
        circuit_config: CircuitConfig | None = None,
        backoff_config: BackoffConfig | None = None,
    ) -> None:
        self._provider = provider
        self._circuit = CircuitBreaker(circuit_config)
        self._backoff_config = backoff_config or BackoffConfig()

    @property
    def name(self) -> str:
        return self._provider.name

    @property
    def models(self) -> list[str]:
        return self._provider.models

    @property
    def circuit(self) -> CircuitBreaker:
        return self._circuit

    @property
    def inner(self) -> LLMProvider:
        return self._provider

    def generate(
        self,
        prompt: str,
        model: str | None = None,
        max_tokens: int = 1024,
        temperature: float = 0.7,
    ) -> LLMResponse:
        """Generate with circuit breaker protection and retry logic."""
        if not self._circuit.allow_request():
            self._circuit.record_rejected()
            return LLMResponse(
                text="",
                model=model or "none",
                provider=self.name,
                metadata={
                    "error": "circuit open",
                    "circuit_state": self._circuit.state.value,
                },
            )

        last_error = ""
        for attempt in range(self._backoff_config.max_retries + 1):
            if attempt > 0:
                delay = compute_backoff(attempt, self._backoff_config)
                time.sleep(delay)

            response = self._provider.generate(
                prompt,
                model=model,
                max_tokens=max_tokens,
                temperature=temperature,
            )

            error = response.metadata.get("error")
            if not error:
                self._circuit.record_success()
                return response

            last_error = str(error)

        # All retries exhausted
        self._circuit.record_failure()
        return LLMResponse(
            text="",
            model=model or "none",
            provider=self.name,
            metadata={
                "error": last_error,
                "retries_exhausted": True,
                "attempts": self._backoff_config.max_retries + 1,
            },
        )

    def health_check(self) -> dict:
        """Health check including circuit state."""
        base = self._provider.health_check()
        base["circuit_state"] = self._circuit.state.value
        base["circuit_stats"] = self._circuit.stats.to_dict()
        return base


class ResilientRouter:
    """Router that wraps providers with circuit breakers.

    Drop-in replacement for ProviderRouter that adds resilience.
    """

    def __init__(
        self,
        providers: list[LLMProvider],
        circuit_config: CircuitConfig | None = None,
        backoff_config: BackoffConfig | None = None,
    ) -> None:
        cc = circuit_config or CircuitConfig()
        bc = backoff_config or BackoffConfig()
        self._resilient: list[ResilientProvider] = [ResilientProvider(p, cc, bc) for p in providers]

    def route(
        self,
        prompt: str,
        model: str | None = None,
        max_tokens: int = 1024,
        temperature: float = 0.7,
    ) -> LLMResponse:
        """Route to first available provider with circuit protection."""
        errors: list[str] = []
        for rp in self._resilient:
            if not rp.circuit.allow_request():
                errors.append(f"{rp.name}: circuit open")
                rp.circuit.record_rejected()
                continue

            response = rp.generate(
                prompt,
                model=model,
                max_tokens=max_tokens,
                temperature=temperature,
            )
            if not response.metadata.get("error"):
                return response
            errors.append(f"{rp.name}: {response.metadata.get('error')}")

        return LLMResponse(
            text="",
            model=model or "none",
            provider="resilient_router",
            metadata={
                "error": "all providers failed or circuit open",
                "details": errors,
            },
        )

    def stats(self) -> dict:
        """Return stats for all wrapped providers."""
        return {rp.name: rp.circuit.stats.to_dict() for rp in self._resilient}

    def reset_all(self) -> None:
        """Reset all circuit breakers."""
        for rp in self._resilient:
            rp.circuit.reset()

    @property
    def providers(self) -> list[ResilientProvider]:
        return list(self._resilient)


__all__ = [
    "BackoffConfig",
    "CircuitBreaker",
    "CircuitConfig",
    "CircuitState",
    "CircuitStats",
    "ResilientProvider",
    "ResilientRouter",
    "compute_backoff",
]
