# utils.py
import time
import random
from functools import wraps
from google.api_core import exceptions
import gradio as gr

def retry_with_exponential_backoff(
    func,
    initial_delay: float = 2,
    exponential_base: float = 2,
    jitter: bool = True,
    max_retries: int = 5,
):
    """
    A decorator to retry a function with exponential backoff for API calls.
    It specifically catches google.api_core.exceptions.ResourceExhausted.
    """

    @wraps(func)
    def wrapper(*args, **kwargs):
        num_retries = 0
        delay = initial_delay

        while True:
            try:
                return func(*args, **kwargs)
            except exceptions.ResourceExhausted as e:
                num_retries += 1
                if num_retries > max_retries:
                    # Using gr.Error to raise a user-facing error in the Gradio UI
                    raise gr.Error(
                        f"Maximum number of retries ({max_retries}) exceeded. The API is still busy. Please try again later."
                    ) from e

                if jitter:
                    delay *= exponential_base * (1 + random.random())
                else:
                    delay *= exponential_base

                # It's helpful to print the delay to the console for debugging
                print(f"Rate limit exceeded. Retrying in {delay:.2f} seconds...")
                time.sleep(delay)
            except Exception as e:
                # Raise other exceptions immediately
                raise gr.Error(f"An unexpected error occurred: {e}")

    return wrapper