from typing import Optional
import asyncio
from collections import deque

__all__ = [
    "Lock",
    "RLock",
    "Waiter",
]


class Lock(asyncio.Lock):
    def acquire_nowait(self) -> bool:
        """Try to acquire the lock without waiting. Return True if successful
        and False if the lock is locked and must be waited for."""
        if not self._locked and (
            self._waiters is None or all(w.cancelled() for w in self._waiters)
        ):
            self._locked = True
            return True
        return False


class RLock(Lock):
    """Re-entrant asyncio lock"""

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._reentrances: int = 0
        self._owner_task: Optional[asyncio.Task] = None

    def acquire_nowait(self) -> bool:
        task = asyncio.current_task()
        if task is self._owner_task:
            self._reentrances += 1
            return True
        result = super().acquire_nowait()
        if result:
            self._owner_task = task
        return result

    acquire_nowait.__doc__ = Lock.acquire_nowait.__doc__

    async def acquire(self) -> bool:
        task = asyncio.current_task()
        if task is self._owner_task:
            self._reentrances += 1
            return True
        result = await super().acquire()
        if result:
            self._owner_task = task
        return result

    def release(self) -> None:
        task = asyncio.current_task()
        if task is self._owner_task:
            if self._reentrances > 0:
                self._reentrances -= 1
                return  # one more release is needed after reaching 0
        else:
            raise RuntimeError(
                f"{type(self).__name__} cannot be released by a task "
                "which has not first acquired it."
            )
        super().release()
        self._owner_task = None


class Waiter:
    """Simple sync across multiple tasks without the risk of forgetting
    to clear the Event flag after waiting."""

    def __init__(self) -> None:
        self._waiters = deque[asyncio.Future]()

    async def wait(self) -> None:
        loop = asyncio.get_running_loop()
        fut = loop.create_future()
        self._waiters.append(fut)
        try:
            await fut
        finally:
            self._waiters.remove(fut)

    def notify(self) -> None:
        if not self._waiters:
            return
        for fut in self._waiters:
            if fut.done():
                continue
            fut.set_result(True)
