"""
cortex/providers/fallback.py — Provider fallback chain with circuit breaker.

Provides a resilient multi-provider LLM call layer for the v1.2.1+ provider
architecture. When the primary provider fails, the chain automatically tries
the next provider. A per-provider circuit breaker prevents hammering broken
providers during outages.

Design
------
ProviderChain
    Ordered list of provider configs. Each entry specifies the provider name,
    base_url, api_key, and model. The chain tries providers in order until one
    succeeds (or all fail).

CircuitBreaker
    Per-provider state machine with three states:

        CLOSED   — Normal. All calls allowed.
        OPEN     — Tripped. Calls to this provider are blocked.
                   Auto-transitions to HALF_OPEN after cooldown_s seconds.
        HALF_OPEN — Probe state. One call is allowed to test recovery.
                   Success → CLOSED. Failure → OPEN (resets cooldown).

    Opens after ``threshold`` consecutive QUOTA_EXHAUSTED or RATE_LIMITED
    failures. Transient errors (5xx, network) and model errors (400/404)
    do NOT count toward the threshold — they are handled by retry logic.

call_with_fallback(prompt, chain, breaker)
    High-level entry point. Tries each provider in chain order, respecting
    breaker state. Logs attempts, failures, and final outcome. Returns the
    first successful CompletionResponse.

Config fields (added to CortexConfig)
--------------------------------------
    provider_fallback_chain   : List[str]  — ordered provider names
    circuit_breaker_threshold : int = 3   — failures before opening (per-chain default)
    circuit_breaker_cooldown_s: int = 60  — seconds before half-open probe

Usage::

    from cortex.providers.fallback import ProviderChain, CircuitBreaker, call_with_fallback

    chain = ProviderChain.from_config(config)
    breaker = CircuitBreaker(
        threshold=config.circuit_breaker_threshold,
        cooldown_s=config.circuit_breaker_cooldown_s,
    )
    response = await call_with_fallback(prompt="Extract memories", chain=chain, breaker=breaker)

Dependencies: stdlib only (no vendor SDK imports).
"""
from __future__ import annotations

import asyncio
import enum
import logging
import time
from dataclasses import dataclass, field
from typing import Any, Callable, Coroutine, List, Optional

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Failure classification (duplicated from cortex.llm.circuit_breaker for
# independence; the old llm/ stack will be removed in a future cleanup)
# ---------------------------------------------------------------------------


class FailureType(enum.Enum):
    """Classification of a provider failure for circuit-breaker accounting."""

    QUOTA_EXHAUSTED = "quota_exhausted"
    """Billing/quota limit hit. Persistent until reset."""

    RATE_LIMITED = "rate_limited"
    """HTTP 429 / transient throttle. Back off and retry."""

    TRANSIENT = "transient"
    """Network error, 5xx, timeout. Retryable; does NOT trip breaker."""

    MODEL_ERROR = "model_error"
    """HTTP 400/404, invalid model. Config issue; does NOT trip breaker."""


def classify_failure(exc: Exception) -> FailureType:
    """Classify a provider exception into a FailureType.

    Inspection order:
        1. HTTP status code (``status_code``, ``status``, ``code`` attribute)
        2. Error message keywords
        3. Exception type (ConnectionError, TimeoutError → TRANSIENT)
        4. Default: TRANSIENT (safe, non-tripping)
    """
    msg = str(exc).lower()

    raw_code = (
        getattr(exc, "status_code", None)
        or getattr(exc, "status", None)
        or getattr(exc, "code", None)
    )
    code: Optional[int] = None
    if raw_code is not None:
        try:
            code = int(raw_code)
        except (ValueError, TypeError):
            code = None

    # --- 429: may be quota exhaustion or rate limit ---
    is_429 = code == 429 or "429" in msg or "rate limit" in msg or "ratelimit" in msg
    if is_429:
        quota_kws = ("quota", "billing", "exceeded your", "usage limit", "insufficient_quota")
        if any(kw in msg for kw in quota_kws):
            return FailureType.QUOTA_EXHAUSTED
        return FailureType.RATE_LIMITED

    # --- Explicit quota / billing keywords ---
    quota_kws = ("quota", "billing", "insufficient_quota", "limit exceeded", "usage limit")
    if any(kw in msg for kw in quota_kws):
        return FailureType.QUOTA_EXHAUSTED

    # --- Model config errors — non-retryable, don't trip breaker ---
    model_kws = ("model not found", "model_not_found", "invalid model", "no such model")
    if code in (400, 404) or any(kw in msg for kw in model_kws):
        return FailureType.MODEL_ERROR

    # --- Transient: network / server errors ---
    if isinstance(exc, (ConnectionError, TimeoutError, asyncio.TimeoutError)):
        return FailureType.TRANSIENT
    if code in (500, 502, 503, 504):
        return FailureType.TRANSIENT
    transient_kws = ("timeout", "connection", "network", "server error", "unavailable")
    if any(kw in msg for kw in transient_kws):
        return FailureType.TRANSIENT

    return FailureType.TRANSIENT  # default: treat unknown as transient


# ---------------------------------------------------------------------------
# Circuit breaker
# ---------------------------------------------------------------------------


class BreakerState(enum.Enum):
    """Circuit breaker FSM states."""

    CLOSED = "closed"
    """Normal operation — all calls allowed."""

    OPEN = "open"
    """Tripped — calls to this provider are blocked."""

    HALF_OPEN = "half_open"
    """Cooldown elapsed — one probe call allowed to test recovery."""


class CircuitBreakerOpenError(Exception):
    """Raised when a call is attempted on a provider with an open circuit."""


@dataclass
class _BreakerState:
    """Per-provider mutable breaker state."""

    state: BreakerState = BreakerState.CLOSED
    consecutive_failures: int = 0
    opened_at: Optional[float] = None
    last_failure_type: Optional[FailureType] = None
    success_count: int = 0
    failure_count: int = 0


class CircuitBreaker:
    """Per-provider circuit breaker.

    Tracks failure rates per provider name and manages the CLOSED →
    OPEN → HALF_OPEN → CLOSED state machine.

    Args:
        threshold:   Consecutive quota/rate-limit failures before opening.
                     Default: 3.
        cooldown_s:  Seconds after opening before auto-transitioning to
                     HALF_OPEN (probe state). Default: 60.
    """

    def __init__(self, threshold: int = 3, cooldown_s: int = 60) -> None:
        self.threshold = max(1, threshold)
        self.cooldown_s = max(1, cooldown_s)
        self._providers: dict[str, _BreakerState] = {}
        # B-010: asyncio.Lock for all state transitions
        self._lock = asyncio.Lock()
        # B-007: per-provider probe locks to prevent concurrent HALF_OPEN probes
        self._probe_locks: dict[str, asyncio.Lock] = {}

    # --- internal helpers ---------------------------------------------------

    def _get(self, provider: str) -> _BreakerState:
        if provider not in self._providers:
            self._providers[provider] = _BreakerState()
        return self._providers[provider]

    def _get_probe_lock(self, provider: str) -> asyncio.Lock:
        """Return the per-provider probe lock, creating if needed."""
        if provider not in self._probe_locks:
            self._probe_locks[provider] = asyncio.Lock()
        return self._probe_locks[provider]

    def _check_recovery(self, breaker: _BreakerState) -> None:
        """Transition OPEN → HALF_OPEN if cooldown has elapsed."""
        if (
            breaker.state == BreakerState.OPEN
            and breaker.opened_at is not None
            and (time.monotonic() - breaker.opened_at) >= self.cooldown_s
        ):
            breaker.state = BreakerState.HALF_OPEN
            logger.info(
                "CircuitBreaker: transitioning to HALF_OPEN after %.0fs cooldown",
                self.cooldown_s,
            )

    # --- public API ---------------------------------------------------------

    async def record_failure(self, provider: str, exc: Exception) -> None:
        """Record a provider failure and potentially open the circuit.

        Only QUOTA_EXHAUSTED and RATE_LIMITED failures count toward the
        threshold. TRANSIENT and MODEL_ERROR failures are logged but do not
        contribute to the breaker count.

        B-008: Any failure during HALF_OPEN (including TRANSIENT) transitions
        the breaker back to OPEN with a reset cooldown timer, since the probe
        call failed regardless of failure classification.

        B-010: All state mutations are serialized via ``self._lock``.
        """
        ftype = classify_failure(exc)
        async with self._lock:
            b = self._get(provider)
            b.failure_count += 1
            b.last_failure_type = ftype

            # B-008: Any failure during HALF_OPEN probe → re-open breaker
            if b.state == BreakerState.HALF_OPEN:
                b.state = BreakerState.OPEN
                b.opened_at = time.monotonic()
                logger.warning(
                    "CircuitBreaker[%s]: HALF_OPEN probe failed (%s: %s) — "
                    "re-opening breaker for %ds.",
                    provider, ftype.value, exc, self.cooldown_s,
                )
                return

            if ftype not in (FailureType.QUOTA_EXHAUSTED, FailureType.RATE_LIMITED):
                logger.debug(
                    "CircuitBreaker[%s]: non-tripping failure %s (%s)",
                    provider, ftype.value, exc,
                )
                return

            b.consecutive_failures += 1
            logger.warning(
                "CircuitBreaker[%s]: %s failure #%d/%d — %s",
                provider, ftype.value, b.consecutive_failures, self.threshold, exc,
            )

            if b.consecutive_failures >= self.threshold and b.state != BreakerState.OPEN:
                b.state = BreakerState.OPEN
                b.opened_at = time.monotonic()
                logger.error(
                    "CircuitBreaker[%s]: OPEN after %d consecutive %s failures. "
                    "Calls blocked for %ds.",
                    provider, b.consecutive_failures, ftype.value, self.cooldown_s,
                )

    async def record_success(self, provider: str) -> None:
        """Record a successful call. Resets failure count; may close breaker.

        B-010: All state mutations are serialized via ``self._lock``.
        """
        async with self._lock:
            b = self._get(provider)
            b.success_count += 1
            prev_state = b.state

            b.consecutive_failures = 0
            b.state = BreakerState.CLOSED
            b.opened_at = None

            if prev_state != BreakerState.CLOSED:
                logger.info(
                    "CircuitBreaker[%s]: CLOSED after successful call (was %s)",
                    provider, prev_state.value,
                )

    async def is_allowed(self, provider: str) -> bool:
        """Return True if a call to *provider* is currently allowed.

        CLOSED → allowed always.
        OPEN   → not allowed (unless cooldown elapsed → HALF_OPEN probe).
        HALF_OPEN → one probe allowed.

        B-010: All state mutations are serialized via ``self._lock``.
        """
        async with self._lock:
            b = self._get(provider)
            self._check_recovery(b)
            return b.state in (BreakerState.CLOSED, BreakerState.HALF_OPEN)

    async def is_probe(self, provider: str) -> bool:
        """Return True if the provider is in HALF_OPEN (probe) state.

        B-010: All state mutations are serialized via ``self._lock``.
        """
        async with self._lock:
            b = self._get(provider)
            self._check_recovery(b)
            return b.state == BreakerState.HALF_OPEN

    async def get_state(self, provider: str) -> BreakerState:
        """Return the current BreakerState for *provider* (with auto-recovery check).

        B-010: All state mutations are serialized via ``self._lock``.
        """
        async with self._lock:
            b = self._get(provider)
            self._check_recovery(b)
            return b.state

    async def reset(self, provider: str) -> None:
        """Manually reset the circuit breaker for *provider* to CLOSED.

        B-010: All state mutations are serialized via ``self._lock``.
        """
        async with self._lock:
            if provider in self._providers:
                b = self._providers[provider]
                b.state = BreakerState.CLOSED
                b.consecutive_failures = 0
                b.opened_at = None
                logger.info("CircuitBreaker[%s]: manually reset to CLOSED", provider)

    async def reset_all(self) -> None:
        """Reset all per-provider breakers.

        B-010: All state mutations are serialized via ``self._lock``.
        """
        async with self._lock:
            self._providers.clear()
            logger.info("CircuitBreaker: all breakers reset")

    async def metrics(self) -> dict[str, dict[str, Any]]:
        """Return per-provider metrics dict.

        B-010: All state mutations are serialized via ``self._lock``.

        Returns:
            Dict keyed by provider name with keys:
                state, consecutive_failures, success_count, failure_count,
                last_failure_type, opened_at.
        """
        async with self._lock:
            result: dict[str, dict[str, Any]] = {}
            for name, b in self._providers.items():
                self._check_recovery(b)
                result[name] = {
                    "state": b.state.value,
                    "consecutive_failures": b.consecutive_failures,
                    "success_count": b.success_count,
                    "failure_count": b.failure_count,
                    "last_failure_type": b.last_failure_type.value if b.last_failure_type else None,
                    "opened_at": b.opened_at,
                }
            return result


# ---------------------------------------------------------------------------
# Provider config entry
# ---------------------------------------------------------------------------


@dataclass
class ProviderConfig:
    """Configuration for a single provider in the fallback chain.

    Attributes:
        name:     Provider name (e.g. "google", "openai", "anthropic").
        base_url: API base URL (OpenAI-compatible endpoint).
        api_key:  API key for this provider.
        model:    Model name to use for completions.
    """

    name: str
    base_url: str = ""
    api_key: str = ""
    model: str = ""

    def is_configured(self) -> bool:
        """Return True if the provider has at least an api_key or base_url."""
        return bool(self.api_key or self.base_url)


# ---------------------------------------------------------------------------
# Provider chain
# ---------------------------------------------------------------------------


@dataclass
class ProviderChain:
    """Ordered list of provider configs for fallback.

    Attributes:
        providers: Ordered list of ProviderConfig entries. The chain tries
                   providers in list order; the first success wins.

    Usage::

        chain = ProviderChain([
            ProviderConfig("google", api_key="...", model="gemini-2.5-flash"),
            ProviderConfig("openai", api_key="...", model="gpt-4.1-mini"),
        ])
    """

    providers: List[ProviderConfig] = field(default_factory=list)

    @classmethod
    def from_config(cls, config: Any) -> "ProviderChain":
        """Build a ProviderChain from a CortexConfig instance.

        Reads ``provider_fallback_chain`` (list of names) and resolves
        per-provider credentials from ``provider_configs`` dict (mapping
        provider name → {api_key, base_url, model}). Falls back to the
        global ``provider_llm_api_key`` / ``provider_llm_base_url`` /
        ``provider_llm_model`` fields for backwards compatibility.

        B-009 fix: Each provider now gets its own credentials instead of
        sharing the same global api_key/model across all providers.

        Args:
            config: CortexConfig (or any object with provider_* attributes).

        Returns:
            ProviderChain with one ProviderConfig per entry in the priority list.
        """
        # New field: provider_fallback_chain. Fall back to llm_provider_priority.
        chain_names: List[str] = (
            getattr(config, "provider_fallback_chain", None)
            or getattr(config, "llm_provider_priority", None)
            or []
        )
        if not chain_names:
            logger.warning("ProviderChain.from_config: no provider_fallback_chain configured")
            return cls(providers=[])

        # B-009: Read per-provider configs dict if available
        per_provider: dict[str, Any] = getattr(config, "provider_configs", {}) or {}

        # Global fallbacks for backwards compatibility
        global_base_url = getattr(config, "provider_llm_base_url", "") or ""
        global_api_key = getattr(config, "provider_llm_api_key", "") or ""
        global_model = getattr(config, "provider_llm_model", "") or ""

        providers: List[ProviderConfig] = []
        for name in chain_names:
            provider_cfg = per_provider.get(name, {})
            if isinstance(provider_cfg, dict):
                p_base_url = provider_cfg.get("base_url", "") or global_base_url
                p_api_key = provider_cfg.get("api_key", "") or global_api_key
                p_model = provider_cfg.get("model", "") or global_model
            else:
                # Support object-like configs with attributes
                p_base_url = getattr(provider_cfg, "base_url", "") or global_base_url
                p_api_key = getattr(provider_cfg, "api_key", "") or global_api_key
                p_model = getattr(provider_cfg, "model", "") or global_model

            pc = ProviderConfig(
                name=name,
                base_url=p_base_url,
                api_key=p_api_key,
                model=p_model,
            )
            providers.append(pc)

        logger.debug(
            "ProviderChain.from_config: built chain with %d provider(s): %s",
            len(providers),
            [p.name for p in providers],
        )
        return cls(providers=providers)

    def provider_names(self) -> List[str]:
        """Return provider names in chain order."""
        return [p.name for p in self.providers]

    def __len__(self) -> int:
        return len(self.providers)

    def __bool__(self) -> bool:
        return bool(self.providers)


# ---------------------------------------------------------------------------
# Fallback call
# ---------------------------------------------------------------------------

# Type alias: a coroutine factory that accepts a ProviderConfig and returns
# any result. Used to inject actual LLM call logic in tests.
ProviderCallable = Callable[[ProviderConfig], Coroutine[Any, Any, Any]]


async def call_with_fallback(
    prompt: str,
    chain: ProviderChain,
    breaker: Optional[CircuitBreaker] = None,
    *,
    call_fn: Optional[ProviderCallable] = None,
) -> Any:
    """Try each provider in chain order; return the first success.

    Circuit breaker integration:
        - Providers with OPEN breakers are skipped (no network call).
        - HALF_OPEN providers get exactly one probe call.
        - Success → record_success() (may close breaker).
        - Failure → record_failure() (may open/keep-open breaker).

    Args:
        prompt:  Prompt string (passed to call_fn for LLM calls; also used
                 for log context).
        chain:   Ordered ProviderChain to try.
        breaker: Optional CircuitBreaker instance for per-provider state.
                 If None, no circuit breaking is applied (all providers tried
                 on every call).
        call_fn: Async callable ``call_fn(provider_config) → result``.
                 Inject the actual LLM call here. If None, raises
                 RuntimeError (useful for testing the chain/breaker logic
                 in isolation).

    Returns:
        The result returned by the first succeeding provider's call_fn.

    Raises:
        RuntimeError: If chain is empty, or all providers fail/are skipped.
    """
    if not chain:
        raise RuntimeError(
            "call_with_fallback: ProviderChain is empty. "
            "Configure at least one provider in provider_fallback_chain."
        )

    if call_fn is None:
        raise RuntimeError(
            "call_with_fallback: call_fn is required. "
            "Pass an async callable that accepts a ProviderConfig and returns a result."
        )

    last_exc: Exception = RuntimeError("No providers available in chain")
    tried: List[str] = []
    skipped: List[str] = []
    succeeded: Optional[str] = None

    for pc in chain.providers:
        name = pc.name

        # --- Check circuit breaker ---
        if breaker is not None and not await breaker.is_allowed(name):
            state = await breaker.get_state(name)
            logger.info(
                "call_with_fallback: skipping %r (circuit breaker %s) — prompt=%r",
                name, state.value, prompt[:80],
            )
            skipped.append(name)
            last_exc = CircuitBreakerOpenError(
                f"Provider {name!r} circuit breaker is {state.value}. "
                f"Cooldown: {getattr(breaker, 'cooldown_s', '?')}s."
            )
            continue

        # B-007: If this provider is in HALF_OPEN, acquire probe lock.
        # Only one concurrent probe is allowed; others skip this provider.
        is_probe = breaker is not None and await breaker.is_probe(name)
        if is_probe:
            probe_lock = breaker._get_probe_lock(name)
            if probe_lock.locked():
                # Another probe is already running — skip this provider
                logger.info(
                    "call_with_fallback: skipping %r (HALF_OPEN probe already in progress)",
                    name,
                )
                skipped.append(name)
                last_exc = CircuitBreakerOpenError(
                    f"Provider {name!r} HALF_OPEN probe already in progress."
                )
                continue

        tried.append(name)
        logger.info(
            "call_with_fallback: trying %r — prompt=%r",
            name, prompt[:80],
        )

        try:
            # B-007: Hold probe lock during the actual call if this is a probe
            if is_probe:
                async with breaker._get_probe_lock(name):
                    result = await call_fn(pc)
            else:
                result = await call_fn(pc)
            if breaker is not None:
                await breaker.record_success(name)
            succeeded = name
            logger.info(
                "call_with_fallback: %r succeeded — tried=%s skipped=%s",
                name, tried, skipped,
            )
            return result

        except Exception as exc:
            last_exc = exc
            if breaker is not None:
                await breaker.record_failure(name, exc)
            logger.warning(
                "call_with_fallback: %r failed (%s) — trying next provider",
                name, exc,
            )

    # All providers failed or skipped
    total = len(chain)
    logger.error(
        "call_with_fallback: all %d provider(s) failed — tried=%s skipped=%s — "
        "last_error=%s",
        total, tried, skipped, last_exc,
    )
    raise last_exc


# ---------------------------------------------------------------------------
# Metrics helper
# ---------------------------------------------------------------------------


async def chain_metrics(breaker: CircuitBreaker) -> dict[str, Any]:
    """Return a dict of per-provider metrics from *breaker*.

    Convenience wrapper around ``CircuitBreaker.metrics()`` with a top-level
    summary for dashboards/health checks.

    Returns::

        {
          "providers": {
            "google": {"state": "closed", "success_count": 42, ...},
            "openai":  {"state": "open",   "success_count": 7,  ...},
          },
          "total_providers": 2,
          "open_count": 1,
          "closed_count": 1,
          "half_open_count": 0,
        }
    """
    raw = await breaker.metrics()
    open_c = sum(1 for v in raw.values() if v["state"] == "open")
    half_c = sum(1 for v in raw.values() if v["state"] == "half_open")
    return {
        "providers": raw,
        "total_providers": len(raw),
        "open_count": open_c,
        "half_open_count": half_c,
        "closed_count": len(raw) - open_c - half_c,
    }
