"""
Circuit Breaker - Fault tolerance pattern for tool execution.

Prevents cascading failures by temporarily blocking calls to failing tools,
allowing them time to recover before retrying.
"""
import asyncio
import logging
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Optional


logger = logging.getLogger(__name__)


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

    CLOSED = "closed"  # Normal operation, requests flow through
    OPEN = "open"  # Circuit tripped, requests are rejected
    HALF_OPEN = "half_open"  # Testing if service has recovered


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

    Attributes:
        failure_threshold: Number of failures before opening circuit (default: 5)
        recovery_timeout_ms: Time to wait before testing recovery (default: 30000ms)
        success_threshold: Successes needed in half-open to close circuit (default: 2)
        half_open_max_calls: Max concurrent calls allowed in half-open state (default: 1)
    """

    failure_threshold: int = 5
    recovery_timeout_ms: int = 30000
    success_threshold: int = 2
    half_open_max_calls: int = 1


@dataclass
class CircuitStats:
    """Statistics for a circuit breaker."""

    state: CircuitState = CircuitState.CLOSED
    failure_count: int = 0
    success_count: int = 0
    last_failure_time: Optional[float] = None
    last_state_change: Optional[float] = None
    total_failures: int = 0
    total_successes: int = 0
    total_rejections: int = 0
    half_open_calls: int = 0


class CircuitOpenError(Exception):
    """Raised when circuit is open and request is rejected."""

    def __init__(
        self,
        tool_name: str,
        time_until_retry: float,
        failure_count: int,
    ):
        self.tool_name = tool_name
        self.time_until_retry = time_until_retry
        self.failure_count = failure_count
        super().__init__(
            f"Circuit open for '{tool_name}': {failure_count} failures. "
            f"Retry in {time_until_retry:.1f}s"
        )


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

    States:
    - CLOSED: Normal operation. Failures increment counter.
    - OPEN: After failure_threshold, rejects all calls for recovery_timeout_ms.
    - HALF_OPEN: After timeout, allows limited calls to test recovery.

    Example:
        ```python
        breaker = CircuitBreaker("send_email", config)

        # Check before calling
        if not breaker.allow_request():
            raise CircuitOpenError(...)

        try:
            result = await tool.handler(args)
            breaker.record_success()
        except Exception as e:
            breaker.record_failure()
            raise
        ```
    """

    def __init__(
        self,
        tool_name: str,
        config: Optional[CircuitBreakerConfig] = None,
    ):
        self.tool_name = tool_name
        self.config = config or CircuitBreakerConfig()
        self._stats = CircuitStats()
        self._lock = asyncio.Lock()

    @property
    def state(self) -> CircuitState:
        """Current circuit state."""
        return self._stats.state

    @property
    def stats(self) -> CircuitStats:
        """Get circuit statistics (read-only snapshot)."""
        return CircuitStats(
            state=self._stats.state,
            failure_count=self._stats.failure_count,
            success_count=self._stats.success_count,
            last_failure_time=self._stats.last_failure_time,
            last_state_change=self._stats.last_state_change,
            total_failures=self._stats.total_failures,
            total_successes=self._stats.total_successes,
            total_rejections=self._stats.total_rejections,
            half_open_calls=self._stats.half_open_calls,
        )

    def allow_request(self) -> bool:
        """
        Check if a request should be allowed through.

        Returns:
            True if request can proceed, False if circuit is open
        """
        now = time.time()

        if self._stats.state == CircuitState.CLOSED:
            return True

        elif self._stats.state == CircuitState.OPEN:
            # Check if recovery timeout has elapsed
            if self._stats.last_failure_time is not None:
                elapsed_ms = (now - self._stats.last_failure_time) * 1000
                if elapsed_ms >= self.config.recovery_timeout_ms:
                    # Transition to half-open
                    self._transition_to(CircuitState.HALF_OPEN)
                    # This call counts against the half-open limit
                    self._stats.half_open_calls += 1
                    return True
            return False

        elif self._stats.state == CircuitState.HALF_OPEN:
            # Allow limited calls in half-open state
            if self._stats.half_open_calls < self.config.half_open_max_calls:
                self._stats.half_open_calls += 1
                return True
            return False

        return False

    def record_success(self) -> None:
        """Record a successful call."""
        self._stats.total_successes += 1

        if self._stats.state == CircuitState.HALF_OPEN:
            self._stats.success_count += 1
            self._stats.half_open_calls = max(0, self._stats.half_open_calls - 1)

            if self._stats.success_count >= self.config.success_threshold:
                # Service recovered, close circuit
                self._transition_to(CircuitState.CLOSED)
                logger.info(f"Circuit closed for '{self.tool_name}': service recovered")

        elif self._stats.state == CircuitState.CLOSED:
            # Reset failure count on success
            self._stats.failure_count = 0

    def record_failure(self) -> None:
        """Record a failed call."""
        now = time.time()
        self._stats.total_failures += 1
        self._stats.last_failure_time = now

        if self._stats.state == CircuitState.HALF_OPEN:
            # Failed during recovery test, reopen circuit
            self._stats.half_open_calls = max(0, self._stats.half_open_calls - 1)
            self._transition_to(CircuitState.OPEN)
            logger.warning(
                f"Circuit reopened for '{self.tool_name}': recovery test failed"
            )

        elif self._stats.state == CircuitState.CLOSED:
            self._stats.failure_count += 1

            if self._stats.failure_count >= self.config.failure_threshold:
                # Threshold reached, open circuit
                self._transition_to(CircuitState.OPEN)
                logger.warning(
                    f"Circuit opened for '{self.tool_name}': "
                    f"{self._stats.failure_count} consecutive failures"
                )

    def record_rejection(self) -> None:
        """Record a rejected request (circuit was open)."""
        self._stats.total_rejections += 1

    def time_until_retry(self) -> float:
        """
        Get seconds until circuit might allow requests again.

        Returns:
            Seconds until retry, or 0 if circuit is not open
        """
        if self._stats.state != CircuitState.OPEN:
            return 0.0

        if self._stats.last_failure_time is None:
            return 0.0

        elapsed_ms = (time.time() - self._stats.last_failure_time) * 1000
        remaining_ms = self.config.recovery_timeout_ms - elapsed_ms
        return max(0.0, remaining_ms / 1000)

    def reset(self) -> None:
        """Reset circuit to closed state."""
        self._transition_to(CircuitState.CLOSED)
        self._stats.failure_count = 0
        self._stats.success_count = 0
        self._stats.half_open_calls = 0
        logger.info(f"Circuit reset for '{self.tool_name}'")

    def _transition_to(self, new_state: CircuitState) -> None:
        """Transition to a new state."""
        old_state = self._stats.state
        self._stats.state = new_state
        self._stats.last_state_change = time.time()

        if new_state == CircuitState.CLOSED:
            self._stats.failure_count = 0
            self._stats.success_count = 0
            self._stats.half_open_calls = 0
        elif new_state == CircuitState.HALF_OPEN:
            self._stats.success_count = 0
            self._stats.half_open_calls = 0

        logger.debug(
            f"Circuit '{self.tool_name}' state: {old_state.value} -> {new_state.value}"
        )


class CircuitBreakerRegistry:
    """
    Registry managing circuit breakers for multiple tools.

    Example:
        ```python
        registry = CircuitBreakerRegistry(
            default_config=CircuitBreakerConfig(failure_threshold=3)
        )

        # Get or create circuit breaker for a tool
        breaker = registry.get("send_email")

        # Configure specific tool differently
        registry.configure("critical_api", CircuitBreakerConfig(
            failure_threshold=10,
            recovery_timeout_ms=60000,
        ))
        ```
    """

    def __init__(
        self,
        default_config: Optional[CircuitBreakerConfig] = None,
        enabled: bool = True,
    ):
        self._default_config = default_config or CircuitBreakerConfig()
        self._breakers: dict[str, CircuitBreaker] = {}
        self._tool_configs: dict[str, CircuitBreakerConfig] = {}
        self._enabled = enabled

    @property
    def enabled(self) -> bool:
        """Check if circuit breakers are enabled."""
        return self._enabled

    @enabled.setter
    def enabled(self, value: bool) -> None:
        """Enable or disable circuit breakers."""
        self._enabled = value

    def configure(
        self,
        tool_name: str,
        config: CircuitBreakerConfig,
    ) -> None:
        """
        Configure circuit breaker for a specific tool.

        Args:
            tool_name: Tool identifier
            config: Circuit breaker configuration
        """
        self._tool_configs[tool_name] = config

        # Update existing breaker if present
        if tool_name in self._breakers:
            self._breakers[tool_name].config = config

    def get(self, tool_name: str) -> CircuitBreaker:
        """
        Get or create circuit breaker for a tool.

        Args:
            tool_name: Tool identifier

        Returns:
            CircuitBreaker instance
        """
        if tool_name not in self._breakers:
            config = self._tool_configs.get(tool_name, self._default_config)
            self._breakers[tool_name] = CircuitBreaker(tool_name, config)

        return self._breakers[tool_name]

    def get_all_stats(self) -> dict[str, CircuitStats]:
        """
        Get statistics for all circuit breakers.

        Returns:
            Dict mapping tool names to their circuit stats
        """
        return {name: breaker.stats for name, breaker in self._breakers.items()}

    def reset(self, tool_name: str) -> None:
        """Reset circuit breaker for a specific tool."""
        if tool_name in self._breakers:
            self._breakers[tool_name].reset()

    def reset_all(self) -> None:
        """Reset all circuit breakers."""
        for breaker in self._breakers.values():
            breaker.reset()
