import asyncio
import functools
import inspect
from collections.abc import Callable

from ..exceptions import APIError

DELAY_FACTOR = 5


def running(retry: int = 0) -> Callable:
    """
    Decorator to check if GeminiClient is running before making a request.
    Supports both regular async functions and async generators.

    Parameters
    ----------
    retry: `int`, optional
        Max number of retries when `gemini_webapi.APIError` is raised.
    """

    def decorator(func):
        if inspect.isasyncgenfunction(func):

            @functools.wraps(func)
            async def wrapper(client, *args, current_retry=None, **kwargs):
                if current_retry is None:
                    current_retry = retry

                try:
                    if not client._running:
                        await client.init(
                            timeout=client.timeout,
                            auto_close=client.auto_close,
                            close_delay=client.close_delay,
                            auto_refresh=client.auto_refresh,
                            refresh_interval=client.refresh_interval,
                            verbose=client.verbose,
                            watchdog_timeout=client.watchdog_timeout,
                        )

                    if not client._running:
                        raise APIError(
                            f"Invalid function call: GeminiClient.{func.__name__}. Client initialization failed."
                        )

                    async for item in func(client, *args, **kwargs):
                        yield item
                except APIError:
                    if current_retry > 0:
                        delay = (retry - current_retry + 1) * DELAY_FACTOR
                        await asyncio.sleep(delay)
                        async for item in wrapper(
                            client, *args, current_retry=current_retry - 1, **kwargs
                        ):
                            yield item
                    else:
                        raise

            return wrapper
        else:

            @functools.wraps(func)
            async def wrapper(client, *args, current_retry=None, **kwargs):
                if current_retry is None:
                    current_retry = retry

                try:
                    if not client._running:
                        await client.init(
                            timeout=client.timeout,
                            auto_close=client.auto_close,
                            close_delay=client.close_delay,
                            auto_refresh=client.auto_refresh,
                            refresh_interval=client.refresh_interval,
                            verbose=client.verbose,
                            watchdog_timeout=client.watchdog_timeout,
                        )

                    if not client._running:
                        raise APIError(
                            f"Invalid function call: GeminiClient.{func.__name__}. Client initialization failed."
                        )

                    return await func(client, *args, **kwargs)
                except APIError:
                    if current_retry > 0:
                        delay = (retry - current_retry + 1) * DELAY_FACTOR
                        await asyncio.sleep(delay)

                        return await wrapper(
                            client, *args, current_retry=current_retry - 1, **kwargs
                        )

                    raise

            return wrapper

    return decorator