"""
Provide utilities for concurrent programming.
"""

import asyncio
import threading
import time
from abc import ABC, abstractmethod
from asyncio import sleep
from math import floor
from threading import Lock
from types import TracebackType
from typing import Self, final

from typing_extensions import override

from betty.asyncio import gather


class _Lock(ABC):
    """
    Provide an asynchronous lock.
    """

    async def __aenter__(self):
        await self.acquire()

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        self.release()

    @abstractmethod
    async def acquire(self, *, wait: bool = True) -> bool:
        """
        Acquire the lock.
        """
        pass

    @abstractmethod
    def release(self) -> None:
        """
        Release the lock.
        """
        pass


async def asynchronize_acquire(lock: Lock, *, wait: bool = True) -> bool:
    """
    Acquire a synchronous lock asynchronously.
    """
    while not lock.acquire(blocking=False):
        if not wait:
            return False
        # Sleeping for zero seconds does not actually sleep, but gives the event
        # loop a chance to progress other tasks while we wait for another chance
        # to acquire the lock.
        await sleep(0)
    return True


@final
class AsynchronizedLock(_Lock):
    """
    Make a sychronous (blocking) lock asynchronous (non-blocking).
    """

    __slots__ = "_lock"

    def __init__(self, lock: Lock):
        self._lock = lock

    @override
    async def acquire(self, *, wait: bool = True) -> bool:
        return await asynchronize_acquire(self._lock, wait=wait)

    @override
    def release(self) -> None:
        self._lock.release()

    @classmethod
    def threading(cls) -> Self:
        """
        Create a new thread-safe, asynchronous lock.
        """
        return cls(threading.Lock())


@final
class MultiLock(_Lock):
    """
    Provide a lock that only acquires if all of the given locks can be acquired.
    """

    __slots__ = "_locked", "_locks"

    def __init__(self, *locks: _Lock):
        self._locks = locks
        self._locked = False

    @override
    async def acquire(self, *, wait: bool = True) -> bool:
        acquisitions = await gather(*(lock.acquire(wait=wait) for lock in self._locks))
        # We require all locks to be acquired, or none at all
        # If one or more fail, release the others.
        if False in acquisitions:
            for lock, acquisition in zip(self._locks, acquisitions):
                if acquisition:
                    lock.release()
            return False
        self._locked = True
        return True

    @override
    def release(self) -> None:
        self._locked = False
        for lock in self._locks:
            lock.release()


@final
class RateLimiter:
    """
    Rate-limit operations.

    This class implements the `Token Bucket algorithm <https://en.wikipedia.org/wiki/Token_bucket>`_.

    This class is thread-safe.
    """

    _PERIOD = 1

    def __init__(self, maximum: int):
        self._lock = AsynchronizedLock.threading()
        self._maximum = maximum
        self._available: int | float = maximum
        # A Token Bucket fills as time passes. However, we want callers to be able to start
        # using the limiter immediately, so we 'preload' the first's period's tokens, and
        # set the last added time to the end of the first period. This ensures there is no
        # needless waiting if the number of tokens consumed in total is less than the limit
        # per period.
        self._last_add = time.monotonic() + self._PERIOD

    def _add_tokens(self):
        now = time.monotonic()
        elapsed = now - self._last_add
        added = elapsed * self._maximum
        possibly_available = floor(self._available + added)
        if possibly_available > 0:
            self._available = min(possibly_available, self._maximum)
            self._last_add = now

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

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        return

    async def wait(self) -> None:
        """
        Wait until an operation may be performed (again).
        """
        async with self._lock:
            while self._available < 1:
                self._add_tokens()
                if self._available < 1:
                    await asyncio.sleep(0.1)
            self._available -= 1
