import asyncio
from datetime import datetime, timedelta
from typing import Union


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.

    Note that `refresh = 0` is just a semaphore.

    :param pod_id: Pod ID
    :param state: PC Pod State
    """

    _restore_time: timedelta
    _restore_last: datetime
    _buckets_total: float
    _buckets_taken: float
    _buckets_pending: float
    _buckets_are_pending: asyncio.Event
    _lock: asyncio.Lock

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

        :param buckets: Number of buckets (concurrent tasks).
        :param refresh: Time (seconds) taken to fill a bucket.
        """

        if isinstance(refresh, timedelta):
            self._restore_time = refresh
        else:
            self._restore_time = timedelta(seconds=refresh)
        self._restore_last = datetime.utcnow()
        self._buckets_total = buckets
        self._buckets_taken = 0.0
        self._buckets_pending = 0.0
        self._buckets_are_pending = asyncio.Event()
        self._buckets_are_pending.set()
        self._lock = asyncio.Lock()

    def _refresh(self):
        now = datetime.utcnow()
        restored_buckets = (now - self._restore_last) / self._restore_time
        self._buckets_pending -= restored_buckets
        self._buckets_pending = max(0, self._buckets_pending)
        self._restore_last = now
        if self._buckets_pending + self._buckets_taken < self._buckets_total:
            return 0
        else:
            taken = self._buckets_pending
            time = taken % 1 * self._restore_time
            return time.total_seconds()

    async def _grab_bucket(self):
        async with self._lock:
            if self._buckets_pending == 0:
                await self._buckets_are_pending.wait()
                self._buckets_are_pending.clear()
            wait_time = self._refresh()
            await asyncio.sleep(wait_time)
            self._buckets_taken += 1

    def _release_bucket(self):
        self._refresh()
        self._buckets_taken -= 1
        self._buckets_pending += 1
        self._buckets_are_pending.set()

    async def wrap(self, co):
        try:
            await self._grab_bucket()
        except asyncio.CancelledError as e:
            await co.throw(e)
        try:
            result = await co
            return result
        finally:
            self._release_bucket()

    async def __aenter__(self) -> None:
        await self._grab_bucket()

    async def __aexit__(self, exc_type, exc_value, traceback):
        self._release_bucket()
