"""Circuit breaker for external service calls."""
from __future__ import annotations

import logging
import time
from dataclasses import dataclass
from typing import Any, TypeVar

_logger = logging.getLogger(__name__)

T = TypeVar("T")


class CircuitOpenError(Exception):
    """Raised when a call is rejected because the circuit is open."""

    def __init__(self, name: str, retry_after: float) -> None:
        self.name = name
        self.retry_after = retry_after
        super().__init__(
            f"Circuit '{name}' is open — retry after {retry_after:.1f}s"
        )


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

    failure_threshold: int = 5
    reset_timeout_seconds: float = 30.0
    half_open_max_calls: int = 1
    excluded_exceptions: tuple[type[BaseException], ...] = ()


class CircuitBreaker:
    """Three-state circuit breaker: closed → open → half-open → closed.

    Usage::

        breaker = CircuitBreaker("mcp-server")

        async def call_mcp():
            async with breaker:
                return await mcp_client.call(...)

    Or wrap a callable directly::

        safe_call = breaker.wrap(mcp_client.call)
        result = await safe_call(...)
    """

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

    def __init__(
        self,
        name: str,
        config: CircuitBreakerConfig | None = None,
    ) -> None:
        self.name = name
        self._config = config or CircuitBreakerConfig()
        self._state = self.CLOSED
        self._failure_count = 0
        self._success_count = 0
        self._last_failure_time = 0.0
        self._half_open_calls = 0

    @property
    def state(self) -> str:
        """Return current state, auto-transitioning open→half_open on timeout."""
        if self._state == self.OPEN:
            elapsed = time.monotonic() - self._last_failure_time
            if elapsed >= self._config.reset_timeout_seconds:
                self._state = self.HALF_OPEN
                self._half_open_calls = 0
                _logger.info("Circuit '%s' → half_open (timeout elapsed)", self.name)
        return self._state

    @property
    def failure_count(self) -> int:
        return self._failure_count

    def _on_success(self) -> None:
        if self._state == self.HALF_OPEN:
            self._success_count += 1
            if self._success_count >= self._config.half_open_max_calls:
                self._state = self.CLOSED
                self._failure_count = 0
                self._success_count = 0
                _logger.info("Circuit '%s' → closed (recovery confirmed)", self.name)
        elif self._state == self.CLOSED:
            self._failure_count = 0

    def _on_failure(self, exc: BaseException) -> None:
        # Don't count excluded exceptions (e.g. validation errors)
        if isinstance(exc, self._config.excluded_exceptions):
            return
        self._failure_count += 1
        self._last_failure_time = time.monotonic()
        if self._state == self.HALF_OPEN:
            self._state = self.OPEN
            _logger.warning("Circuit '%s' → open (half-open probe failed)", self.name)
        elif (
            self._state == self.CLOSED
            and self._failure_count >= self._config.failure_threshold
        ):
            self._state = self.OPEN
            _logger.warning(
                "Circuit '%s' → open (threshold %d reached)",
                self.name, self._config.failure_threshold,
            )

    def _check_allowed(self) -> None:
        state = self.state  # triggers auto-transition
        if state == self.OPEN:
            retry_after = self._config.reset_timeout_seconds - (
                time.monotonic() - self._last_failure_time
            )
            raise CircuitOpenError(self.name, max(0.0, retry_after))
        if state == self.HALF_OPEN:
            if self._half_open_calls >= self._config.half_open_max_calls:
                raise CircuitOpenError(self.name, self._config.reset_timeout_seconds)
            self._half_open_calls += 1

    async def __aenter__(self) -> CircuitBreaker:
        self._check_allowed()
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        if exc_val is None:
            self._on_success()
        else:
            self._on_failure(exc_val)

    def __enter__(self) -> CircuitBreaker:
        self._check_allowed()
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        if exc_val is None:
            self._on_success()
        else:
            self._on_failure(exc_val)

    def reset(self) -> None:
        """Manually reset the circuit to closed state."""
        self._state = self.CLOSED
        self._failure_count = 0
        self._success_count = 0
        self._half_open_calls = 0


__all__ = [
    "CircuitBreaker",
    "CircuitBreakerConfig",
    "CircuitOpenError",
]
