"""Connection management with automatic retry, recovery, and circuit breaker patterns.

Features:
- Exponential backoff retry strategy
- Circuit breaker for failing hosts
- Automatic reconnection on connection loss
- Connection health monitoring
- Graceful degradation with informative errors
"""

import asyncio
import logging
import time
from dataclasses import dataclass
from enum import Enum
from typing import Optional, Dict, Tuple

logger = logging.getLogger(__name__)


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

    CLOSED = "closed"  # Operating normally
    OPEN = "open"  # Failing, reject requests
    HALF_OPEN = "half_open"  # Testing recovery


@dataclass
class RetryConfig:
    """Configuration for retry strategy."""

    max_attempts: int = 3  # Total attempts (including initial)
    initial_delay_s: float = 0.1  # Initial backoff delay
    max_delay_s: float = 30.0  # Maximum backoff delay
    exponential_base: float = 2.0  # Backoff multiplier
    jitter: float = 0.1  # Random jitter (fraction of delay)


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

    failure_threshold: int = 5  # Failures before opening circuit
    recovery_timeout_s: float = 60.0  # Time before trying again
    success_threshold: int = 2  # Successes needed to close circuit


class CircuitBreaker:
    """Circuit breaker for managing failing connections."""

    def __init__(self, host: str, config: Optional[CircuitBreakerConfig] = None):
        """
        Initialize circuit breaker.

        Args:
            host: Host name/address
            config: Circuit breaker configuration
        """
        self.host = host
        self.config = config or CircuitBreakerConfig()

        self.state = CircuitState.CLOSED
        self.failure_count = 0
        self.success_count = 0
        self.last_failure_time: Optional[float] = None
        self.last_state_change_time = time.time()

    def record_success(self) -> None:
        """Record a successful operation."""
        self.failure_count = 0

        if self.state == CircuitState.HALF_OPEN:
            self.success_count += 1
            if self.success_count >= self.config.success_threshold:
                self._close_circuit()
        elif self.state == CircuitState.CLOSED:
            self.success_count = 0

    def record_failure(self) -> None:
        """Record a failed operation."""
        self.failure_count += 1
        self.last_failure_time = time.time()

        if self.state == CircuitState.CLOSED:
            if self.failure_count >= self.config.failure_threshold:
                self._open_circuit()

        elif self.state == CircuitState.HALF_OPEN:
            self._open_circuit()  # Failed recovery attempt

    def should_attempt(self) -> bool:
        """Check if operation should be attempted."""
        if self.state == CircuitState.CLOSED:
            return True

        if self.state == CircuitState.OPEN:
            # Check if recovery timeout elapsed
            if self.last_failure_time:
                elapsed = time.time() - self.last_failure_time
                if elapsed >= self.config.recovery_timeout_s:
                    self._try_recover()
                    return True
            return False

        # HALF_OPEN: attempt recovery
        return True

    def _open_circuit(self) -> None:
        """Transition to OPEN state."""
        self.state = CircuitState.OPEN
        self.success_count = 0
        self.last_state_change_time = time.time()
        logger.warning(
            f"Circuit breaker OPENED for {self.host} after {self.failure_count} failures"
        )

    def _try_recover(self) -> None:
        """Transition to HALF_OPEN state to test recovery."""
        self.state = CircuitState.HALF_OPEN
        self.success_count = 0
        self.last_state_change_time = time.time()
        logger.info(f"Circuit breaker HALF_OPEN for {self.host}, testing recovery")

    def _close_circuit(self) -> None:
        """Transition to CLOSED state."""
        self.state = CircuitState.CLOSED
        self.failure_count = 0
        self.success_count = 0
        self.last_state_change_time = time.time()
        logger.info(f"Circuit breaker CLOSED for {self.host}, service recovered")

    def get_state(self) -> str:
        """Get current state as string."""
        return self.state.value


class ConnectionManager:
    """Manages connection retry and recovery strategies."""

    def __init__(
        self,
        retry_config: Optional[RetryConfig] = None,
        circuit_breaker_config: Optional[CircuitBreakerConfig] = None,
    ):
        """
        Initialize connection manager.

        Args:
            retry_config: Retry configuration
            circuit_breaker_config: Circuit breaker configuration
        """
        self.retry_config = retry_config or RetryConfig()
        self.circuit_breaker_config = circuit_breaker_config or CircuitBreakerConfig()
        self._circuit_breakers: Dict[str, CircuitBreaker] = {}

    def get_circuit_breaker(self, host: str) -> CircuitBreaker:
        """Get or create circuit breaker for host."""
        if host not in self._circuit_breakers:
            self._circuit_breakers[host] = CircuitBreaker(host, self.circuit_breaker_config)
        return self._circuit_breakers[host]

    async def execute_with_retry(
        self,
        host: str,
        operation: callable,
        *args,
        **kwargs,
    ):
        """
        Execute operation with automatic retry and recovery.

        Args:
            host: Host being accessed
            operation: Async callable to execute
            *args: Positional arguments for operation
            **kwargs: Keyword arguments for operation

        Returns:
            Result from operation

        Raises:
            ConnectionError: If all retries exhausted
            CircuitBreakerOpen: If circuit is OPEN and not recovering
        """
        breaker = self.get_circuit_breaker(host)

        # Check circuit breaker
        if not breaker.should_attempt():
            raise ConnectionError(
                f"Circuit breaker is OPEN for {host} ({breaker.state.value}). "
                f"Service unavailable. Retry after {self.circuit_breaker_config.recovery_timeout_s}s"
            )

        last_error: Optional[Exception] = None
        delay = self.retry_config.initial_delay_s

        for attempt in range(1, self.retry_config.max_attempts + 1):
            try:
                logger.debug(f"Attempt {attempt}/{self.retry_config.max_attempts} for {host}")

                result = await operation(*args, **kwargs)
                breaker.record_success()
                return result

            except Exception as e:
                last_error = e
                breaker.record_failure()

                if attempt < self.retry_config.max_attempts:
                    # Calculate backoff with jitter
                    import random

                    jitter = delay * self.retry_config.jitter * random.random()
                    wait_time = min(delay + jitter, self.retry_config.max_delay_s)

                    logger.warning(
                        f"Attempt {attempt} failed for {host}: {e}. "
                        f"Retrying in {wait_time:.1f}s ({breaker.state.value})"
                    )

                    await asyncio.sleep(wait_time)
                    delay *= self.retry_config.exponential_base
                else:
                    logger.error(f"All {self.retry_config.max_attempts} attempts failed for {host}")

        # All retries exhausted
        assert last_error is not None
        raise ConnectionError(
            f"Failed to connect to {host} after {self.retry_config.max_attempts} attempts: {last_error}"
        ) from last_error

    def get_host_status(self, host: str) -> Dict[str, any]:
        """Get status of connection to host."""
        if host not in self._circuit_breakers:
            return {"host": host, "status": "unknown"}

        breaker = self._circuit_breakers[host]
        return {
            "host": host,
            "status": breaker.get_state(),
            "failures": breaker.failure_count,
            "successes": breaker.success_count,
            "last_failure": breaker.last_failure_time,
            "last_state_change": breaker.last_state_change_time,
        }

    def reset_host(self, host: str) -> None:
        """Manually reset circuit breaker for host."""
        if host in self._circuit_breakers:
            breaker = self._circuit_breakers[host]
            old_state = breaker.state.value
            breaker._close_circuit()
            logger.info(f"Manually reset {host} from {old_state} to {breaker.get_state()}")

    def reset_all(self) -> None:
        """Reset all circuit breakers."""
        for host, breaker in self._circuit_breakers.items():
            old_state = breaker.state.value
            breaker._close_circuit()
            logger.info(f"Reset {host} from {old_state} to {breaker.get_state()}")
