"""Small synthetic checks only; model/real-source runs belong on Worker 006."""

import importlib
import io
import json
import threading
import zipfile
from pathlib import Path

import numpy as np
import pytest


@pytest.fixture
def pilot(monkeypatch):
    root = Path(__file__).resolve().parents[1]
    monkeypatch.syspath_prepend(
        str(root / "experiments/perception/worker/streaming_profile_stage1")
    )
    return lambda name: importlib.import_module(name)


def test_pending_overflow_is_explicit_and_does_not_evict_active(pilot):
    queue = pilot("pilot_queue").Mailbox(capacity=2, byte_limit=100)

    def make(sequence):
        return {"sequence": sequence, "payload_bytes": 10}

    queue.put(make(0))
    active = queue.take()
    for sequence in (1, 2, 3):
        queue.put(make(sequence))
    assert queue.dropped == [{"sequence": 1, "reason": "pending-overflow"}]
    assert queue.bytes == 30
    assert queue.peak_pending == 2
    queue.release(active)
    assert queue.take()["sequence"] == 2
    assert queue.take()["sequence"] == 3


def test_active_bytes_count_towards_memory_limit(pilot):
    queue = pilot("pilot_queue").Mailbox(byte_limit=20)
    queue.put({"sequence": 0, "payload_bytes": 20})
    active = queue.take()
    queue.put({"sequence": 1, "payload_bytes": 1})
    assert queue.dropped == [{"sequence": 1, "reason": "byte-budget"}]
    assert queue.peak_bytes == 20
    queue.release(active)
    queue.finish()
    assert queue.take() is None


def test_numpy_member_is_read_incrementally_and_bounded(pilot):
    data = io.BytesIO()
    np.savez_compressed(data, points=np.arange(300, dtype="<f8").reshape(100, 3))
    data.seek(0)
    with zipfile.ZipFile(data) as archive:
        rows = pilot("pilot_source").NpyRows(archive, "points", columns=(3,), dtype="<f8")
        np.testing.assert_array_equal(rows.take(2), np.arange(6).reshape(2, 3))
        assert rows.count == 100 and rows.position == 2 and rows.bytes_read == 48
        with pytest.raises(ValueError, match="budget"):
            rows.take(99)


def test_original_camera_clock_not_fixed_frame_rate(pilot, tmp_path):
    path = tmp_path / "camera.jsonl"
    stamps = [123000, 127777, 987654]
    rows = [
        {
            "schema_version": "missioncore.camera-recording-index/v1",
            "kind": "media",
            "sequence": i + 1,
            "host_monotonic_ns": stamp,
        }
        for i, stamp in enumerate(stamps)
    ]
    path.write_text("\n".join(json.dumps(row) for row in rows))
    events = list(pilot("pilot_source").camera_events(path, 2))
    assert [e.time_ns for e in events] == stamps[:2]
    assert [e.sequence for e in events] == [0, 1]


def test_ipc_preserves_binary_and_rejects_oversized_header(pilot):
    ipc = pilot("pilot_ipc")
    stream = io.BytesIO()
    ipc.send(stream, {"op": "infer"}, b"\x00\xff")
    stream.seek(0)
    assert ipc.receive(stream) == ({"op": "infer"}, b"\x00\xff")
    with pytest.raises(ValueError, match="header"):
        ipc.receive(io.BytesIO((65537).to_bytes(4, "little")))


def test_nearest_rank_tail_metrics_keep_outlier(pilot):
    metrics = pilot("run_joint_pilot").distribution([1] * 99 + [200])
    assert metrics["p99"] == 1
    assert metrics["max"] == 200
    assert pilot("run_joint_pilot").distribution([1, 200])["p99"] == 200


def test_body_history_is_bounded_and_never_reads_future(pilot):
    store = pilot("pilot_graph").CurrentStore(None)
    for index in range(100):
        store.frame_id = f"frame-{index:06d}"
        store.body = index
        store.remember_body()
    assert len(store.body_history) == 64
    assert store.body_frame_for_frame("frame-000098") == 98
    assert store.body_frame_for_frame("frame-000000") is None
    with pytest.raises(ValueError, match="future"):
        store.body_frame_for_frame("frame-000100")


def test_vectorized_grid_lookup_preserves_exact_cells_and_holes(pilot):
    grid = np.asarray([[-2, -1], [0, 0], [2, 1]])
    points = np.asarray([[x, y] for x in range(-5, 6) for y in range(-5, 6)])
    lookup = {tuple(row): i for i, row in enumerate(grid)}
    expected = np.asarray([lookup.get(tuple(row), -1) for row in points])
    np.testing.assert_array_equal(pilot("pilot_graph").grid_indices(points, grid), expected)


def test_ddrnet_layout_preserves_class_order_and_ties(pilot):
    class Scores:
        def __init__(self, array):
            self.array = array

        def permute(self, *axes):
            return Scores(self.array.transpose(axes))

        def contiguous(self):
            return Scores(np.ascontiguousarray(self.array))

    # Include saturated/equal probabilities: first class must win every tie.
    scores = Scores(np.array([[[[1.0, 0.5]], [[1.0, 0.5]], [[0.0, 0.9]], [[1.0, 0.9]]]]))
    function = pilot("pilot_ddrnet_runtime").layout_scores
    reference, reference_dim = function(scores, "reference")
    candidate, candidate_dim = function(scores, "channels-last")
    np.testing.assert_array_equal(
        reference.array.argmax(axis=reference_dim), candidate.array.argmax(axis=candidate_dim)
    )
    np.testing.assert_array_equal(candidate.array.argmax(axis=candidate_dim), [[[0, 2]]])
    assert candidate.array.flags.c_contiguous
    with pytest.raises(ValueError, match="unknown"):
        function(scores, "silent-change")


def test_gpu_stage_preserves_order_and_reserves_bounded_output_before_compute(pilot):
    mailbox = pilot("pilot_queue").Mailbox(capacity=2, byte_limit=30)
    started = [threading.Event(), threading.Event()]
    stop = threading.Event()

    def compute(bundle):
        started[bundle["sequence"]].set()
        return bundle["sequence"] * 10

    stage = pilot("pilot_scheduler").GpuStage(mailbox, compute, stop)
    try:
        mailbox.put({"sequence": 0, "payload_bytes": 10})
        assert started[0].wait(1)
        mailbox.put({"sequence": 1, "payload_bytes": 10})
        mailbox.finish()
        assert not started[1].wait(0.05)  # No hidden completed third slot.
        assert mailbox.bytes == 20  # GPU completion cannot release CPU-owned input.
        first, result = stage.take()
        assert (first["sequence"], result) == (0, 0)
        assert started[1].wait(1)
        second, result = stage.take()
        assert (second["sequence"], result) == (1, 10)
        assert mailbox.bytes == 20
        mailbox.release(first)
        mailbox.release(second)
        assert stage.take() is None
        assert mailbox.bytes == 0 and not mailbox.dropped
        assert stage.peak_pending <= 1
    finally:
        assert stage.close()


def test_gpu_stage_propagates_failure_and_stops_waiting_for_input(pilot):
    mailbox = pilot("pilot_queue").Mailbox(capacity=1)

    def compute(bundle):
        raise ValueError("synthetic GPU failure")

    stage = pilot("pilot_scheduler").GpuStage(mailbox, compute, threading.Event())
    mailbox.put({"sequence": 0, "payload_bytes": 10})
    try:
        with pytest.raises(RuntimeError, match="synthetic GPU failure"):
            stage.take()
    finally:
        assert stage.close()
    waiting = pilot("pilot_scheduler").GpuStage(
        pilot("pilot_queue").Mailbox(capacity=1), compute, threading.Event()
    )
    assert waiting.close()


def test_completed_gpu_and_ingress_share_two_pending_slots(pilot):
    mailbox = pilot("pilot_queue").Mailbox(capacity=2, byte_limit=100)
    mailbox.put({"sequence": 0, "payload_bytes": 10})
    active = mailbox.take()
    for sequence in (1, 2):
        mailbox.put({"sequence": sequence, "payload_bytes": 10})
    mailbox.reserve_completed()
    assert mailbox.dropped == [{"sequence": 1, "reason": "handoff-overflow"}]
    assert mailbox.external_pending + len(mailbox.pending) == 2
    mailbox.put({"sequence": 3, "payload_bytes": 10})
    assert mailbox.dropped[-1] == {"sequence": 2, "reason": "pending-overflow"}
    assert mailbox.peak_pending == 2 and mailbox.bytes == 20
    mailbox.take_completed()
    mailbox.release(active)
    last = mailbox.take()
    assert last["sequence"] == 3
    mailbox.release(last)
    assert mailbox.bytes == 0
    with pytest.raises(ValueError, match="accounting"):
        mailbox.take_completed()