# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Resilience for Familiar

Provides:
- Retry decorator with exponential backoff
- Rate limit handling for Anthropic/OpenAI
- Circuit breaker pattern for failing fast
- Timeout handling

Usage:
    from familiar.core.resilience import with_retry, RetryConfig

    @with_retry()
    def call_llm():
        return client.messages.create(...)

    # Or with custom config:
    @with_retry(RetryConfig(max_attempts=5, base_delay=2.0))
    def call_llm():
        ...

Circuit Breaker Usage:
    from familiar.core.resilience import get_circuit_breaker, CircuitBreakerError

    breaker = get_circuit_breaker("anthropic")

    try:
        with breaker:
            response = client.messages.create(...)
    except CircuitBreakerError as e:
        # Circuit is open, fail fast
        print(f"Service unavailable, retry in {e.time_until_retry:.0f}s")
"""

import functools
import logging
import threading
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Callable, Optional

from .observability import Span, Trace

logger = logging.getLogger(__name__)


class RetryableError(Exception):
    """Base class for errors that should trigger a retry."""

    retry_after: Optional[float] = None


class RateLimitError(RetryableError):
    """Rate limit exceeded - should retry after delay."""

    def __init__(self, message: str, retry_after: Optional[float] = None):
        super().__init__(message)
        self.retry_after = retry_after


class APIConnectionError(RetryableError):
    """Connection failed - should retry."""

    pass


class APITimeoutError(RetryableError):
    """Request timed out - should retry."""

    pass


class OverloadedError(RetryableError):
    """Service overloaded - should retry with backoff."""

    pass


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

    def __init__(self, message: str, circuit_name: str, time_until_retry: float):
        super().__init__(message)
        self.circuit_name = circuit_name
        self.time_until_retry = time_until_retry


class CircuitState(str, Enum):
    """States of a circuit breaker."""

    CLOSED = "closed"  # Normal operation
    OPEN = "open"  # Failing fast, not allowing requests
    HALF_OPEN = "half_open"  # Testing if service recovered


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

    failure_threshold: int = 5  # Failures before opening
    recovery_timeout: float = 30.0  # Seconds before trying again
    success_threshold: int = 2  # Successes in half-open before closing
    # Which exceptions count as failures (empty = all exceptions)
    failure_exceptions: tuple = field(default_factory=tuple)


class CircuitBreaker:
    """
    Circuit breaker pattern implementation.

    Prevents cascading failures by failing fast when a service is unhealthy.

    States:
        CLOSED: Normal operation, requests pass through
        OPEN: Service unhealthy, requests fail immediately
        HALF_OPEN: Testing recovery, limited requests allowed

    Usage:
        breaker = CircuitBreaker("anthropic_api")

        try:
            with breaker:
                response = client.messages.create(...)
        except CircuitBreakerError:
            # Circuit is open, use fallback
            return cached_response()

    Or as decorator:
        @breaker
        def call_api():
            return client.messages.create(...)
    """

    def __init__(self, name: str, config: Optional[CircuitBreakerConfig] = None):
        self.name = name
        self.config = config or CircuitBreakerConfig()

        self._state = CircuitState.CLOSED
        self._failure_count = 0
        self._success_count = 0
        self._last_failure_time: Optional[float] = None
        import threading

        self._lock = threading.Lock()

    @property
    def state(self) -> CircuitState:
        """Current circuit state, with automatic transition from OPEN to HALF_OPEN."""
        with self._lock:
            if self._state == CircuitState.OPEN:
                # Check if recovery timeout has passed
                if self._last_failure_time is not None:
                    elapsed = time.time() - self._last_failure_time
                    if elapsed >= self.config.recovery_timeout:
                        logger.info(f"Circuit '{self.name}' transitioning to HALF_OPEN")
                        self._state = CircuitState.HALF_OPEN
                        self._success_count = 0
            return self._state

    @property
    def is_closed(self) -> bool:
        return self.state == CircuitState.CLOSED

    @property
    def is_open(self) -> bool:
        return self.state == CircuitState.OPEN

    @property
    def time_until_retry(self) -> float:
        """Seconds until circuit might allow requests (0 if not open)."""
        if self._state != CircuitState.OPEN or self._last_failure_time is None:
            return 0.0
        elapsed = time.time() - self._last_failure_time
        remaining = self.config.recovery_timeout - elapsed
        return max(0.0, remaining)

    def record_success(self):
        """Record a successful call."""
        with self._lock:
            if self._state == CircuitState.HALF_OPEN:
                self._success_count += 1
                if self._success_count >= self.config.success_threshold:
                    logger.info(f"Circuit '{self.name}' closing after recovery")
                    self._state = CircuitState.CLOSED
                    self._failure_count = 0
            elif self._state == CircuitState.CLOSED:
                # Reset failure count on success
                self._failure_count = 0

    def record_failure(self, exception: Optional[Exception] = None):
        """Record a failed call."""
        with self._lock:
            # Check if this exception should count as a failure
            if exception and self.config.failure_exceptions:
                if not isinstance(exception, self.config.failure_exceptions):
                    return  # Don't count this as a failure

            self._failure_count += 1
            self._last_failure_time = time.time()

            if self._state == CircuitState.HALF_OPEN:
                # Any failure in half-open reopens the circuit
                logger.warning(f"Circuit '{self.name}' reopening after failure in half-open")
                self._state = CircuitState.OPEN
                self._success_count = 0

            elif self._state == CircuitState.CLOSED:
                if self._failure_count >= self.config.failure_threshold:
                    logger.warning(
                        f"Circuit '{self.name}' opening after {self._failure_count} failures"
                    )
                    self._state = CircuitState.OPEN

    def allow_request(self) -> bool:
        """Check if a request should be allowed through."""
        state = self.state  # This may transition OPEN -> HALF_OPEN

        if state == CircuitState.CLOSED:
            return True
        elif state == CircuitState.OPEN:
            return False
        elif state == CircuitState.HALF_OPEN:
            # Allow request for testing
            return True
        return False

    def reset(self):
        """Manually reset the circuit to closed state."""
        with self._lock:
            self._state = CircuitState.CLOSED
            self._failure_count = 0
            self._success_count = 0
            self._last_failure_time = None
            logger.info(f"Circuit '{self.name}' manually reset")

    def __enter__(self):
        """Context manager entry - check if request allowed."""
        if not self.allow_request():
            raise CircuitBreakerError(
                f"Circuit '{self.name}' is open",
                circuit_name=self.name,
                time_until_retry=self.time_until_retry,
            )
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - record success or failure."""
        if exc_val is None:
            self.record_success()
        else:
            self.record_failure(exc_val)
        return False  # Don't suppress exceptions

    def __call__(self, func: Callable) -> Callable:
        """Use as a decorator."""

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            with self:
                return func(*args, **kwargs)

        return wrapper

    def get_status(self) -> dict:
        """Get current circuit breaker status."""
        return {
            "name": self.name,
            "state": self.state.value,
            "failure_count": self._failure_count,
            "success_count": self._success_count,
            "time_until_retry": self.time_until_retry,
            "config": {
                "failure_threshold": self.config.failure_threshold,
                "recovery_timeout": self.config.recovery_timeout,
                "success_threshold": self.config.success_threshold,
            },
        }


# Global circuit breaker registry
_circuit_breakers: dict[str, CircuitBreaker] = {}
_circuit_breakers_lock = threading.Lock()


def get_circuit_breaker(name: str, config: Optional[CircuitBreakerConfig] = None) -> CircuitBreaker:
    """
    Get or create a named circuit breaker.

    Circuit breakers are shared by name, so multiple calls with the
    same name return the same instance.

    Args:
        name: Unique identifier for the circuit (e.g., "anthropic", "openai")
        config: Configuration (only used if creating new circuit)

    Returns:
        The circuit breaker instance
    """
    with _circuit_breakers_lock:
        if name not in _circuit_breakers:
            _circuit_breakers[name] = CircuitBreaker(name, config)
        return _circuit_breakers[name]


def reset_circuit_breaker(name: str):
    """Reset a circuit breaker by name."""
    with _circuit_breakers_lock:
        if name in _circuit_breakers:
            _circuit_breakers[name].reset()


def get_all_circuit_status() -> dict[str, dict]:
    """Get status of all circuit breakers."""
    return {name: cb.get_status() for name, cb in _circuit_breakers.items()}


class ProviderCapacityTracker:
    """
    Track per-provider rate limit events in a sliding window.

    Complements circuit breakers (which track all failures with threshold=5).
    This is rate-limit-specific with a lower threshold and longer window,
    enabling proactive fallback before burning retries.
    """

    def __init__(self, window_seconds: float = 300.0, threshold: int = 2):
        self._window = window_seconds
        self._threshold = threshold
        self._events: dict[str, list[float]] = {}  # provider -> [timestamps]
        self._lock = __import__("threading").Lock()

    def record_rate_limit(self, provider_name: str) -> None:
        """Record a rate limit event for a provider."""
        provider_name = provider_name.lower()
        now = time.time()
        with self._lock:
            if provider_name not in self._events:
                self._events[provider_name] = []
            self._events[provider_name].append(now)
            self._prune(provider_name, now)

    def is_rate_limited(self, provider_name: str) -> bool:
        """Check if a provider has hit the rate limit threshold in the current window."""
        provider_name = provider_name.lower()
        now = time.time()
        with self._lock:
            if provider_name not in self._events:
                return False
            self._prune(provider_name, now)
            return len(self._events[provider_name]) >= self._threshold

    def _prune(self, provider_name: str, now: float) -> None:
        """Remove events outside the sliding window. Must hold self._lock."""
        cutoff = now - self._window
        self._events[provider_name] = [t for t in self._events[provider_name] if t > cutoff]

    def get_status(self) -> dict[str, dict]:
        """Return diagnostics for all tracked providers."""
        now = time.time()
        with self._lock:
            result = {}
            for name, events in self._events.items():
                active = [t for t in events if t > now - self._window]
                result[name] = {
                    "events_in_window": len(active),
                    "threshold": self._threshold,
                    "is_rate_limited": len(active) >= self._threshold,
                    "window_seconds": self._window,
                }
            return result


# Map provider exceptions to our retryable types
def classify_exception(exc: Exception) -> tuple[bool, Optional[float]]:
    """
    Determine if an exception is retryable and extract retry_after if available.

    Returns: (is_retryable, retry_after_seconds)
    """
    exc_type = type(exc).__name__
    exc_module = type(exc).__module__

    # Anthropic exceptions
    if "anthropic" in exc_module:
        if exc_type == "RateLimitError":
            # Try to extract retry-after from response headers
            retry_after = getattr(exc, "retry_after", None)
            if retry_after is None and hasattr(exc, "response"):
                retry_after = exc.response.headers.get("retry-after")
                if retry_after:
                    retry_after = float(retry_after)
            return True, retry_after or 60.0

        if exc_type == "APIConnectionError":
            return True, None

        if exc_type == "APITimeoutError":
            return True, None

        if exc_type == "OverloadedError":
            return True, 30.0

        if exc_type == "InternalServerError":
            return True, 5.0

    # OpenAI exceptions
    if "openai" in exc_module:
        if exc_type == "RateLimitError":
            retry_after = None
            if hasattr(exc, "response") and exc.response:
                retry_after = exc.response.headers.get("retry-after")
                if retry_after:
                    retry_after = float(retry_after)
            return True, retry_after or 60.0

        if exc_type == "APIConnectionError":
            return True, None

        if exc_type == "APITimeoutError":
            return True, None

        if exc_type == "InternalServerError":
            return True, 5.0

    # Generic connection errors
    if exc_type in ("ConnectionError", "TimeoutError", "ConnectionRefusedError"):
        return True, None

    # urllib3/requests errors
    if "urllib3" in exc_module or "requests" in exc_module:
        if "timeout" in str(exc).lower() or "connection" in str(exc).lower():
            return True, None

    # Familiar's own wrapped provider exceptions
    if "familiar" in exc_module:
        if exc_type == "ProviderRateLimitError":
            retry_after = getattr(exc, "retry_after", None)
            return True, retry_after or 60.0
        if exc_type in ("ProviderConnectionError", "ProviderTimeoutError"):
            return True, None

    return False, None


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

    max_attempts: int = 3
    base_delay: float = 1.0  # Initial delay in seconds
    max_delay: float = 60.0  # Maximum delay between retries
    exponential_base: float = 2.0  # Multiplier for exponential backoff
    jitter: float = 0.1  # Random jitter factor (0-1)

    # Which exceptions to retry on (empty = use classify_exception)
    retry_exceptions: tuple = field(default_factory=tuple)

    # Callback for custom retry logic
    on_retry: Optional[Callable[[Exception, int, float], None]] = None

    def calculate_delay(self, attempt: int, retry_after: Optional[float] = None) -> float:
        """Calculate delay before next retry."""
        import random

        if retry_after is not None:
            # Use server-provided retry-after, but cap it
            return min(retry_after, self.max_delay)

        # Exponential backoff: base * (exponential_base ^ attempt)
        delay = self.base_delay * (self.exponential_base ** (attempt - 1))

        # Add jitter
        if self.jitter > 0:
            jitter_range = delay * self.jitter
            delay += random.uniform(-jitter_range, jitter_range)

        # Cap at max_delay
        return min(delay, self.max_delay)


# Default configuration
DEFAULT_RETRY_CONFIG = RetryConfig()


def with_retry(
    config: Optional[RetryConfig] = None,
    trace: Optional[Trace] = None,
    span_name: Optional[str] = None,
):
    """
    Decorator that adds retry logic with exponential backoff.

    Args:
        config: Retry configuration (uses defaults if None)
        trace: Optional trace for observability
        span_name: Name for the span (defaults to function name)

    Usage:
        @with_retry()
        def call_api():
            return client.request()

        @with_retry(RetryConfig(max_attempts=5))
        def call_flaky_service():
            ...
    """
    config = config or DEFAULT_RETRY_CONFIG

    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # Extract trace from kwargs if provided at call time
            active_trace = kwargs.pop("_trace", None) or trace
            active_span_name = span_name or func.__name__

            last_exception = None

            for attempt in range(1, config.max_attempts + 1):
                # Create span for this attempt if we have a trace
                span = None
                if active_trace:
                    span = Span.create_compat(
                        name=active_span_name, trace_id=active_trace.trace_id, attempt=attempt
                    )

                try:
                    result = func(*args, **kwargs)

                    # Success - record span and return
                    if span:
                        span.finish()
                        span.status = "ok"
                        active_trace.add_span(span)

                    return result

                except Exception as e:
                    last_exception = e
                    is_retryable, retry_after = classify_exception(e)

                    # Check if we should retry
                    should_retry = is_retryable or (
                        config.retry_exceptions and isinstance(e, config.retry_exceptions)
                    )

                    if not should_retry or attempt >= config.max_attempts:
                        # Final failure
                        if span:
                            span.set_error(e, status="error")
                            span.finish()
                            active_trace.add_span(span)

                        logger.error(
                            f"[{active_trace.trace_id if active_trace else 'no-trace'}] "
                            f"{active_span_name} failed after {attempt} attempt(s): {e}"
                        )
                        raise

                    # Calculate delay
                    delay = config.calculate_delay(attempt, retry_after)

                    # Record retry span
                    if span:
                        span.set_error(e, status="retry")
                        span.metadata["retry_delay"] = delay
                        span.finish()
                        active_trace.add_span(span)

                    logger.warning(
                        f"[{active_trace.trace_id if active_trace else 'no-trace'}] "
                        f"{active_span_name} attempt {attempt} failed: {type(e).__name__}. "
                        f"Retrying in {delay:.1f}s..."
                    )

                    # Callback hook
                    if config.on_retry:
                        config.on_retry(e, attempt, delay)

                    time.sleep(delay)

            # Should not reach here, but just in case
            raise last_exception

        return wrapper

    return decorator


class RetryableOperation:
    """
    Context manager for retryable operations with full observability.

    More flexible than the decorator for complex scenarios.

    Usage:
        with RetryableOperation(trace, "llm_call", config) as op:
            for attempt in op.attempts():
                with attempt:
                    result = provider.chat(...)
                    attempt.span.set_tokens(result.usage, model)
    """

    def __init__(self, trace: Optional[Trace], name: str, config: Optional[RetryConfig] = None):
        self.trace = trace
        self.name = name
        self.config = config or DEFAULT_RETRY_CONFIG
        self._current_attempt = 0
        self._success = False
        self._last_error: Optional[Exception] = None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    def attempts(self):
        """Generator yielding attempt contexts."""
        for attempt_num in range(1, self.config.max_attempts + 1):
            self._current_attempt = attempt_num
            yield AttemptContext(self, attempt_num)

            if self._success:
                break

    @property
    def succeeded(self) -> bool:
        return self._success

    @property
    def last_error(self) -> Optional[Exception]:
        return self._last_error


class AttemptContext:
    """Context for a single retry attempt."""

    def __init__(self, operation: RetryableOperation, attempt: int):
        self.operation = operation
        self.attempt = attempt
        self.span: Optional[Span] = None
        self._delay_applied = False

    def __enter__(self):
        # Create span if we have a trace
        if self.operation.trace:
            self.span = Span.create_compat(
                name=self.operation.name,
                trace_id=self.operation.trace.trace_id,
                attempt=self.attempt,
            )
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        trace = self.operation.trace
        if exc_val is None:
            # Success
            self.operation._success = True
            if self.span and trace:
                self.span.status = "ok"
                self.span.finish()
                trace.add_span(self.span)
            return False

        # Error occurred
        self.operation._last_error = exc_val
        is_retryable, retry_after = classify_exception(exc_val)

        config = self.operation.config
        should_retry = is_retryable and self.attempt < config.max_attempts

        if self.span and trace:
            self.span.set_error(exc_val, status="retry" if should_retry else "error")
            self.span.finish()
            trace.add_span(self.span)

        if not should_retry:
            # Don't suppress - let it propagate
            return False

        # Apply delay before next attempt
        delay = config.calculate_delay(self.attempt, retry_after)

        logger.warning(
            f"[{trace.trace_id if trace else 'no-trace'}] "
            f"{self.operation.name} attempt {self.attempt} failed: {type(exc_val).__name__}. "
            f"Retrying in {delay:.1f}s..."
        )

        if config.on_retry:
            config.on_retry(exc_val, self.attempt, delay)

        time.sleep(delay)

        # Suppress exception to allow retry
        return True


def retry_on_rate_limit(func: Callable) -> Callable:
    """
    Simple decorator specifically for rate limit handling.

    Convenience wrapper around with_retry with sensible defaults for API calls.
    """
    return with_retry(
        RetryConfig(
            max_attempts=3,
            base_delay=1.0,
            max_delay=120.0,  # Allow longer delays for rate limits
            exponential_base=2.0,
        )
    )(func)
