"""
Utility module providing request‑rate control for the agent.

This file implements a simple throttling layer for OpenAI API calls, combining
a token‑bucket to pace outgoing requests and a reactive backoff that interprets
retry hints from 429 errors. It helps the agent run continuously without
tripping rate limits during long investigations.
"""


import re
from typing import Optional
import time
import random
from dataclasses import dataclass
from openai import RateLimitError, APIError


@dataclass
class BackoffConfig:
    initial_delay: float = 1.0      # seconds
    max_delay: float = 60.0         # seconds
    jitter: float = 0.25            # +/- 25% jitter


class TokenBucket:
    """
    Classic token bucket implementation.
    rate: tokens added per second
    capacity: maximum tokens in the bucket
    """
    def __init__(self, rate: float, capacity: float):
        self.rate = rate
        self.capacity = capacity
        self.tokens = capacity
        self.timestamp = time.time()

    def consume(self, tokens: float = 1.0) -> bool:
        now = time.time()
        elapsed = now - self.timestamp
        self.timestamp = now

        # Refill tokens
        self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)

        if self.tokens >= tokens:
            self.tokens -= tokens
            return True
        return False


class Throttler:
    """
    Rate‑limits API calls using two layers of protection:
    a token‑bucket for proactive pacing, and an exponential backoff
    that reacts to server‑side 429 errors. If the API returns a
    retry hint (e.g., “Please try again in 2.3s”), that delay is
    extracted and used directly before retrying the request.
    """
    def __init__(
        self,
        rate_per_minute: float,
        bucket_capacity: Optional[float] = None,
        backoff: BackoffConfig = BackoffConfig(),
    ):
        # Convert TPM to tokens per second
        rate_per_second = rate_per_minute / 60.0

        if bucket_capacity is None:
            bucket_capacity = rate_per_minute  # allow 1 minute of burst

        self.bucket = TokenBucket(rate=rate_per_second, capacity=bucket_capacity)
        self.backoff = backoff
        self.current_delay = backoff.initial_delay

    # ----------------------------------------------------------------------
    # Proactive throttling
    # ----------------------------------------------------------------------
    def wait_for_token(self):
        """Block until a token is available."""
        while not self.bucket.consume():
            print("...Throttling. Proactive")
            time.sleep(0.01)

    # ----------------------------------------------------------------------
    # Reactive throttling
    # ----------------------------------------------------------------------
    def backoff_wait(self, delay=None):
        """Wait using exponential backoff with jitter. If the actual delay needed is known it can be supplied."""
        if delay is None:
            jitter_factor = 1 + random.uniform(-self.backoff.jitter, self.backoff.jitter)
            delay = self.current_delay * jitter_factor

        print(f"...Throttling {delay}s. Reactive")
        time.sleep(delay)

        # Increase delay for next time
        if delay is None:
            self.current_delay = min(self.current_delay * 2, self.backoff.max_delay)

    def reset_backoff(self):
        """Reset backoff after a successful request."""
        self.current_delay = self.backoff.initial_delay



    def parse_retry_after(self, message: str) -> Optional[float]:
        """
        Extracts the floating‑point wait time from an OpenAI rate‑limit message.
        Returns the number of seconds as a float, or None if not found.
        """
        pattern = r"Please try again in ([0-9]*\.?[0-9]+)s"
        match = re.search(pattern, message)
        if match:
            return float(match.group(1))
        return None

    # ----------------------------------------------------------------------
    # Wrapper for API calls
    # ----------------------------------------------------------------------
    def run(self, func, *args, **kwargs):
        """
        Execute a function with throttling.
        - Waits for token
        - Runs the function
        - If RateLimitError occurs, applies backoff and retries
        """
        while True:
            # Proactive throttle
            self.wait_for_token()

            try:
                result = func(*args, **kwargs)
                self.reset_backoff()
                return result

            except RateLimitError as e:
                delay = self.parse_retry_after(str(e))
                self.backoff_wait(delay)
                continue

            # there might be other exceptions that trip for 429. Catch them here.
            except APIError as e:
                if e.status_code == 429: # type: ignore
                    delay = self.parse_retry_after(str(e))
                    self.backoff_wait(delay)
                    continue
                raise