"""Rate limiter for LLM API calls."""

from __future__ import annotations

import logging
import time
from collections import deque
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)


@dataclass
class RateLimiter:
    """Token bucket rate limiter for LLM API calls.

    Limits both requests per minute (RPM) and tokens per minute (TPM).
    """

    max_requests_per_minute: int = 60
    max_tokens_per_minute: int = 100_000
    _request_timestamps: deque = field(default_factory=deque)
    _token_counts: deque = field(default_factory=deque)

    def _cleanup_old_entries(self) -> None:
        """Remove entries older than 60 seconds."""
        cutoff = time.monotonic() - 60.0
        while self._request_timestamps and self._request_timestamps[0] < cutoff:
            self._request_timestamps.popleft()
        while self._token_counts and self._token_counts[0][0] < cutoff:
            self._token_counts.popleft()

    @property
    def current_rpm(self) -> int:
        """Current requests in the last minute."""
        self._cleanup_old_entries()
        return len(self._request_timestamps)

    @property
    def current_tpm(self) -> int:
        """Current tokens used in the last minute."""
        self._cleanup_old_entries()
        return sum(count for _, count in self._token_counts)

    def wait_if_needed(self, estimated_tokens: int = 0) -> float:
        """Wait if rate limit would be exceeded. Returns seconds waited."""
        self._cleanup_old_entries()
        total_wait = 0.0

        # Check RPM
        if len(self._request_timestamps) >= self.max_requests_per_minute:
            oldest = self._request_timestamps[0]
            wait_time = 60.0 - (time.monotonic() - oldest) + 0.1
            if wait_time > 0:
                logger.info("Rate limit: waiting %.1fs (RPM limit reached)", wait_time)
                time.sleep(wait_time)
                total_wait += wait_time
                self._cleanup_old_entries()

        # Check TPM
        current_tokens = sum(count for _, count in self._token_counts)
        if current_tokens + estimated_tokens > self.max_tokens_per_minute and self._token_counts:
            oldest = self._token_counts[0][0]
            wait_time = 60.0 - (time.monotonic() - oldest) + 0.1
            if wait_time > 0:
                logger.info("Rate limit: waiting %.1fs (TPM limit reached)", wait_time)
                time.sleep(wait_time)
                total_wait += wait_time
                self._cleanup_old_entries()

        return total_wait

    def record_request(self, tokens_used: int = 0) -> None:
        """Record a completed request."""
        now = time.monotonic()
        self._request_timestamps.append(now)
        if tokens_used > 0:
            self._token_counts.append((now, tokens_used))

    def can_proceed(self, estimated_tokens: int = 0) -> bool:
        """Check if a request can proceed without waiting."""
        self._cleanup_old_entries()
        if len(self._request_timestamps) >= self.max_requests_per_minute:
            return False
        current_tokens = sum(count for _, count in self._token_counts)
        return current_tokens + estimated_tokens <= self.max_tokens_per_minute


# Global rate limiter instance (shared across all LLM calls)
_global_limiter: RateLimiter | None = None


def get_rate_limiter(
    max_rpm: int = 60,
    max_tpm: int = 100_000,
) -> RateLimiter:
    """Get or create the global rate limiter."""
    global _global_limiter
    if _global_limiter is None:
        _global_limiter = RateLimiter(
            max_requests_per_minute=max_rpm,
            max_tokens_per_minute=max_tpm,
        )
    return _global_limiter


def reset_rate_limiter() -> None:
    """Reset the global rate limiter (useful for testing)."""
    global _global_limiter
    _global_limiter = None
