"""Fixes to aiortc's SCTP stack, applied on import (aiortc 1.15).

The first two are the node's (`meshbay_node/transport/webrtc/sctp_fix.py`, where
why each is needed is written down); this is a copy because the client does not
depend on the node, and `meshbay-common` does not depend on aiortc (the hub uses
it without). `test_sctp_fix.py` in each package holds its copy to aiortc itself.

In a download this side is the receiver of the node's 1 MiB `file_chunk`
messages, about 870 SCTP packets each, and two receiver costs set the ceiling.
Measured on loopback, 400 MB, one connection:

- aiortc walked every fragment of the message in progress after each one
  arrived, 153 million steps: two thirds of the client's CPU. Fixed by the
  reassembly fix below — 20 MB/s became 28;
- it acknowledged every packet. RFC 4960 §6.2 asks for one SACK every second
  packet (Chromium does that), and each SACK costs the node a DTLS decryption
  and a parse on the one core that also sends: more of its time than sending
  the data. One every second packet — 28 MB/s became 39, the node's own figure
  for the transfer 49.

A SACK still goes at once when something is out of order or duplicated (that is
the loss signal the sender is waiting for) or the packet carried anything other
than data; a lone packet is acknowledged within `SACK_DELAY_S`.
"""

import asyncio
import contextlib

from aiortc.rtcsctptransport import (
    SCTP_DATA_LAST_FRAG,
    DataChunk,
    InboundStream,
    InitChunk,
    RTCSctpTransport,
    parse_packet,
)

# RFC 4960 §6.2: at most 500 ms; 200 is what implementations use.
SACK_DELAY_S = 0.2


def _flight_size_increase(self, chunk) -> None:
    chunk._acked = False
    self._flight_size += chunk._book_size


RTCSctpTransport._flight_size_increase = _flight_size_increase


_add_chunk = InboundStream.add_chunk
_pop_messages = InboundStream.pop_messages
_prune_chunks = InboundStream.prune_chunks


def add_chunk(self, chunk) -> None:
    _add_chunk(self, chunk)
    if chunk.flags & SCTP_DATA_LAST_FRAG:
        self._last_frags = getattr(self, "_last_frags", 0) + 1


def pop_messages(self):
    # Only ever too high (a pruned last fragment), never too low: a message
    # leaves the reassembly with exactly one last fragment, and only here.
    if not getattr(self, "_last_frags", 0):
        return
    for message in _pop_messages(self):
        self._last_frags -= 1
        yield message


def prune_chunks(self, tsn: int) -> int:
    size = _prune_chunks(self, tsn)
    self._last_frags = sum(1 for c in self.reassembly if c.flags & SCTP_DATA_LAST_FRAG)
    return size


InboundStream.add_chunk = add_chunk
InboundStream.pop_messages = pop_messages
InboundStream.prune_chunks = prune_chunks


async def _handle_data(self, data: bytes) -> None:
    """aiortc's own, down to the SACK at the end: that one may wait for a
    second packet."""
    try:
        _, _, verification_tag, chunks = parse_packet(data)
    except ValueError:
        return
    if any(isinstance(x, InitChunk) for x in chunks):
        assert len(chunks) == 1
        expected_tag = 0
    else:
        expected_tag = self._local_verification_tag
    if verification_tag != expected_tag:
        return
    for chunk in chunks:
        await self._receive_chunk(chunk)
    if not self._sack_needed:
        return
    self._sack_packets = getattr(self, "_sack_packets", 0) + 1
    if (self._sack_packets >= 2 or self._sack_misordered or self._sack_duplicates
            or not all(isinstance(c, DataChunk) for c in chunks)):
        await _sack_now(self)
    elif getattr(self, "_sack_timer", None) is None:
        self._sack_timer = self._loop.call_later(SACK_DELAY_S, _sack_later, self)


async def _sack_now(self) -> None:
    timer = getattr(self, "_sack_timer", None)
    if timer is not None:
        timer.cancel()
        self._sack_timer = None
    self._sack_packets = 0
    if self._sack_needed and self._association_state != self.State.CLOSED:
        # The DTLS transport may be gone under a timer that fired late.
        with contextlib.suppress(ConnectionError):
            await self._send_sack()


def _sack_later(self) -> None:
    self._sack_timer = None
    # Held: a task nobody references can be collected before it runs.
    self._sack_task = asyncio.ensure_future(_sack_now(self))


RTCSctpTransport._handle_data = _handle_data