"""The client's copy of the node's aiortc fixes (client/sctp_fix.py).

The copy exists because neither package may depend on the other; it is held to
the node's by comparing the code itself, docstrings aside, and the two effects
that set a download's speed here are replayed against aiortc.
"""

import ast
import asyncio
from pathlib import Path

import pytest
from aiortc.rtcsctptransport import (
    SCTP_DATA_FIRST_FRAG,
    SCTP_DATA_LAST_FRAG,
    DataChunk,
    InboundStream,
    RTCSctpTransport,
    serialize_packet,
)
from meshbay_cli.client import sctp_fix

NODE_COPY = (Path(__file__).resolve().parents[2] / "meshbay-node" / "src" / "meshbay_node"
             / "transport" / "webrtc" / "sctp_fix.py")


def _code(path: Path) -> str:
    tree = ast.parse(path.read_text(encoding="utf-8"))
    tree.body = tree.body[1:]          # the module docstring
    return ast.dump(tree)


@pytest.mark.skipif(not NODE_COPY.exists(), reason="not in the source tree")
def test_the_code_is_the_nodes():
    assert _code(Path(sctp_fix.__file__)) == _code(NODE_COPY)


def _chunks(n_fragments: int, count: int, tsn: int = 1000) -> list[DataChunk]:
    out = []
    for seq in range(count):
        for i in range(n_fragments):
            c = DataChunk()
            c.flags = ((SCTP_DATA_FIRST_FRAG if i == 0 else 0)
                       | (SCTP_DATA_LAST_FRAG if i == n_fragments - 1 else 0))
            c.tsn, c.stream_id, c.stream_seq, c.protocol = tsn, 1, seq, 53
            c.user_data = b"x"
            tsn += 1
            out.append(c)
    return out


def test_a_megabyte_is_walked_once(monkeypatch):
    walks = []
    original = sctp_fix._pop_messages
    monkeypatch.setattr(sctp_fix, "_pop_messages",
                        lambda s: (walks.append(1), original(s))[1])
    s = InboundStream()
    got = 0
    for c in _chunks(870, 3):
        s.add_chunk(c)
        got += len(list(s.pop_messages()))
    assert got == 3 and len(walks) == 3


def test_every_second_packet_is_acknowledged():
    class _Dtls:
        state = "connected"

        def _register_data_receiver(self, *a):
            pass

    async def go():
        s = RTCSctpTransport(_Dtls())
        s._association_state = s.State.ESTABLISHED
        s._last_received_tsn = 999
        sacks = []

        async def send_sack():
            sacks.append(s._last_received_tsn)
            s._sack_needed = False

        async def deliver(*_a):
            pass

        s._send_sack, s._receive = send_sack, deliver
        for c in _chunks(1, 6):
            await s._handle_data(serialize_packet(1, 1, s._local_verification_tag, c))
        return sacks
    assert asyncio.run(go()) == [1001, 1003, 1005]