"""
Starvex Middleware - Rate limiting and cost control for AI guardrail checks

Provides middleware components for controlling API usage:
- Token bucket rate limiting per API key
- Budget tracking with request count and cost caps
- Combined guard middleware for pre/post check enforcement

Example:
    from starvex.observability import RateLimiter, BudgetTracker, GuardMiddleware

    # Standalone rate limiter
    limiter = RateLimiter(rate=10.0, burst=20)
    if limiter.check_rate_limit("sv_live_abc"):
        # proceed with request
        ...

    # Combined middleware
    middleware = GuardMiddleware(rate=10.0, burst=20)
    middleware.budget_tracker.set_budget("sv_live_abc", max_requests=1000, max_cost=50.0)

    allowed, reason = middleware.before_check("sv_live_abc")
    if allowed:
        # run guardrail check, then record outcome
        middleware.after_check("sv_live_abc", latency_ms=42.5, blocked=False)
"""

import logging
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Dict, Optional, Tuple

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Token-bucket rate limiter
# ---------------------------------------------------------------------------


@dataclass
class _TokenBucket:
    """Internal token bucket state for a single API key."""

    tokens: float
    max_tokens: float
    refill_rate: float  # tokens per second
    last_refill: float = field(default_factory=time.monotonic)


class RateLimiter:
    """
    Per-API-key rate limiter using the token bucket algorithm.

    Each API key gets its own bucket that refills at *rate* tokens per second
    up to a maximum of *burst* tokens.  A single request consumes one token.

    The implementation is thread-safe; concurrent calls from different threads
    will never corrupt internal state.

    Args:
        rate: Sustained request rate (requests per second).
        burst: Maximum burst size (peak requests before throttling).
    """

    def __init__(self, rate: float = 10.0, burst: int = 20) -> None:
        if rate <= 0:
            raise ValueError("rate must be positive")
        if burst <= 0:
            raise ValueError("burst must be positive")

        self._rate = rate
        self._burst = burst
        self._buckets: Dict[str, _TokenBucket] = {}
        self._lock = threading.Lock()

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

    def check_rate_limit(self, api_key: str) -> bool:
        """
        Attempt to consume one token for *api_key*.

        Returns ``True`` if the request is allowed (a token was available),
        ``False`` if the caller should be throttled.
        """
        with self._lock:
            bucket = self._get_or_create_bucket(api_key)
            self._refill(bucket)

            if bucket.tokens >= 1.0:
                bucket.tokens -= 1.0
                return True

            logger.debug("Rate limit exceeded for key %.8s...", api_key)
            return False

    def get_usage(self, api_key: str) -> Dict[str, Any]:
        """
        Return current rate-limiter state for *api_key*.

        Keys in the returned dict:
        - ``tokens_remaining``: float, tokens currently available
        - ``max_tokens``: float, burst capacity
        - ``refill_rate``: float, tokens added per second
        - ``is_throttled``: bool, ``True`` when no tokens are available
        """
        with self._lock:
            bucket = self._get_or_create_bucket(api_key)
            self._refill(bucket)

            return {
                "tokens_remaining": round(bucket.tokens, 4),
                "max_tokens": bucket.max_tokens,
                "refill_rate": bucket.refill_rate,
                "is_throttled": bucket.tokens < 1.0,
            }

    # -- internals ----------------------------------------------------------

    def _get_or_create_bucket(self, api_key: str) -> _TokenBucket:
        """Return the bucket for *api_key*, creating one if necessary."""
        if api_key not in self._buckets:
            self._buckets[api_key] = _TokenBucket(
                tokens=float(self._burst),
                max_tokens=float(self._burst),
                refill_rate=self._rate,
            )
        return self._buckets[api_key]

    @staticmethod
    def _refill(bucket: _TokenBucket) -> None:
        """Add tokens that have accumulated since the last refill."""
        now = time.monotonic()
        elapsed = now - bucket.last_refill
        if elapsed > 0:
            bucket.tokens = min(
                bucket.max_tokens,
                bucket.tokens + elapsed * bucket.refill_rate,
            )
            bucket.last_refill = now


# ---------------------------------------------------------------------------
# Budget tracker
# ---------------------------------------------------------------------------


@dataclass
class _BudgetEntry:
    """Internal per-key budget and usage counters."""

    max_requests: int = 0
    max_cost: float = 0.0
    requests_used: int = 0
    cost_used: float = 0.0


class BudgetTracker:
    """
    Per-API-key budget tracker for request counts and estimated cost.

    Call :meth:`set_budget` to define spending limits for a key, then use
    :meth:`record_usage` after each guardrail check to accumulate spend.
    :meth:`check_budget` returns ``False`` once either limit is reached.

    The tracker is thread-safe.
    """

    def __init__(self) -> None:
        self._budgets: Dict[str, _BudgetEntry] = {}
        self._lock = threading.Lock()

    # -- configuration ------------------------------------------------------

    def set_budget(
        self,
        api_key: str,
        max_requests: int = 0,
        max_cost: float = 0.0,
    ) -> None:
        """
        Define (or update) the budget for *api_key*.

        Args:
            api_key: The API key to configure.
            max_requests: Maximum allowed request count.  0 means unlimited.
            max_cost: Maximum allowed accumulated cost.  0.0 means unlimited.
        """
        with self._lock:
            entry = self._get_or_create(api_key)
            entry.max_requests = max_requests
            entry.max_cost = max_cost
            logger.debug(
                "Budget set for %.8s...: max_requests=%d, max_cost=%.4f",
                api_key,
                max_requests,
                max_cost,
            )

    # -- usage recording ----------------------------------------------------

    def record_usage(self, api_key: str, cost: float = 0.0) -> None:
        """
        Record one request and its associated *cost* for *api_key*.

        Args:
            api_key: The API key that made the request.
            cost: Estimated dollar cost of this request.
        """
        with self._lock:
            entry = self._get_or_create(api_key)
            entry.requests_used += 1
            entry.cost_used += cost

    # -- budget checks ------------------------------------------------------

    def check_budget(self, api_key: str) -> bool:
        """
        Return ``True`` if *api_key* is still within its budget.

        A key with no explicit budget (never passed to :meth:`set_budget`)
        is always considered within budget.

        Returns ``False`` when either the request count or the cost cap has
        been reached or exceeded.
        """
        with self._lock:
            entry = self._get_or_create(api_key)

            if entry.max_requests > 0 and entry.requests_used >= entry.max_requests:
                logger.debug(
                    "Request budget exhausted for %.8s... (%d/%d)",
                    api_key,
                    entry.requests_used,
                    entry.max_requests,
                )
                return False

            if entry.max_cost > 0 and entry.cost_used >= entry.max_cost:
                logger.debug(
                    "Cost budget exhausted for %.8s... (%.4f/%.4f)",
                    api_key,
                    entry.cost_used,
                    entry.max_cost,
                )
                return False

            return True

    # -- reporting ----------------------------------------------------------

    def get_usage_report(self, api_key: str) -> Dict[str, Any]:
        """
        Return a snapshot of budget and usage for *api_key*.

        Keys in the returned dict:
        - ``requests_used``: int
        - ``max_requests``: int (0 = unlimited)
        - ``cost_used``: float
        - ``max_cost``: float (0.0 = unlimited)
        - ``requests_remaining``: int or ``None`` if unlimited
        - ``cost_remaining``: float or ``None`` if unlimited
        - ``over_budget``: bool
        """
        with self._lock:
            entry = self._get_or_create(api_key)

            requests_remaining: Optional[int] = None
            if entry.max_requests > 0:
                requests_remaining = max(0, entry.max_requests - entry.requests_used)

            cost_remaining: Optional[float] = None
            if entry.max_cost > 0:
                cost_remaining = round(max(0.0, entry.max_cost - entry.cost_used), 6)

            over_budget = (
                (entry.max_requests > 0 and entry.requests_used >= entry.max_requests)
                or (entry.max_cost > 0 and entry.cost_used >= entry.max_cost)
            )

            return {
                "requests_used": entry.requests_used,
                "max_requests": entry.max_requests,
                "cost_used": round(entry.cost_used, 6),
                "max_cost": entry.max_cost,
                "requests_remaining": requests_remaining,
                "cost_remaining": cost_remaining,
                "over_budget": over_budget,
            }

    # -- internals ----------------------------------------------------------

    def _get_or_create(self, api_key: str) -> _BudgetEntry:
        if api_key not in self._budgets:
            self._budgets[api_key] = _BudgetEntry()
        return self._budgets[api_key]


# ---------------------------------------------------------------------------
# Combined guard middleware
# ---------------------------------------------------------------------------


class GuardMiddleware:
    """
    Middleware that combines rate limiting and budget control.

    Use :meth:`before_check` before running a guardrail check and
    :meth:`after_check` afterwards to record metrics.  :meth:`get_metrics`
    returns a comprehensive view of all collected data.

    Args:
        rate: Sustained request rate (requests per second) for the rate limiter.
        burst: Maximum burst size for the rate limiter.
        default_cost: Cost attributed to each request when no explicit cost
            is given to :meth:`after_check`.
    """

    def __init__(
        self,
        rate: float = 10.0,
        burst: int = 20,
        default_cost: float = 0.001,
    ) -> None:
        self.rate_limiter = RateLimiter(rate=rate, burst=burst)
        self.budget_tracker = BudgetTracker()
        self._default_cost = default_cost

        # Per-key metrics accumulators (protected by _lock)
        self._metrics: Dict[str, Dict[str, Any]] = {}
        self._lock = threading.Lock()

    # -- pre-check ----------------------------------------------------------

    def before_check(self, api_key: str) -> Tuple[bool, str]:
        """
        Evaluate whether a guardrail check should proceed for *api_key*.

        Returns:
            A ``(allowed, reason)`` tuple.  *allowed* is ``True`` when the
            request may proceed; *reason* is an empty string on success or a
            human-readable explanation on denial.
        """
        # Rate limit comes first -- it is the cheapest check.
        if not self.rate_limiter.check_rate_limit(api_key):
            self._record_denied(api_key, "rate_limited")
            return False, "Rate limit exceeded"

        if not self.budget_tracker.check_budget(api_key):
            self._record_denied(api_key, "over_budget")
            return False, "Budget limit exceeded"

        return True, ""

    # -- post-check ---------------------------------------------------------

    def after_check(
        self,
        api_key: str,
        latency_ms: float = 0.0,
        blocked: bool = False,
        cost: Optional[float] = None,
    ) -> None:
        """
        Record the outcome of a guardrail check for *api_key*.

        Args:
            api_key: The API key that made the request.
            latency_ms: End-to-end latency of the check in milliseconds.
            blocked: Whether the check resulted in the request being blocked.
            cost: Estimated cost.  Falls back to *default_cost* if ``None``.
        """
        actual_cost = cost if cost is not None else self._default_cost
        self.budget_tracker.record_usage(api_key, actual_cost)

        with self._lock:
            m = self._get_or_create_metrics(api_key)
            m["total_checks"] += 1
            m["total_latency_ms"] += latency_ms
            m["latencies"].append(latency_ms)

            if blocked:
                m["blocked_checks"] += 1
            else:
                m["passed_checks"] += 1

    # -- reporting ----------------------------------------------------------

    def get_metrics(self, api_key: str) -> Dict[str, Any]:
        """
        Return comprehensive usage metrics for *api_key*.

        The returned dict merges rate limiter state, budget report, and
        check-level statistics (counts, latency percentiles, etc.).
        """
        with self._lock:
            m = self._get_or_create_metrics(api_key)

            total = m["total_checks"]
            avg_latency = m["total_latency_ms"] / total if total > 0 else 0.0

            latencies = sorted(m["latencies"])
            p50 = self._percentile(latencies, 50)
            p95 = self._percentile(latencies, 95)
            p99 = self._percentile(latencies, 99)

            check_stats: Dict[str, Any] = {
                "total_checks": total,
                "passed_checks": m["passed_checks"],
                "blocked_checks": m["blocked_checks"],
                "denied_checks": m["denied_checks"],
                "denied_reasons": dict(m["denied_reasons"]),
                "avg_latency_ms": round(avg_latency, 4),
                "p50_latency_ms": round(p50, 4),
                "p95_latency_ms": round(p95, 4),
                "p99_latency_ms": round(p99, 4),
            }

        # Merge in rate limiter and budget data (each acquires its own lock)
        rate_info = self.rate_limiter.get_usage(api_key)
        budget_info = self.budget_tracker.get_usage_report(api_key)

        return {
            "rate_limiter": rate_info,
            "budget": budget_info,
            "checks": check_stats,
        }

    # -- internals ----------------------------------------------------------

    def _record_denied(self, api_key: str, reason: str) -> None:
        """Increment denial counters under the lock."""
        with self._lock:
            m = self._get_or_create_metrics(api_key)
            m["denied_checks"] += 1
            m["denied_reasons"][reason] = m["denied_reasons"].get(reason, 0) + 1

    def _get_or_create_metrics(self, api_key: str) -> Dict[str, Any]:
        if api_key not in self._metrics:
            self._metrics[api_key] = {
                "total_checks": 0,
                "passed_checks": 0,
                "blocked_checks": 0,
                "denied_checks": 0,
                "denied_reasons": {},
                "total_latency_ms": 0.0,
                "latencies": [],
            }
        return self._metrics[api_key]

    @staticmethod
    def _percentile(sorted_values: list, pct: float) -> float:
        """Return the *pct*-th percentile from a **sorted** list of numbers."""
        if not sorted_values:
            return 0.0
        k = (len(sorted_values) - 1) * (pct / 100.0)
        f = int(k)
        c = f + 1
        if c >= len(sorted_values):
            return float(sorted_values[f])
        d = k - f
        return float(sorted_values[f]) + d * (float(sorted_values[c]) - float(sorted_values[f]))
