import time
from smolagents import CodeAgent, ActionStep


class RateLimiter:
    """
    Counts API calls and sleeps for a fixed number of seconds once a
    call threshold is reached, then resets the counter.
    """
    def __init__(self, calls_per_wait: int = 10, seconds_to_wait: int = 60):
        self.calls_per_wait = calls_per_wait
        self.seconds_to_wait = seconds_to_wait
        self._call_count = 0

    def increment_and_sleep_if_needed(self, memory_step: ActionStep, agent: CodeAgent):
        """
        Increments the call counter, and once it reaches the configured
        threshold, sleeps for the configured duration and resets the counter.
        """
        self._call_count += 1
        if self._call_count >= self.calls_per_wait:
            print(
                f"Call count reached {self.calls_per_wait}, "
                f"sleeping for {self.seconds_to_wait} seconds"
            )
            time.sleep(self.seconds_to_wait)
            self._call_count = 0