"""
cortex/providers/rate_limiter.py — Per-provider async rate limiter for the
extraction pipeline.

Prevents burst LLM calls from exhausting API quotas and triggering 429 errors.

Algorithm: Token Bucket (dual-bucket variant)
  - One bucket tracks *requests* (RPM — requests per minute).
  - One bucket tracks *tokens* (TPM — tokens per minute).
  - Both buckets are per-provider and operate independently.
  - A call blocks until *both* buckets have capacity.
  - On HTTP 429 responses, exponential backoff with jitter is applied.

Usage::

    limiter = ProviderRateLimiter.from_config(config)

    async with limiter.acquire("openai", tokens=500):
        response = await llm.complete(...)

    # After a 429 response, record it to trigger backoff:
    await limiter.record_429("openai", retry_after=30.0)

Config (in cortex.toml under [providers.rate_limits])::

    [providers.rate_limits.openai]
    requests_per_minute = 60
    tokens_per_minute = 90_000

    [providers.rate_limits.anthropic]
    requests_per_minute = 50
    tokens_per_minute = 100_000

Or via CortexConfig.provider_rate_limits dict:
    config.provider_rate_limits = {
        "openai": {"requests_per_minute": 60, "tokens_per_minute": 90_000},
    }

Metrics are exposed via ProviderRateLimiter.metrics(provider).
"""
from __future__ import annotations

import asyncio
import logging
import math
import random
import time
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import Any, AsyncIterator, Dict, Optional

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Defaults
# ---------------------------------------------------------------------------

#: Conservative defaults when no per-provider config is supplied.
_DEFAULT_RPM = 60
_DEFAULT_TPM = 90_000

#: Backoff: base * 2^attempt + jitter, capped at this many seconds.
_MAX_BACKOFF_SECONDS = 64.0
_BACKOFF_BASE = 1.0
_JITTER_FRACTION = 0.25  # ±25% jitter


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------


@dataclass
class RateLimitConfig:
    """Per-provider rate limit configuration.

    Args:
        requests_per_minute: Maximum requests per minute (0 = unlimited).
        tokens_per_minute:   Maximum tokens per minute (0 = unlimited).
        burst_requests:      Burst headroom for requests (default = RPM).
        burst_tokens:        Burst headroom for tokens (default = TPM).
    """

    requests_per_minute: int = _DEFAULT_RPM
    tokens_per_minute: int = _DEFAULT_TPM
    burst_requests: int = 0   # 0 → default to requests_per_minute
    burst_tokens: int = 0     # 0 → default to tokens_per_minute

    def __post_init__(self) -> None:
        # B-025: When RPM=0 (unlimited), don't default burst to 1 — use a
        # large value so rate limiting is effectively skipped.
        if self.burst_requests <= 0:
            if self.requests_per_minute <= 0:
                self.burst_requests = int(1e9)  # effectively unlimited
            else:
                self.burst_requests = self.requests_per_minute
        if self.burst_tokens <= 0:
            if self.tokens_per_minute <= 0:
                self.burst_tokens = int(1e9)  # effectively unlimited
            else:
                self.burst_tokens = self.tokens_per_minute

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "RateLimitConfig":
        return cls(
            requests_per_minute=int(d.get("requests_per_minute", _DEFAULT_RPM)),
            tokens_per_minute=int(d.get("tokens_per_minute", _DEFAULT_TPM)),
            burst_requests=int(d.get("burst_requests", 0)),
            burst_tokens=int(d.get("burst_tokens", 0)),
        )


# ---------------------------------------------------------------------------
# Metrics
# ---------------------------------------------------------------------------


@dataclass
class RateLimitMetrics:
    """Counters accumulated per provider."""

    requests_allowed: int = 0
    """Total requests that passed through without waiting."""

    requests_throttled: int = 0
    """Total requests that were delayed (token bucket was empty)."""

    total_wait_seconds: float = 0.0
    """Cumulative wall-clock seconds spent waiting for capacity."""

    backoffs_triggered: int = 0
    """Number of times a 429 response triggered an exponential backoff."""

    total_backoff_seconds: float = 0.0
    """Cumulative seconds spent in 429 backoff sleep."""


# ---------------------------------------------------------------------------
# Internal bucket
# ---------------------------------------------------------------------------


class _AsyncTokenBucket:
    """Single async token bucket (not thread-safe; use under asyncio only)."""

    def __init__(self, capacity: float, refill_rate: float) -> None:
        """
        Args:
            capacity:     Maximum tokens (burst size).
            refill_rate:  Tokens added per second.
        """
        self._capacity = capacity
        self._refill_rate = refill_rate
        self._tokens = float(capacity)
        self._last_refill = time.monotonic()

    def _refill(self) -> None:
        now = time.monotonic()
        elapsed = now - self._last_refill
        self._tokens = min(
            self._capacity,
            self._tokens + elapsed * self._refill_rate,
        )
        self._last_refill = now

    def available(self) -> float:
        """Current available tokens (after refill)."""
        self._refill()
        return self._tokens

    def try_consume(self, amount: float) -> tuple[bool, float]:
        """
        Attempt to consume *amount* tokens.

        B-023: Negative amounts are clamped to 0 to prevent inflating the
        bucket.

        Returns:
            (ok, wait_seconds) — if ok=False, wait_seconds > 0 indicates how
            long to wait before retrying.
        """
        amount = max(0.0, amount)  # B-023: guard against negative amounts
        self._refill()
        if self._tokens >= amount:
            self._tokens -= amount
            return True, 0.0
        tokens_needed = amount - self._tokens
        wait = tokens_needed / self._refill_rate
        return False, wait

    async def acquire(self, amount: float, metrics: Optional[RateLimitMetrics] = None) -> None:
        """Block until *amount* tokens are available, then consume them."""
        while True:
            ok, wait = self.try_consume(amount)
            if ok:
                return
            if metrics is not None:
                metrics.requests_throttled += 1
                metrics.total_wait_seconds += wait
            logger.debug(
                "Rate bucket throttle: need %.1f tokens, waiting %.3fs",
                amount,
                wait,
            )
            await asyncio.sleep(wait)


# ---------------------------------------------------------------------------
# Per-provider state
# ---------------------------------------------------------------------------


@dataclass
class _ProviderState:
    config: RateLimitConfig
    req_bucket: _AsyncTokenBucket
    tok_bucket: _AsyncTokenBucket
    metrics: RateLimitMetrics = field(default_factory=RateLimitMetrics)

    # Backoff state
    backoff_until: float = 0.0   # monotonic time; sleep until this if >now
    consecutive_429s: int = 0

    @classmethod
    def from_config(cls, cfg: RateLimitConfig) -> "_ProviderState":
        req_rate = cfg.requests_per_minute / 60.0 if cfg.requests_per_minute > 0 else 1e9
        tok_rate = cfg.tokens_per_minute / 60.0 if cfg.tokens_per_minute > 0 else 1e9
        return cls(
            config=cfg,
            req_bucket=_AsyncTokenBucket(
                capacity=float(cfg.burst_requests),
                refill_rate=req_rate,
            ),
            tok_bucket=_AsyncTokenBucket(
                capacity=float(cfg.burst_tokens),
                refill_rate=tok_rate,
            ),
        )


# ---------------------------------------------------------------------------
# Public limiter
# ---------------------------------------------------------------------------


class ProviderRateLimiter:
    """
    Async per-provider rate limiter for LLM calls in the extraction pipeline.

    Features:
    - Token bucket with dual limits: requests/min + tokens/min
    - Per-provider isolation (each provider has its own buckets)
    - Automatic backoff on 429 responses (exponential + jitter)
    - Metrics: throttled requests, wait times, backoff events
    - Configurable via CortexConfig.provider_rate_limits dict

    Thread safety: designed for single-event-loop async use. The bucket
    operations are not thread-safe; use a separate limiter per event loop
    thread if needed (standard asyncio pattern).

    Example::

        limiter = ProviderRateLimiter()
        limiter.configure("openai", RateLimitConfig(requests_per_minute=60))

        async with limiter.acquire("openai", tokens=500):
            resp = await llm.complete(...)
    """

    def __init__(
        self,
        defaults: Optional[RateLimitConfig] = None,
    ) -> None:
        self._defaults = defaults or RateLimitConfig()
        self._providers: Dict[str, _ProviderState] = {}

    # ------------------------------------------------------------------
    # Configuration
    # ------------------------------------------------------------------

    def configure(self, provider: str, cfg: RateLimitConfig) -> None:
        """Add or replace the rate limit configuration for *provider*."""
        self._providers[provider] = _ProviderState.from_config(cfg)
        logger.debug(
            "ProviderRateLimiter: configured %s rpm=%d tpm=%d",
            provider,
            cfg.requests_per_minute,
            cfg.tokens_per_minute,
        )

    def configure_from_dict(self, provider: str, d: Dict[str, Any]) -> None:
        """Configure *provider* from a plain dict (e.g., from TOML config)."""
        self.configure(provider, RateLimitConfig.from_dict(d))

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

        Reads ``config.provider_rate_limits`` (dict mapping provider name →
        config dict) if present.  Falls back to a permissive default.
        """
        limiter = cls()
        rate_limits: Dict[str, Any] = getattr(config, "provider_rate_limits", {}) or {}
        for provider, limits in rate_limits.items():
            if isinstance(limits, dict):
                limiter.configure_from_dict(provider, limits)
            elif isinstance(limits, RateLimitConfig):
                limiter.configure(provider, limits)
        return limiter

    def _get_state(self, provider: str) -> _ProviderState:
        """Return the state for *provider*, creating it from defaults if needed."""
        if provider not in self._providers:
            state = _ProviderState.from_config(self._defaults)
            self._providers[provider] = state
            logger.debug(
                "ProviderRateLimiter: auto-configured %s with defaults",
                provider,
            )
        return self._providers[provider]

    # ------------------------------------------------------------------
    # Core acquire
    # ------------------------------------------------------------------

    @asynccontextmanager
    async def acquire(
        self,
        provider: str,
        tokens: int = 0,
    ) -> AsyncIterator[None]:
        """
        Async context manager — acquire capacity before making an LLM call.

        Blocks until:
        1. Any active 429 backoff for this provider has expired.
        2. The request-bucket has ≥1 slot.
        3. The token-bucket has ≥*tokens* capacity (if tokens > 0).

        Usage::

            async with limiter.acquire("openai", tokens=500):
                response = await llm.complete(...)

        Args:
            provider: Provider identifier (e.g. "openai", "anthropic").
            tokens:   Estimated input+output tokens for this call.  Pass 0
                      to skip token-bucket throttling (requests-only mode).
        """
        state = self._get_state(provider)
        m = state.metrics

        # 1. Wait out any active 429 backoff
        now = time.monotonic()
        if state.backoff_until > now:
            sleep_for = state.backoff_until - now
            logger.info(
                "ProviderRateLimiter[%s]: 429 backoff — sleeping %.2fs",
                provider,
                sleep_for,
            )
            m.total_backoff_seconds += sleep_for
            await asyncio.sleep(sleep_for)

        # 2. Request-bucket (1 request per call)
        if state.config.requests_per_minute > 0:
            await state.req_bucket.acquire(1.0, m)

        # 3. Token-bucket (optional)
        if tokens > 0 and state.config.tokens_per_minute > 0:
            await state.tok_bucket.acquire(float(tokens), m)

        m.requests_allowed += 1
        try:
            yield
        finally:
            pass  # no cleanup needed

    # ------------------------------------------------------------------
    # 429 handling
    # ------------------------------------------------------------------

    async def record_429(
        self,
        provider: str,
        retry_after: Optional[float] = None,
    ) -> None:
        """
        Record a 429 response from *provider* and schedule a backoff.

        Uses exponential backoff with jitter unless the server supplied a
        ``Retry-After`` value.

        Args:
            provider:    Provider that returned 429.
            retry_after: Seconds from server's Retry-After header (optional).
        """
        state = self._get_state(provider)
        m = state.metrics
        m.backoffs_triggered += 1
        state.consecutive_429s += 1

        if retry_after and retry_after > 0:
            backoff = retry_after
        else:
            # Exponential backoff with ±25% jitter
            attempt = min(state.consecutive_429s - 1, 8)  # cap exponent
            raw = _BACKOFF_BASE * (2 ** attempt)
            jitter = raw * _JITTER_FRACTION * (2 * random.random() - 1)
            backoff = min(raw + jitter, _MAX_BACKOFF_SECONDS)

        backoff = max(backoff, 0.0)
        state.backoff_until = time.monotonic() + backoff
        logger.warning(
            "ProviderRateLimiter[%s]: 429 received (attempt %d) — "
            "backoff %.2fs",
            provider,
            state.consecutive_429s,
            backoff,
        )

    def record_success(self, provider: str) -> None:
        """Reset the consecutive-429 counter and clear active backoff after a successful call.

        B-024: Also clears backoff_until so subsequent requests don't
        needlessly wait after a successful recovery.
        """
        if provider in self._providers:
            state = self._providers[provider]
            state.consecutive_429s = 0
            state.backoff_until = 0.0  # B-024: clear active backoff

    # ------------------------------------------------------------------
    # Metrics & introspection
    # ------------------------------------------------------------------

    def metrics(self, provider: str) -> RateLimitMetrics:
        """Return a copy of the current metrics for *provider*."""
        state = self._get_state(provider)
        m = state.metrics
        # Return a snapshot copy
        return RateLimitMetrics(
            requests_allowed=m.requests_allowed,
            requests_throttled=m.requests_throttled,
            total_wait_seconds=m.total_wait_seconds,
            backoffs_triggered=m.backoffs_triggered,
            total_backoff_seconds=m.total_backoff_seconds,
        )

    def all_metrics(self) -> Dict[str, RateLimitMetrics]:
        """Return a dict of metrics snapshots for all configured providers."""
        return {p: self.metrics(p) for p in self._providers}

    def status(self) -> Dict[str, Any]:
        """Return a human-readable status dict for all providers."""
        out: Dict[str, Any] = {}
        now = time.monotonic()
        for name, state in self._providers.items():
            m = state.metrics
            backoff_remaining = max(0.0, state.backoff_until - now)
            out[name] = {
                "rpm": state.config.requests_per_minute,
                "tpm": state.config.tokens_per_minute,
                "requests_allowed": m.requests_allowed,
                "requests_throttled": m.requests_throttled,
                "total_wait_s": round(m.total_wait_seconds, 3),
                "backoffs_triggered": m.backoffs_triggered,
                "consecutive_429s": state.consecutive_429s,
                "backoff_remaining_s": round(backoff_remaining, 3),
            }
        return out

    def reset(self, provider: str) -> None:
        """Remove cached state for *provider* (forces re-init on next use)."""
        self._providers.pop(provider, None)

    def reset_all(self) -> None:
        """Clear all provider state."""
        self._providers.clear()
