"""
On Windows the node's datagrams go straight through the socket, not one
overlapped operation per datagram and per turn of the loop (`proactor_udp.py`).

Run on a real ProactorEventLoop, which only Windows has.
"""

import asyncio
import socket
import sys
import time

import pytest
from meshbay_node.transport.webrtc import proactor_udp

pytestmark = pytest.mark.skipif(sys.platform != "win32", reason="the proactor is Windows'")

N = 2000


class _Collect(asyncio.DatagramProtocol):
    def __init__(self, raise_on: int = -1):
        self.got: list[int] = []
        self.raise_on = raise_on
        self.all = asyncio.get_running_loop().create_future()

    def datagram_received(self, data, addr):
        self.got.append(int.from_bytes(data[:4], "big"))
        if len(self.got) == N and not self.all.done():
            self.all.set_result(None)
        if len(self.got) == self.raise_on:
            raise RuntimeError("a protocol's own bug")


class _Refusing(socket.socket):
    """A socket that refuses a send every so often, as a full buffer does, and
    can answer one read with the reset an ICMP port unreachable leaves. The
    overlapped path works by handle and never calls these."""

    def __init__(self, every: int = 0, reset_once: bool = False):
        super().__init__(socket.AF_INET, socket.SOCK_DGRAM)
        self.every, self.calls, self.reset_once = every, 0, reset_once

    def sendto(self, data, addr):
        self.calls += 1
        if self.every and self.calls % self.every == 0:
            raise BlockingIOError
        return super().sendto(data, addr)

    def recvfrom(self, n):
        if self.reset_once:
            self.reset_once = False
            raise ConnectionResetError(10054, "port unreachable")
        return super().recvfrom(n)


def _busy_turns(ms: float):
    """A loop with other work to do: every turn takes `ms` longer."""
    async def busy():
        while True:
            end = time.perf_counter() + ms / 1000
            while time.perf_counter() < end:
                pass
            await asyncio.sleep(0)
    return asyncio.ensure_future(busy())


async def _exchange(*, refuse_every=0, reset_once=False, raise_on=-1, busy_ms=0.0):
    """N datagrams from one proactor transport to another on loopback: what
    arrived in what order, and how many overlapped sends and reads it took."""
    loop = asyncio.get_running_loop()
    rsock = _Refusing(reset_once=reset_once)
    rsock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 8 << 20)
    rsock.bind(("127.0.0.1", 0))
    rtr, rx = await loop.create_datagram_endpoint(lambda: _Collect(raise_on), sock=rsock)
    tsock = _Refusing(refuse_every)
    tsock.bind(("127.0.0.1", 0))
    tx, _ = await loop.create_datagram_endpoint(asyncio.DatagramProtocol, sock=tsock)
    counts = {"sendto": 0, "recvfrom": 0}
    originals = {name: getattr(loop._proactor, name) for name in counts}

    def counting(name):
        def call(*a, **kw):
            counts[name] += 1
            return originals[name](*a, **kw)
        return call
    for name in counts:
        setattr(loop._proactor, name, counting(name))
    errors = []
    loop.set_exception_handler(lambda _loop, ctx: errors.append(ctx))
    busy = _busy_turns(busy_ms) if busy_ms else None
    try:
        dest = rsock.getsockname()
        for i in range(N):
            tx.sendto(i.to_bytes(4, "big") * 50, dest)
            if i % 50 == 49:
                await asyncio.sleep(0)
        await asyncio.wait_for(rx.all, 20)
        return rx.got, counts, errors
    finally:
        if busy:
            busy.cancel()
        for name, fn in originals.items():
            setattr(loop._proactor, name, fn)
        tx.close()
        rtr.close()


def _run(**kw):
    return asyncio.run(_exchange(**kw), loop_factory=asyncio.ProactorEventLoop)


def test_the_transport_is_patched():
    cls = asyncio.proactor_events._ProactorDatagramTransport
    assert cls.sendto is proactor_udp.sendto
    assert cls._loop_reading is proactor_udp._loop_reading


def test_datagrams_go_straight_out_in_order():
    got, counts, _ = _run()
    assert got == list(range(N))
    assert counts["sendto"] == 0


def test_what_the_socket_refuses_is_queued_and_order_is_kept():
    """A refusal sends that datagram the overlapped way, and the ones after it
    wait behind it rather than overtake it."""
    got, counts, _ = _run(refuse_every=97)
    assert got == list(range(N))
    assert counts["sendto"] >= N // 97


def test_a_busy_loop_reads_everything_waiting_in_order():
    """Unpatched, this loop read one datagram a turn: about 2,000 a second
    whatever was waiting, and the kernel dropped the rest."""
    got, counts, _ = _run(busy_ms=2)
    assert got == list(range(N))
    assert counts["recvfrom"] < N // 10


def test_a_port_unreachable_on_a_direct_read_is_skipped():
    got, _, errors = _run(reset_once=True, busy_ms=2)
    assert got == list(range(N))
    assert not errors


def test_a_protocol_that_raises_does_not_end_the_reading():
    got, _, errors = _run(raise_on=10)
    assert got == list(range(N))
    assert [str(e["exception"]) for e in errors] == ["a protocol's own bug"]