"""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
from types import SimpleNamespace

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_interrupted_decode_accounts_for_released_but_not_admitted_camera(pilot, monkeypatch):
    module = pilot("run_joint_pilot")
    event = SimpleNamespace(time_ns=1_000_000_000, channel="camera", sequence=0, value={})
    archive = SimpleNamespace(counters=lambda: {}, close=lambda: None)
    monkeypatch.setattr(module, "SensorArchive", lambda _path: archive)
    monkeypatch.setattr(module, "camera_events", lambda *args: iter([event]))
    monkeypatch.setattr(module, "merged_events", lambda *args: iter([event]))

    def interrupted(*args):
        raise EOFError("decoder stopped by supervisor")

    monkeypatch.setattr(module, "send", interrupted)
    queue = pilot("pilot_queue").Mailbox()
    report = {}
    stop = SimpleNamespace(is_set=lambda: False, wait=lambda _seconds: False)
    args = SimpleNamespace(sensor_archive="unused", camera_index="unused", frames=1)
    module.produce(args, SimpleNamespace(stdin=None), queue, stop, report)
    assert report["arrivals"]["camera"] == 1
    assert report["failed_camera_sequences"] == [0]
    assert "decoder stopped" in queue.error and queue.dropped_count == 0
    assert queue.quiescent


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_layer_refresh_cadence_is_source_clocked_and_fail_closed(pilot):
    refresh = pilot("run_joint_pilot").layer_refresh_due

    assert refresh(None, 1_000_000_000, 50)
    assert not refresh(1_000_000_000, 1_049_999_999, 50)
    assert refresh(1_000_000_000, 1_050_000_000, 50)
    with pytest.raises(ValueError, match="backwards"):
        refresh(1_000_000_000, 999_999_999, 50)


def test_native_ddrnet_preprocess_matches_pillow_nearest_exactly(pilot):
    from PIL import Image

    bgr = np.random.default_rng(41).integers(0, 256, (600, 800, 3), dtype=np.uint8)
    rgb = Image.fromarray(bgr[:, :, ::-1]).crop((100, 0, 700, 600))
    resized = rgb.resize((512, 512), Image.Resampling.NEAREST)
    expected = (np.asarray(resized, dtype=np.float32) / 255.0).transpose(2, 0, 1)[None]
    actual = pilot("pilot_ddrnet_numpy").preprocess_bgr(bgr)
    np.testing.assert_array_equal(actual, expected)
    assert actual.flags.c_contiguous and actual.dtype == np.dtype("<f4")
    with pytest.raises(ValueError, match="contract"):
        pilot("pilot_ddrnet_numpy").preprocess_bgr(bgr[:512])


def test_triton_statistics_exclude_warmup_and_reject_counter_reset(pilot):
    names = ("success", "fail", "queue", "compute_input", "compute_infer", "compute_output")
    before = {"model": {name: {"count": "8", "ns": "1000000"} for name in names}}
    after = {"model": {name: {"count": "10", "ns": "5000000"} for name in names}}
    delta = pilot("pilot_triton_stats").stats_delta(before, after)
    assert delta["model"]["compute_infer"] == {"count": 2, "total_ms": 4, "mean_ms": 2}
    with pytest.raises(RuntimeError, match="backwards"):
        pilot("pilot_triton_stats").stats_delta(after, before)


def test_ddrnet_triton_transport_is_bounded_and_validates_output_identity(pilot):
    module = pilot("pilot_ddrnet_triton")
    mask = np.zeros((512, 512), np.uint8)
    descriptor = {
        "outputs": [
            {
                "name": "mask",
                "datatype": "UINT8",
                "shape": [1, 512, 512],
                "parameters": {"binary_data_size": mask.nbytes},
            }
        ]
    }
    header = json.dumps(descriptor).encode()

    class Response:
        status = 200
        payload = header + mask.tobytes()

        def read(self, limit):
            assert limit == module.MAX_RESPONSE_BYTES + 1
            return self.payload[:limit]

        def getheader(self, _name):
            return str(len(header))

    response = Response()
    backend = module.TritonDdrnetMaskHttpInferenceBackend()
    backend.connection = SimpleNamespace(
        request=lambda *args, **kwargs: None,
        getresponse=lambda: response,
        close=lambda: None,
    )
    tensor = np.zeros((1, 3, 512, 512), np.float32)
    np.testing.assert_array_equal(backend.infer(tensor), mask)
    response.payload = b"x" * (module.MAX_RESPONSE_BYTES + 1)
    with pytest.raises(RuntimeError, match="budget"):
        backend.infer(tensor)
    response.payload = header.replace(b'"UINT8"', b'"FP32" ') + mask.tobytes()
    with pytest.raises(RuntimeError, match="contract"):
        backend.infer(tensor)


def test_nvml_sampler_reuses_one_session_without_spawning_processes(pilot):
    class Function:
        def __init__(self, callback):
            self.callback = callback
            self.calls = 0

        def __call__(self, *args):
            self.calls += 1
            return self.callback(*args)

    def set_handle(_index, pointer):
        pointer._obj.value = 7
        return 0

    def set_memory(_handle, pointer):
        pointer._obj.total = 24 * 1024 * 1024
        pointer._obj.free = 20 * 1024 * 1024
        pointer._obj.used = 3 * 1024 * 1024 + 1
        return 0

    def set_utilization(_handle, pointer):
        pointer._obj.gpu = 42
        pointer._obj.memory = 3
        return 0

    library = SimpleNamespace(
        nvmlInit_v2=Function(lambda: 0),
        nvmlShutdown=Function(lambda: 0),
        nvmlDeviceGetHandleByIndex_v2=Function(set_handle),
        nvmlDeviceGetMemoryInfo=Function(set_memory),
        nvmlDeviceGetUtilizationRates=Function(set_utilization),
    )
    sampler = pilot("pilot_telemetry").NvmlSampler(library)

    assert sampler.sample() == {"gpu_used_mib": 4, "gpu_utilization": 42}
    assert sampler.sample() == {"gpu_used_mib": 4, "gpu_utilization": 42}
    sampler.close()
    sampler.close()

    assert library.nvmlInit_v2.calls == 1
    assert library.nvmlShutdown.calls == 1
    assert library.nvmlDeviceGetMemoryInfo.calls == 2
    assert library.nvmlDeviceGetUtilizationRates.calls == 2


def test_material_vote_tie_never_allows_by_class_order(pilot):
    choose = pilot("pilot_graph").select_material
    # unknown=0, hard_surface=1, bare_soil=2, grass=3
    votes = np.array(
        [[0, 0, 0, 2], [0, 1, 0, 1], [0, 2, 0, 1], [0, 0, 0, 0], [1, 1, 0, 0], [0, 0, 2, 2]]
    )
    np.testing.assert_array_equal(choose(votes), [3, 0, 1, 0, 0, 0])
    with pytest.raises(ValueError, match="invalid material"):
        choose(np.array([[0, -1, 0]]))


def test_numeric_diagnostic_counts_transitions_and_score_ties(pilot):
    module = pilot("diagnose_ddrnet_numeric")
    reference, candidate = np.array([[3, 3], [1, 0]]), np.array([[1, 3], [0, 0]])
    assert module.compare_masks(reference, reference)["different_pixels"] == 0
    assert module.compare_masks(reference, candidate) == {
        "different_pixels": 2,
        "transitions": [
            {"reference": 1, "candidate": 0, "pixels": 1},
            {"reference": 3, "candidate": 1, "pixels": 1},
        ],
    }
    logits = np.array([[[[100.0]], [[101.0]]]], dtype=np.float32)
    details = module.score_details(logits, np.ones_like(logits), np.array([[0]]), [(0, 0)])
    assert details[0]["selected"] == 0 and details[0]["logit_margin"] == 1
    assert details[0]["max_score_ties"] == [0, 1]


def test_optional_nvml_device_state_distinguishes_unavailable_from_zero(pilot):
    class Function:
        def __init__(self, code, value):
            self.code, self.value = code, value

        def __call__(self, *args):
            args[-1]._obj.value = self.value
            return self.code

    module = pilot("pilot_telemetry")
    sampler = object.__new__(module.NvmlSampler)
    sampler.closed, sampler.handle = False, None
    sampler.library = SimpleNamespace(
        nvmlDeviceGetClockInfo=Function(0, 2685),
        nvmlDeviceGetPerformanceState=Function(0, 0),
        nvmlDeviceGetPowerUsage=Function(3, 0),
    )
    state = sampler.sample_device_state()
    assert state["sm_clock_mhz"] == 2685 and state["pstate"] == 0
    assert state["power_mw"] is None and state["device_state_errors"] == {"power_mw": 3}
    del sampler.library.nvmlDeviceGetClockInfo
    assert sampler.sample_device_state()["sm_clock_mhz"] is None
    sampler.closed = True
    with pytest.raises(RuntimeError, match="closed"):
        sampler.sample_device_state()


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_waiting_consumer_uses_direct_handoff_without_evicting_ingress(pilot):
    mailbox = pilot("pilot_queue").Mailbox(capacity=2, byte_limit=100)
    compute_started = threading.Event()
    release_compute = threading.Event()

    def compute(bundle):
        if bundle["sequence"] == 0:
            compute_started.set()
            assert release_compute.wait(1)
        return bundle["sequence"] * 10

    stage = pilot("pilot_scheduler").GpuStage(mailbox, compute, threading.Event())
    bundles = [{"sequence": sequence, "payload_bytes": 10} for sequence in range(3)]
    mailbox.put(bundles[0])
    assert compute_started.wait(1)
    mailbox.put(bundles[1])
    mailbox.put(bundles[2])
    mailbox.finish()
    received = []
    consumer = threading.Thread(target=lambda: received.append(stage.take()))
    consumer.start()
    with stage.output_condition:
        assert stage.output_condition.wait_for(lambda: stage.consumer_waiting, timeout=1)
    release_compute.set()
    consumer.join(timeout=1)
    assert not consumer.is_alive()
    assert received[0][0]["sequence"] == 0
    assert not mailbox.dropped
    assert stage.direct_handoffs == 1
    for expected in (1, 2):
        bundle, result = stage.take()
        assert (bundle["sequence"], result) == (expected, expected * 10)
    for bundle in bundles:
        mailbox.release(bundle)
    assert stage.take() is None
    assert mailbox.bytes == 0 and 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()


def _sensor(pilot, channel, milliseconds, sequence=0):
    return pilot("pilot_source").SensorEvent(
        milliseconds * 1_000_000, channel, sequence, (np.zeros((1, 3)), None)
    )


def test_preroll_history_cannot_poison_the_first_fresh_pair(pilot):
    pose = _sensor(pilot, "pose", 975)
    points = tuple(_sensor(pilot, "points", t, i) for i, t in enumerate((590, 738, 808, 910, 980)))
    binding = pilot("pilot_sensor_binding").bind_sensors(1_000_000_000, pose, points)
    assert binding.available
    assert binding.increments == points[-2:]
    assert binding.history_only == points[:3]
    assert binding.pose_age_ns == 25_000_000
    assert binding.oldest_point_age_ns == 90_000_000
    assert binding.binding_age_ns == 65_000_000
    assert binding.document()["preroll_history_only"][0] == {
        "sequence": 0,
        "host_monotonic_ns": 590_000_000,
        "points": 1,
    }


def test_after_preroll_admission_does_not_silently_filter_bad_increments(pilot):
    pose = _sensor(pilot, "pose", 975)
    points = (_sensor(pilot, "points", 730), _sensor(pilot, "points", 980, 1))
    binding = pilot("pilot_sensor_binding").bind_sensors(
        1_000_000_000, pose, points, previous_camera_time_ns=700_000_000
    )
    assert binding.increments == points and not binding.history_only
    assert not binding.available
    assert set(binding.reasons) == {"oldest-points-too-old", "point-pose-skew"}


def test_pose_reuse_keeps_age_and_empty_increment_is_not_new_geometry(pilot):
    function = pilot("pilot_sensor_binding").bind_sensors
    pose = _sensor(pilot, "pose", 975)
    binding = function(1_000_000_000, pose, (), previous_camera_time_ns=980_000_000)
    assert binding.pose_state == "held" and binding.pose_age_ns == 25_000_000
    assert binding.points_state == "unavailable" and not binding.available
    assert binding.reasons == ("point-increment-unavailable",)
    expired = function(1_080_000_000, pose, (), previous_camera_time_ns=1_000_000_000)
    assert expired.pose_state == "stale" and "pose-too-old" in expired.reasons
    missing = function(1_000_000_000, None, ())
    assert missing.pose_state == "unavailable" and not missing.available


def test_sensor_binding_preserves_exact_limits_and_rejects_lookahead(pilot):
    function = pilot("pilot_sensor_binding").bind_sensors
    pose = _sensor(pilot, "pose", 900)
    points = (_sensor(pilot, "points", 800), _sensor(pilot, "points", 900, 1))
    assert function(1_000_000_000, pose, points, previous_camera_time_ns=700_000_000).available
    assert not function(1_000_000_001, pose, points, previous_camera_time_ns=700_000_000).available
    with pytest.raises(ValueError, match="future pose"):
        function(1_000_000_000, _sensor(pilot, "pose", 1001), ())
    with pytest.raises(ValueError, match="future points"):
        function(1_000_000_000, pose, (_sensor(pilot, "points", 1001),))
    with pytest.raises(ValueError, match="backwards"):
        function(1_000_000_000, pose, points[::-1])


def test_producer_keeps_preroll_history_for_tgs_and_publishes_sensor_ages(pilot, monkeypatch):
    runner = pilot("run_joint_pilot")
    events = [
        _sensor(pilot, "points", 590),
        _sensor(pilot, "pose", 975),
        _sensor(pilot, "points", 980, 1),
        pilot("pilot_source").SensorEvent(
            1_000_000_000, "camera", 0, {"sequence": 1, "host_epoch_ns": 2_000_000_000}
        ),
    ]

    class Archive:
        def __init__(self, path):
            pass

        def counters(self):
            return {}

        def close(self):
            pass

    class NoWait:
        def is_set(self):
            return False

        def wait(self, seconds):
            return False

    monkeypatch.setattr(runner, "SensorArchive", Archive)
    monkeypatch.setattr(runner, "camera_events", lambda *args: iter(events[-1:]))
    monkeypatch.setattr(runner, "merged_events", lambda *args: iter(events))
    raw = io.BytesIO()
    pilot("pilot_ipc").send(raw, {"decode_ms": 0.0}, bytes(600 * 800 * 3))
    raw.seek(0)
    decoder = SimpleNamespace(stdin=io.BytesIO(), stdout=raw)
    mailbox = pilot("pilot_queue").Mailbox()
    report = {}
    runner.produce(
        SimpleNamespace(sensor_archive="unused", camera_index="unused", frames=1),
        decoder,
        mailbox,
        NoWait(),
        report,
    )
    assert mailbox.error is None
    bundle = mailbox.take()
    assert bundle["available"] and len(bundle["points"]) == 1
    assert len(bundle["rolling_points"]) == 2
    np.testing.assert_array_equal(bundle["rolling_times"], [590_000_000, 980_000_000])
    assert bundle["lineage"]["point_increments"][0]["host_monotonic_ns"] == 980_000_000
    assert bundle["sensor_binding"]["preroll_history_only"][0]["host_monotonic_ns"] == 590_000_000
    assert report["preroll_history_only_points"] == 1
    mailbox.release(bundle)
    assert mailbox.bytes == 0 and mailbox.take() is None