import asyncio
import time
from collections.abc import Callable

from pydantic import BaseModel


class CallableInfo(BaseModel):
    backoff_seconds: int
    last_call: float | None = None
    is_async: bool


class BackoffService(BaseModel):
    """Use this service to implement a backoff strategy on random callables.
    The callable will be called the first time without delay but every subsequent
    call may be hold off for a given time

    Examples:
         >>> def callable(x): print(f"foo {x}")
         >>> backoff_service = BackoffService()
         >>> backoff_service.add_callable(callable, backoff_time=3)
         >>> backoff_service.call(callable, "bar") # prints "foo bar"
         >>> backoff_service.call(callable, "baz") # prints nothing
         >>> # wait 3 seconds before calling callable again
         >>> backoff_service.call(callable, "qux") # prints "foo qux"
    """

    callables: dict[Callable, CallableInfo] = {}

    def add_callable(self, func: Callable, backoff_seconds: int):
        """Adds a callable to the local mapper with its backoff configuration.

        Args:
            func: The function to be registered
            backoff_seconds: The number of seconds to wait between successive calls
        """
        self.callables[func] = CallableInfo(
            backoff_seconds=backoff_seconds, is_async=self._is_async(func)
        )

    @staticmethod
    def _is_async(func: Callable) -> bool:
        """Checks if the callable is async"""
        return asyncio.iscoroutinefunction(func)

    def _can_call(self, func: Callable) -> bool:
        """Determines if enough time has passed since the last call"""
        if func not in self.callables:
            raise ValueError(f"Function {func} not registered with backoff service")

        func_info = self.callables[func]
        last_call = func_info.last_call

        if last_call is None:
            return True

        elapsed = time.time() - last_call
        return elapsed >= func_info.backoff_seconds

    def call(self, func: Callable, *args, **kwargs):
        """Calls the callable with arguments and returns its result if it isn't held off

        Args:
            func: The function to call
            *args: Positional arguments for the function
            **kwargs: Keyword arguments for the function

        Returns:
            Optional[Any]: The result of the function call if executed, None if held off
        """
        if func not in self.callables:
            raise ValueError(f"Function {func} not registered with backoff service")

        if self.callables[func].is_async:
            raise ValueError(
                "Cannot call async function with .call(), use .call_async() instead"
            )

        if not self._can_call(func):
            return None

        result = func(*args, **kwargs)
        self.callables[func].last_call = time.time()
        return result

    async def call_async(self, func: Callable, *args, **kwargs):
        """Same as .call(...) but for async functions

        Args:
            func: The async function to call
            *args: Positional arguments for the function
            **kwargs: Keyword arguments for the function

        Returns:
            Optional[Any]: The result of the async function call if executed, None if held off
        """
        if func not in self.callables:
            raise ValueError(f"Function {func} not registered with backoff service")

        if not self.callables[func].is_async:
            raise ValueError(
                "Cannot call sync function with .call_async(), use .call() instead"
            )

        if not self._can_call(func):
            return None

        result = await func(*args, **kwargs)
        self.callables[func].last_call = time.time()
        return result
