"""A fake clock for the whole suite, so every test finishes in milliseconds of real time.

FakeClock.now() is the virtual time; FakeClock.sleep(d) parks the caller until the virtual
time reaches now + d. Nothing ever waits for real: a driver task lets every runnable coroutine
run until only sleepers are left, then jumps the clock to the earliest wake-up. Because the
bucket, the client and the transport all take `clock`/`sleep` as parameters, the same code
that runs in production on asyncio.sleep runs here on virtual time — with exact arithmetic.

Tests are plain `def`s that call `run(coro)`; no pytest-asyncio needed.
"""

from __future__ import annotations

import asyncio
import heapq
from collections.abc import Coroutine
from typing import Any

import pytest

# How many loop iterations the driver yields before advancing time. A task reaches its next
# fake sleep within a handful of iterations (a semaphore hand-off is one); this is a margin.
DRAIN = 25


class FakeClock:
    def __init__(self) -> None:
        self.t = 0.0
        self._heap: list[tuple[float, int, asyncio.Future[None]]] = []
        self._seq = 0
        self._driver: asyncio.Task[None] | None = None

    def now(self) -> float:
        return self.t

    async def sleep(self, delay: float) -> None:
        loop = asyncio.get_running_loop()
        fut: asyncio.Future[None] = loop.create_future()
        heapq.heappush(self._heap, (self.t + max(0.0, delay), self._seq, fut))
        self._seq += 1
        if self._driver is None or self._driver.done():
            self._driver = loop.create_task(self._drive())
        await fut  # a cancelled task raises CancelledError right here, like asyncio.sleep

    async def _drive(self) -> None:
        while self._heap:
            for _ in range(DRAIN):
                await asyncio.sleep(0)
            when, _, fut = heapq.heappop(self._heap)
            self.t = max(self.t, when)
            if not fut.done():
                fut.set_result(None)


# Seconds of REAL time a single test may take. The clock is fake, so a correct solution needs
# milliseconds; a test still running after this long is stuck — typically a busy-wait
# (`while not enough: await asyncio.sleep(0)`) that nothing ever wakes, because only the
# injected `sleep` advances the fake clock. Without this guard such a solution hangs pytest.
REAL_TIMEOUT = 5.0


def run(coro: Coroutine[Any, Any, Any]) -> Any:
    """asyncio.run with a fresh loop per test (the fake clock lives outside the loop)."""

    async def guarded() -> Any:
        try:
            return await asyncio.wait_for(coro, timeout=REAL_TIMEOUT)
        except TimeoutError:  # what wait_for raises on 3.11+ (asyncio.TimeoutError is its alias)
            raise AssertionError(
                f"still running after {REAL_TIMEOUT:.0f} s of REAL time — nothing advanced the "
                "fake clock. Wait with the injected `sleep`, not a busy loop or asyncio.sleep."
            ) from None

    return asyncio.run(guarded())


@pytest.fixture
def clock() -> FakeClock:
    return FakeClock()