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

**The sender leaked bytes in flight.** aiortc counts the bytes in flight and
sends nothing new while that count is at or above the congestion window. A chunk
the peer acknowledged in a gap block is marked `_acked` and taken out of the
count. If the T3 timer then expires, every chunk still queued is sent again,
that one included, and goes back into the count with `_acked` still set. When
the cumulative ack finally covers it, the SACK handler skips it because it reads
as acknowledged already, and its bytes stay counted for the life of the
association.

Each such loss leaks a chunk. On a lossy path (a phone on Wi-Fi) the leak grows
past the window, which T3 itself keeps small. From then on the node sends
nothing on that connection: nothing is outstanding, so no timer runs to unstick
it, and the association still reads as established. The peer keeps talking and
the node keeps answering into a queue that never moves. Live symptom: the photo
backup stalled after a few dozen photos, its slot granted and the grant never
delivered (docs/MESHBAY_DESIGN.md §9.12). The fix is the invariant the code
assumed: a chunk being sent again is in flight, so it is not acknowledged.

**The receiver reassembled in quadratic time.** After every DATA chunk,
`InboundStream.pop_messages` walks the reassembly list from its start, through
every fragment of the message still being assembled. A message of n fragments
costs n²/2 steps, and a 1 MiB `file_chunk` is about 870 fragments: measured on a
receiving aiortc, 153 million steps for 400 MB, two thirds of its CPU, and the
whole of why it could not take more than 20 MB/s on loopback. No message can
come out of a reassembly holding no last fragment, so the walk is skipped until
one is there; when one is, the walk is aiortc's own, unchanged.

The node receives uploads in 48 KB messages, where the walk is short; the
command-line client receives 1 MiB ones and carries the same fix
(`meshbay_cli/client/sctp_fix.py`).

**The receiver acknowledged every packet.** RFC 4960 §6.2 asks for one SACK
every second packet, and Chromium does that; aiortc sent one at the end of
every packet of data. Each SACK costs its receiver a DTLS decryption and a parse
on the one core that also does everything else — on a node sending a download
to the command-line client, more time than sending the data. In an upload the
node is the receiver and pays the other half: the encryption and the send of a
SACK per packet. A SACK still goes at once when something is out of order or
duplicated (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`.

`test_sctp_fix.py` replays all three against aiortc itself.
"""

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