import asyncio
from contextlib import AbstractAsyncContextManager
from datetime import datetime, timedelta
from logging import getLogger
from typing import Any, Awaitable, TypeVar, Union

from typing_extensions import Never, deprecated, override

_background_tasks: set[asyncio.Task[Any]] = set()

log = getLogger(__name__)


class _SleepUntil(object):
    def __init__(self, dt: datetime | timedelta):
        self._dt: datetime = datetime.fromtimestamp(0)
        self.update(dt)
        self._lock = asyncio.Lock()

    def update(self, dt: datetime | timedelta) -> None:
        if isinstance(dt, timedelta):
            dt = datetime.now() + dt

        if dt > self._dt:
            self._dt = dt

    async def __call__(self) -> None:
        async with self._lock:
            while (delta := self._dt - datetime.now()) > timedelta(seconds=0):
                await asyncio.sleep(delta.total_seconds())


T = TypeVar("T")


class Throttler(object):
    """
    BETA: Provides throttling using the Leaky Bucket Algorithm.

    https://en.wikipedia.org/wiki/Leaky_bucket

    It can be imagined like so:

    There are a number of buckets, full of water. There is a tap, that can fill
    a bucket in (refresh) seconds. For a task to start, it must flip a filled
    bucket. It will keep the bucket overturned until it completes. When it
    completes, the bucket will wait to be filled so that another task can use
    it.
    """

    def __init__(
        self,
        buckets: int,
        refresh: Union[float, timedelta],
        stagger: Union[float, timedelta] = timedelta(0.0),
    ):
        """
        Create a throttle object.

        :param buckets: Number of buckets (concurrent tasks).
        :param refresh: Time (seconds) taken to fill a bucket.
        :param stagger: Stagger task start times.
        """
        if isinstance(refresh, timedelta):
            self._restore_time = refresh
        else:
            self._restore_time = timedelta(seconds=refresh)

        if isinstance(stagger, timedelta):
            self._stagger = stagger
        else:
            self._stagger = timedelta(seconds=stagger)

        self._stagger_wait = _SleepUntil(timedelta(seconds=0))

        self._filler_task = asyncio.create_task(self._filler())
        _background_tasks.add(self._filler_task)  # For GC
        self._filler_task.add_done_callback(_background_tasks.discard)

        self._buckets_total = buckets
        self._full_buckets = asyncio.Queue[None]()
        for i in range(self._buckets_total):
            self._full_buckets.put_nowait(None)

        self._empty_buckets = asyncio.Queue[None]()

    def __del__(self) -> None:
        self._filler_task.cancel()

    async def _filler(self) -> Never:
        while True:
            await self._empty_buckets.get()
            await asyncio.sleep(self._restore_time.total_seconds())
            self._full_buckets.put_nowait(None)

    async def bucket(self) -> "Bucket":
        await self._stagger_wait()
        await self._full_buckets.get()
        self._stagger_wait.update(self._stagger)

        return Bucket(self)

    @deprecated("Use .bucket")
    async def wrap(self, co: Awaitable[T]) -> T:
        async with self:
            return await co

    @deprecated("Use .bucket")
    async def __aenter__(self) -> None:
        await self._full_buckets.get()

    @deprecated("Use .bucket")
    async def __aexit__(self, exc_type, exc_value, traceback):
        self._empty_buckets.put_nowait(None)


class Bucket(AbstractAsyncContextManager[None]):
    def __init__(self, throttler: Throttler):
        self._used: bool = False
        self._throttler: Throttler = throttler

    async def __call__(self, co: Awaitable[T]) -> T:
        async with self:
            return await co

    @override
    async def __aenter__(self) -> None:
        assert self._used is False

    @override
    async def __aexit__(self, exc_type, exc_value, traceback):
        self._throttler._empty_buckets.put_nowait(None)
        self._used = True

    def __del__(self) -> None:
        if self._used is False:
            log.warn("Deleting unused bucket")
