"""
Circuit Breaker Pattern for groknroll.

Provides circuit breaker functionality for external service calls (LLM APIs, MCP, LSP)
to prevent cascading failures and enable graceful degradation.

States:
- CLOSED: Normal operation, requests pass through
- OPEN: Service is failing, requests fail fast
- HALF_OPEN: Testing if service has recovered

Example:
    breaker = CircuitBreaker(
        failure_threshold=5,
        recovery_timeout=30,
        expected_exception=LMHandlerError
    )

    @breaker
    async def call_llm_api(prompt):
        return await client.completion(prompt)
"""

import asyncio
import time
from enum import Enum
from functools import wraps
from typing import Any, Callable, Optional, Type


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

    CLOSED = "closed"  # Normal operation
    OPEN = "open"  # Failing, reject requests
    HALF_OPEN = "half_open"  # Testing recovery


class CircuitBreakerError(Exception):
    """Raised when circuit breaker is open."""

    def __init__(self, message: str, state: CircuitState, service_name: str):
        super().__init__(message)
        self.state = state
        self.service_name = service_name


class CircuitBreaker:
    """
    Circuit breaker for external service calls.

    Tracks failures and automatically opens the circuit when threshold is exceeded.
    After a recovery timeout, tries to recover by testing the service in HALF_OPEN state.
    """

    def __init__(
        self,
        failure_threshold: int = 5,
        recovery_timeout: float = 30.0,
        success_threshold: int = 2,
        expected_exception: Type[Exception] = Exception,
        service_name: str = "service",
        half_open_max_calls: int = 1,
    ):
        """
        Initialize circuit breaker.

        Args:
            failure_threshold: Number of failures before opening circuit
            recovery_timeout: Seconds to wait before attempting recovery
            success_threshold: Successes needed in HALF_OPEN to close circuit
            expected_exception: Exception type to catch and count as failure
            service_name: Name of the service for error messages
            half_open_max_calls: Max concurrent calls allowed in HALF_OPEN state
        """
        self.failure_threshold = failure_threshold
        self.recovery_timeout = recovery_timeout
        self.success_threshold = success_threshold
        self.expected_exception = expected_exception
        self.service_name = service_name
        self.half_open_max_calls = half_open_max_calls

        # State tracking
        self.state = CircuitState.CLOSED
        self.failure_count = 0
        self.success_count = 0
        self.last_failure_time: Optional[float] = None
        self.half_open_calls = 0

        # Thread safety
        self._lock = asyncio.Lock()

    async def call(self, func: Callable[..., Any], *args, **kwargs) -> Any:
        """
        Execute function with circuit breaker protection.

        Args:
            func: Function to call (can be sync or async)
            *args: Positional arguments to pass to func
            **kwargs: Keyword arguments to pass to func

        Returns:
            Function result

        Raises:
            CircuitBreakerError: If circuit is open
            Exception: If function raises unexpected exception
        """
        async with self._lock:
            # Check if circuit should transition from OPEN to HALF_OPEN
            if self.state == CircuitState.OPEN:
                if self._should_attempt_reset():
                    self.state = CircuitState.HALF_OPEN
                    self.half_open_calls = 0
                    self.success_count = 0
                else:
                    raise CircuitBreakerError(
                        f"Circuit breaker is OPEN for {self.service_name}. "
                        f"Retry after {self.recovery_timeout}s.",
                        state=self.state,
                        service_name=self.service_name,
                    )

            # In HALF_OPEN, limit concurrent calls
            if self.state == CircuitState.HALF_OPEN:
                if self.half_open_calls >= self.half_open_max_calls:
                    raise CircuitBreakerError(
                        f"Circuit breaker is HALF_OPEN for {self.service_name}. "
                        f"Testing recovery, please wait.",
                        state=self.state,
                        service_name=self.service_name,
                    )
                self.half_open_calls += 1

        # Execute the function
        try:
            # Handle both sync and async functions
            if asyncio.iscoroutinefunction(func):
                result = await func(*args, **kwargs)
            else:
                result = func(*args, **kwargs)

            # Success - update state
            async with self._lock:
                self._on_success()

            return result

        except self.expected_exception as e:
            # Expected failure - update state
            async with self._lock:
                self._on_failure()
            raise

        except Exception:
            # Unexpected exception - don't count as circuit failure
            async with self._lock:
                if self.state == CircuitState.HALF_OPEN:
                    self.half_open_calls -= 1
            raise

    def _should_attempt_reset(self) -> bool:
        """Check if enough time has passed to attempt recovery."""
        if self.last_failure_time is None:
            return True
        return time.time() - self.last_failure_time >= self.recovery_timeout

    def _on_success(self):
        """Handle successful call."""
        if self.state == CircuitState.HALF_OPEN:
            self.success_count += 1
            self.half_open_calls -= 1

            if self.success_count >= self.success_threshold:
                # Recovery successful - close circuit
                self.state = CircuitState.CLOSED
                self.failure_count = 0
                self.success_count = 0
        else:
            # In CLOSED state, reset failure count on success
            self.failure_count = 0

    def _on_failure(self):
        """Handle failed call."""
        self.failure_count += 1
        self.last_failure_time = time.time()

        if self.state == CircuitState.HALF_OPEN:
            # Recovery failed - reopen circuit
            self.state = CircuitState.OPEN
            self.success_count = 0
            self.half_open_calls = 0

        elif self.state == CircuitState.CLOSED:
            # Check if threshold exceeded
            if self.failure_count >= self.failure_threshold:
                self.state = CircuitState.OPEN

    def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]:
        """Decorator interface for circuit breaker."""

        @wraps(func)
        async def wrapper(*args, **kwargs):
            return await self.call(func, *args, **kwargs)

        return wrapper

    def reset(self):
        """Manually reset circuit breaker to CLOSED state."""
        self.state = CircuitState.CLOSED
        self.failure_count = 0
        self.success_count = 0
        self.last_failure_time = None
        self.half_open_calls = 0

    @property
    def is_open(self) -> bool:
        """Check if circuit is open."""
        return self.state == CircuitState.OPEN

    @property
    def is_closed(self) -> bool:
        """Check if circuit is closed."""
        return self.state == CircuitState.CLOSED

    def get_stats(self) -> dict[str, Any]:
        """Get circuit breaker statistics."""
        return {
            "state": self.state.value,
            "failure_count": self.failure_count,
            "success_count": self.success_count,
            "last_failure_time": self.last_failure_time,
            "service_name": self.service_name,
        }


# Global circuit breakers for common services
_circuit_breakers: dict[str, CircuitBreaker] = {}


def get_circuit_breaker(
    service_name: str,
    failure_threshold: int = 5,
    recovery_timeout: float = 30.0,
    **kwargs,
) -> CircuitBreaker:
    """
    Get or create a circuit breaker for a service.

    Args:
        service_name: Unique service identifier
        failure_threshold: Number of failures before opening circuit
        recovery_timeout: Seconds to wait before attempting recovery
        **kwargs: Additional CircuitBreaker parameters

    Returns:
        CircuitBreaker instance for the service
    """
    if service_name not in _circuit_breakers:
        _circuit_breakers[service_name] = CircuitBreaker(
            failure_threshold=failure_threshold,
            recovery_timeout=recovery_timeout,
            service_name=service_name,
            **kwargs,
        )
    return _circuit_breakers[service_name]


def reset_all_circuit_breakers():
    """Reset all circuit breakers to CLOSED state."""
    for breaker in _circuit_breakers.values():
        breaker.reset()
