"""Coverage for bounded, ordered prefill transfer."""

from threading import Event

import pytest

from meshnet_gateway.prefill_backpressure import (
    BoundedPrefillSender,
    DEFAULT_PREFILL_CHUNK_TOKENS,
    DEFAULT_PREFILL_MAX_IN_FLIGHT,
    PrefillTransferLimits,
)
from meshnet_gateway.server import _BinaryActivation, _post_binary_forward


def test_limits_have_safe_defaults_and_keep_legacy_chunk_setting(monkeypatch):
    monkeypatch.delenv("MESHNET_PREFILL_CHUNK_TOKENS", raising=False)
    monkeypatch.delenv("MESHNET_PREFILL_MAX_IN_FLIGHT", raising=False)
    monkeypatch.setenv("MESHNET_CHUNK_TOKENS", "64")

    limits = PrefillTransferLimits.from_env()

    assert limits.chunk_tokens == 64
    assert limits.max_in_flight == DEFAULT_PREFILL_MAX_IN_FLIGHT
    assert PrefillTransferLimits().chunk_tokens == DEFAULT_PREFILL_CHUNK_TOKENS
    assert limits.max_buffered_bytes == limits.max_chunk_bytes


def test_slow_consumer_applies_backpressure_preserves_order_and_bounds_bytes():
    sender = BoundedPrefillSender(
        PrefillTransferLimits(chunk_tokens=2, max_in_flight=4, max_chunk_bytes=8)
    )
    sent: list[int] = []
    produced: list[int] = []

    def chunks():
        for index in range(3):
            produced.append(index)
            yield bytes([index]) * 8

    def forward(chunk: bytes) -> int:
        sent.append(chunk[0])
        # The next item cannot have been constructed while this consumer waits.
        assert produced[-1] == chunk[0]
        assert len(produced) == chunk[0] + 1
        assert sender.buffered_bytes == 8
        assert sender.in_flight == 1
        return chunk[0]

    assert sender.send(chunks(), body_size=len, forward=forward) == [0, 1, 2]
    assert sent == [0, 1, 2]
    assert sender.peak_buffered_bytes <= sender.limits.max_buffered_bytes
    assert sender.peak_in_flight == 1
    assert sender.buffered_bytes == sender.in_flight == 0


def test_failure_and_cancellation_release_owned_buffers():
    sender = BoundedPrefillSender(PrefillTransferLimits())
    with pytest.raises(RuntimeError, match="route lost"):
        sender.send([b"one", b"two"], body_size=len, forward=lambda _: (_ for _ in ()).throw(RuntimeError("route lost")))
    assert sender.closed
    assert sender.buffered_bytes == sender.in_flight == 0

    cancelled = Event()
    cancelled.set()
    sender = BoundedPrefillSender(PrefillTransferLimits())
    assert sender.send([b"never"], body_size=len, forward=lambda _: None, cancelled=cancelled) == []
    assert sender.buffered_bytes == sender.in_flight == 0


def test_legacy_single_chunk_peer_response_uses_outgoing_metadata(monkeypatch):
    class Response:
        headers = {"Content-Type": "application/octet-stream"}

        def read(self):
            return b"\0" * (1 * 1 * 64 * 2)

        def __enter__(self):
            return self

        def __exit__(self, *_):
            return False

    monkeypatch.setattr("urllib.request.urlopen", lambda *_args, **_kwargs: Response())
    activation = _BinaryActivation(
        body=b"\0" * (1 * 1 * 64 * 2), shape=[1, 1, 64], dtype="bfloat16",
        session="legacy-session", chunk_index=0, chunk_total=1, encoding=None, headers={},
    )

    response = _post_binary_forward("http://legacy/forward", activation, hop_index=0, timeout=1.0)

    assert response.session == activation.session
    assert response.chunk_index == response.chunk_total - 1 == 0