"""IPC/SHM tests: encode/decode round-trips, corruption handling, ring semantics.

Pins the agreed contract:
  * encode -> decode is lossless for raw/preprocessed traces and result payloads;
  * a corrupt/truncated frame raises a single, catchable ValueError (never a bare
    struct.error / UnicodeDecodeError);
  * the ring is latest-wins: on overflow the oldest slot is overwritten and a slow
    reader keeps the freshest frames;
  * opening a missing/incompatible ring fails fast.
"""

from __future__ import annotations

import struct
import threading
import time
import unittest
from contextlib import suppress
from pathlib import Path
from typing import cast
from unittest import mock

import numpy as np

from python_app.models.dataset_model import (
    ComboKey,
    SweepCollection,
    TraceData,
)
from python_app.orchestration.shm.decoder import (
    PREPROC_MAGIC,
    RAW_MAGIC,
    RESULT_MAGIC,
    decode_result_collection,
    decode_trace_collection,
)
from python_app.orchestration.shm.ring_reader import ShmRingReader
from python_app.orchestration.shm.ring_writer import ShmRingWriter
from python_app.storage.npz.serialize import serialize_trace_collection


def _trace(in_pos: int, out_pos: int, n: int, *, capture_ns: tuple[int, int] = (0, 0)) -> TraceData:
    """Build a trace with float32-exact data so round-trips compare exactly."""
    freq = np.arange(n, dtype=np.float32) + 1.0
    s21 = (-np.arange(n, dtype=np.float32) + 2.0j * np.arange(n, dtype=np.float32)).astype(np.complex64)
    return TraceData(
        combo=ComboKey(input=in_pos, output=out_pos),
        frequency_hz=freq,
        s21=s21,
        capture_start_ns=capture_ns[0],
        capture_end_ns=capture_ns[1],
    )


class TraceCollectionRoundTripTest(unittest.TestCase):
    def _assert_round_trips(self, magic: int) -> None:
        collection = SweepCollection(
            collection_id=7,
            monotonic_ns=123,
            traces=[_trace(0, 0, 4, capture_ns=(11, 13)), _trace(3, 1, 2, capture_ns=(15, 19))],
            capture_start_ns=10,
            capture_end_ns=20,
        )
        decoded = decode_trace_collection(serialize_trace_collection(collection, magic), magic)
        self.assertEqual(decoded.collection_id, 7)
        self.assertEqual(decoded.monotonic_ns, 123)
        self.assertEqual((decoded.capture_start_ns, decoded.capture_end_ns), (10, 20))
        self.assertEqual(len(decoded.traces), 2)
        for original, got in zip(collection.traces, decoded.traces):
            self.assertEqual((got.combo.input, got.combo.output), (original.combo.input, original.combo.output))
            self.assertTrue(np.array_equal(got.frequency_hz, original.frequency_hz))
            self.assertTrue(np.array_equal(got.s21, original.s21))
            self.assertEqual(
                (got.capture_start_ns, got.capture_end_ns),
                (original.capture_start_ns, original.capture_end_ns),
            )

    def test_raw_round_trips(self) -> None:
        self._assert_round_trips(RAW_MAGIC)

    def test_preprocessed_round_trips(self) -> None:
        self._assert_round_trips(PREPROC_MAGIC)

    def test_s21_only_wire_uses_literal_raw3_prp3_and_twelve_bytes_per_point(self) -> None:
        collection = SweepCollection(7, 123, [_trace(1, 2, 3)], capture_start_ns=100, capture_end_ns=120)
        for magic, label in ((RAW_MAGIC, b"RAW3"), (PREPROC_MAGIC, b"PRP3")):
            with self.subTest(magic=label):
                wire = serialize_trace_collection(collection, magic)
                expected = label + struct.pack("<QQI", 7, 123, 1)
                expected += struct.pack("<III3f6f", 1, 2, 3, 1., 2., 3., 0., 0., -1., 2., -2., 4.)
                expected += struct.pack("<QQIQQ", 100, 120, 1, 0, 0)
                self.assertEqual(wire, expected)
                self.assertEqual(len(wire), 44 + 28 + 12 * 3)

    def test_old_trace_magics_are_rejected(self) -> None:
        collection = SweepCollection(1, 2, [_trace(0, 0, 2)])
        for magic, old in ((RAW_MAGIC, b"RAW2"), (PREPROC_MAGIC, b"PPR2"), (PREPROC_MAGIC, b"PRP2")):
            with self.subTest(old=old), self.assertRaisesRegex(ValueError, "magic"):
                decode_trace_collection(old + serialize_trace_collection(collection, magic)[4:], magic)

    def test_complex_components_are_decoded_without_arithmetic_changes(self) -> None:
        words = [0x3f800000, 0x7fc01234, 0x80000000, 0xbf800000]  # (1, NaN), (-0, -1)
        samples = struct.pack("<4I", *words)
        wire = b"RAW3" + struct.pack("<QQIIII2f", 1, 2, 1, 0, 0, 2, 10., 20.) + samples + struct.pack("<QQIQQ", 0, 0, 1, 0, 0)
        decoded = decode_trace_collection(wire, RAW_MAGIC).traces[0]
        self.assertEqual(decoded.s21.dtype, np.dtype("<c8"))
        np.testing.assert_array_equal(decoded.s21.view("<u4"), words)
        result_wire = struct.pack("<I6QfII", RESULT_MAGIC, 1, 2, 0, 0, 1, 2, 0., 1, 0)
        result_wire += struct.pack("<BHI2f", 1, 0, 2, 10., 20.) + samples
        trace = decode_result_collection(result_wire).collection_payloads[0].trace
        np.testing.assert_array_equal(trace.view("<u4"), words)

    def test_point_count_cannot_read_beyond_the_available_arrays(self) -> None:
        for magic in (RAW_MAGIC, PREPROC_MAGIC):
            for point_count in (3, 0xffffffff):
                with self.subTest(magic=magic, points=point_count), self.assertRaisesRegex(ValueError, "truncated"):
                    wire = struct.pack("<IQQIIII6f", magic, 1, 2, 1, 0, 0, point_count, 1., 2., 3., 4., 5., 6.)
                    decode_trace_collection(wire, magic)

    def test_empty_traces_round_trip(self) -> None:
        collection = SweepCollection(collection_id=1, monotonic_ns=2, traces=[])
        decoded = decode_trace_collection(serialize_trace_collection(collection, RAW_MAGIC), RAW_MAGIC)
        self.assertEqual(decoded.traces, [])

    def test_live_raw_requires_capture_window_but_calibration_bundles_do_not(self) -> None:
        collection = SweepCollection(1, 100, [_trace(0, 0, 2)])
        for start, end in ((0, 0), (0, 1), (10, 9)):
            collection.capture_start_ns, collection.capture_end_ns = start, end
            with self.subTest(window=(start, end)), self.assertRaisesRegex(ValueError, "capture window"):
                serialize_trace_collection(collection, RAW_MAGIC, require_capture_window=True)
        collection.capture_start_ns = collection.capture_end_ns = 10
        serialize_trace_collection(collection, RAW_MAGIC, require_capture_window=True)
        collection.capture_start_ns = collection.capture_end_ns = 0
        decoded = decode_trace_collection(serialize_trace_collection(collection, RAW_MAGIC), RAW_MAGIC)
        self.assertEqual((decoded.capture_start_ns, decoded.capture_end_ns), (0, 0))

    def test_timing_trailer_is_required(self) -> None:
        collection = SweepCollection(4, 5, [_trace(1, 0, 3, capture_ns=(7, 9))],
                                     capture_start_ns=6, capture_end_ns=10)
        for magic in (RAW_MAGIC, PREPROC_MAGIC):
            payload = serialize_trace_collection(collection, magic)
            for removed_bytes in (20, 36):
                with self.subTest(magic=magic, removed_bytes=removed_bytes), self.assertRaisesRegex(ValueError, "truncated"):
                    decode_trace_collection(payload[:-removed_bytes], magic)

    def test_per_trace_window_count_mismatch_is_rejected(self) -> None:
        collection = SweepCollection(
            collection_id=4, monotonic_ns=5, traces=[_trace(1, 0, 3, capture_ns=(7, 9))]
        )
        payload = serialize_trace_collection(collection, RAW_MAGIC)
        # Overwrite the window-table count (u32 before the single 16-byte pair).
        corrupt = payload[:-20] + struct.pack("<I", 2) + payload[-16:]
        with self.assertRaises(ValueError):
            decode_trace_collection(corrupt, RAW_MAGIC)


class ResultCollectionDecodeTest(unittest.TestCase):
    def test_empty_tables_and_images_retain_their_declared_dimensions(self) -> None:
        for rows, columns in ((0, 3), (2, 0), (0, 0)):
            for kind in (3, 4):
                with self.subTest(shape=(rows, columns), kind=kind):
                    wire = struct.pack("<I6QfII", RESULT_MAGIC, 1, 2, 0, 0, 1, 2, 0., 1, 0)
                    wire += struct.pack("<BHII", kind, 0, columns, rows)
                    if kind == 3:
                        wire += bytes(4 * (rows + columns))  # zero-valued axes
                    if kind == 4 and rows and not columns:
                        with self.assertRaises(ValueError):
                            decode_result_collection(wire)
                    else:
                        decoded = decode_result_collection(wire).collection_payloads[0]
                        self.assertEqual((decoded.image if kind == 3 else decoded.table).shape, (rows, columns))

    def test_all_payload_kinds_decode_from_wire_format(self) -> None:
        # A fixed frame independent of Python encoding preserves the C++ wire contract.
        wire = struct.pack("<I6QfII", RESULT_MAGIC, 9, 42, 1000, 2**63 + 7, 2**60 + 3, 2**60 + 9, 1.25, 2, 1)
        wire += struct.pack("<BH", 3, 15) + b"gpr_accumulator"
        wire += struct.pack("<II11f", 3, 2, 0, 1, 2, 0, 1, 0, 1, 2, 3, 4, 5)
        wire += struct.pack("<BH", 4, 10) + b"gpr_points"
        wire += struct.pack("<II6f", 3, 2, 0, 1, 2, 3, 4, 5)
        wire += struct.pack("<III", 1, 0, 2)  # input, output, payload count
        wire += struct.pack("<BH", 1, 5) + b"bscan"
        wire += struct.pack("<I6f", 2, 1, 2, 1, 1, 2, -2)
        wire += struct.pack("<BH", 2, 3) + b"snr" + struct.pack("<f", 2.5)
        decoded = decode_result_collection(wire)
        self.assertEqual((decoded.collection_id, decoded.monotonic_ns, decoded.processing_duration_ns), (9, 42, 1000))
        self.assertEqual(wire[:4], b"RSL3")
        self.assertEqual(struct.calcsize("<I6QfII"), 64)
        self.assertEqual((decoded.settings_revision, decoded.capture_start_ns, decoded.capture_end_ns, decoded.speed_m_s),
                         (2**63 + 7, 2**60 + 3, 2**60 + 9, 1.25))
        for invalid in (b"RSL2" + wire[4:], wire[:63], wire + b"\0"):
            with self.subTest(invalid_size=len(invalid)), self.assertRaises(ValueError):
                decode_result_collection(invalid)
        image, table = decoded.collection_payloads
        self.assertEqual(image.processing_name, "gpr_accumulator")
        self.assertEqual(table.processing_name, "gpr_points")
        np.testing.assert_array_equal(image.image, np.arange(6, dtype=np.float32).reshape(2, 3))
        np.testing.assert_array_equal(image.image_x_axis, [0, 1, 2])
        np.testing.assert_array_equal(image.image_y_axis, [0, 1])
        np.testing.assert_array_equal(table.table, [[0, 1, 2], [3, 4, 5]])
        block = decoded.blocks[0]
        self.assertEqual(block.combo, ComboKey(1, 0))
        np.testing.assert_array_equal(block.payloads[0].frequency_hz, [1, 2])
        np.testing.assert_array_equal(block.payloads[0].trace, [1 + 1j, 2 - 2j])
        self.assertEqual(block.payloads[1].scalar_value, 2.5)


class CorruptFrameTest(unittest.TestCase):
    """Any corruption surfaces as ValueError (the single catchable contract)."""

    def _valid_trace_bytes(self) -> bytes:
        return serialize_trace_collection(
            SweepCollection(collection_id=1, monotonic_ns=1, traces=[_trace(0, 0, 3)]), RAW_MAGIC
        )

    def test_bad_magic(self) -> None:
        corrupt = b"\x00\x00\x00\x00" + self._valid_trace_bytes()[4:]
        with self.assertRaises(ValueError):
            decode_trace_collection(corrupt, RAW_MAGIC)

    def test_truncated_buffer(self) -> None:
        with self.assertRaises(ValueError):
            decode_trace_collection(self._valid_trace_bytes()[:18], RAW_MAGIC)

    def test_absurd_trace_count(self) -> None:
        # Claims 1e6 traces but supplies none -> the first trace read runs off the end.
        corrupt = struct.pack("<IQQI", RAW_MAGIC, 1, 1, 1_000_000)
        with self.assertRaises(ValueError):
            decode_trace_collection(corrupt, RAW_MAGIC)

    def test_unsupported_payload_kind(self) -> None:
        corrupt = struct.pack("<I6QfII", RESULT_MAGIC, 1, 1, 0, 0, 1, 2, 0., 1, 0) + struct.pack("<BH", 99, 0)
        with self.assertRaises(ValueError):
            decode_result_collection(corrupt)

    def test_invalid_utf8_name(self) -> None:
        corrupt = struct.pack("<I6QfII", RESULT_MAGIC, 1, 1, 0, 0, 1, 2, 0., 1, 0) + struct.pack("<BH", 2, 1) + b"\xff"
        with self.assertRaises(ValueError):
            decode_result_collection(corrupt)


class _RingTestCase(unittest.TestCase):
    """Base case that creates/cleans named /dev/shm rings."""

    def _ring_name(self, suffix: str = "") -> str:
        return f"/radar_test_{self._testMethodName}{suffix}"

    def _writer(self, name: str, capacity: int, slot_size: int) -> ShmRingWriter:
        with suppress(OSError):
            (Path("/dev/shm") / name[1:]).unlink()  # drop a leftover from a crashed run
        writer = ShmRingWriter(name, capacity, slot_size)
        self.addCleanup(self._cleanup, name, writer)
        return writer

    def _reader(self, name: str, **kw: float) -> ShmRingReader:
        reader = ShmRingReader(name, **kw)
        self.addCleanup(self._safe_close, reader)
        return reader

    @staticmethod
    def _safe_close(obj: ShmRingReader | ShmRingWriter) -> None:
        with suppress(Exception):
            obj.close()

    @staticmethod
    def _cleanup(name: str, writer: ShmRingWriter) -> None:
        with suppress(Exception):
            writer.close()
        with suppress(OSError):
            (Path("/dev/shm") / name[1:]).unlink()


class RingRoundTripTest(_RingTestCase):
    def test_fifo_round_trip(self) -> None:
        name = self._ring_name()
        writer = self._writer(name, capacity=8, slot_size=64)
        reader = self._reader(name)
        payloads = [f"frame{i}".encode() for i in range(5)]
        for p in payloads:
            self.assertTrue(writer.push(p))
        self.assertEqual([reader.pop_payload() for _ in payloads], payloads)
        self.assertIsNone(reader.pop_payload())  # empty afterwards

    def test_oversized_payload_rejected(self) -> None:
        writer = self._writer(self._ring_name(), capacity=4, slot_size=8)
        self.assertFalse(writer.push(b"x" * 9))  # larger than the slot


class RingOverflowTest(_RingTestCase):
    def test_overwrite_during_read_discards_torn_payload(self) -> None:
        from python_app.orchestration.shm import atomics

        name = self._ring_name()
        writer = self._writer(name, capacity=1, slot_size=8)
        reader = self._reader(name)
        writer.push(b"old")
        fence = atomics.fence
        overwritten = False

        def overwrite_after_copy():
            nonlocal overwritten
            fence()
            if not overwritten:
                overwritten = True
                writer.push(b"new")

        with mock.patch.object(atomics, "fence", side_effect=overwrite_after_copy):
            self.assertIsNone(reader.pop_payload())
        self.assertEqual(reader.pop_payload(), b"new")

    def test_reader_cannot_rewind_cursor_after_producer_laps_during_commit(self) -> None:
        from python_app.orchestration.shm import atomics

        name = self._ring_name()
        writer = self._writer(name, capacity=4, slot_size=64)
        reader = self._reader(name)
        writer.push(b"f0")
        advance = atomics.advance

        def lap_before_reader_commit(buffer, offset, desired):
            if buffer is reader._mmap and desired == 1:
                for i in range(1, 6):
                    writer.push(f"f{i}".encode())
            advance(buffer, offset, desired)

        with mock.patch.object(atomics, "advance", side_effect=lap_before_reader_commit):
            self.assertEqual(reader.pop_payload(), b"f0")
        self.assertEqual(reader._read_u64(32), 2)
        self.assertEqual([reader.pop_payload() for _ in range(4)], [b"f2", b"f3", b"f4", b"f5"])

    def test_latest_wins_drops_oldest(self) -> None:
        name = self._ring_name()
        writer = self._writer(name, capacity=4, slot_size=64)
        reader = self._reader(name)
        for i in range(7):  # 3 more than capacity
            self.assertTrue(writer.push(f"f{i}".encode()))
        # The 4 newest survive; the 3 oldest were overwritten.
        survivors = []
        while (item := reader.pop_payload()) is not None:
            survivors.append(item)
        self.assertEqual(survivors, [b"f3", b"f4", b"f5", b"f6"])


class RingOpenTest(_RingTestCase):
    def test_same_size_segment_is_reused_only_with_matching_header(self) -> None:
        for corrupt_header in (False, True):
            with self.subTest(corrupt_header=corrupt_header):
                name = self._ring_name(str(corrupt_header))
                writer = self._writer(name, capacity=1, slot_size=8)
                reader = self._reader(name)
                writer.push(b"old")
                if corrupt_header:
                    writer._write_u32(8, 99)
                path = Path("/dev/shm") / name[1:]
                old_inode = path.stat().st_ino
                replacement = ShmRingWriter(name, 1, 8)
                self.addCleanup(self._safe_close, replacement)
                new_reader = self._reader(name)
                self.assertEqual(path.stat().st_ino == old_inode, not corrupt_header)
                self.assertEqual(new_reader.pop_payload(), None if corrupt_header else b"old")
                if corrupt_header:
                    self.assertEqual(reader.pop_payload(), b"old")
                replacement.push(b"new")
                self.assertEqual(new_reader.pop_payload(), b"new")
                replacement.close()
                replacement.close()

    def test_invalid_geometry_and_names_fail_before_opening_files(self) -> None:
        for capacity, slot in ((True, 8), (1.5, 8), (1, 9), (2 ** 32, 8)):
            with self.subTest(capacity=capacity, slot=slot), self.assertRaises(ValueError):
                ShmRingWriter(self._ring_name(), cast(int, capacity), slot)
        for name in ("/", "/../tmp/ring", "/nested/ring", "/bad\x00ring"):
            with self.subTest(name=name):
                with self.assertRaises(ValueError):
                    ShmRingWriter(name, 1, 8)
                with self.assertRaises(ValueError):
                    ShmRingReader(name, open_timeout_s=0)

    def test_failed_mapping_closes_opened_file(self) -> None:
        from python_app.orchestration.shm import ring_writer

        name = self._ring_name()
        opened = []
        fdopen = ring_writer.os.fdopen

        def track_fd(fd, *args, **kwargs):
            handle = fdopen(fd, *args, **kwargs)
            opened.append(handle)
            return handle

        self.addCleanup(lambda: (Path('/dev/shm') / name[1:]).unlink(missing_ok=True))
        with mock.patch.object(ring_writer.os, 'fdopen', side_effect=track_fd), \
                mock.patch.object(ring_writer.mmap, 'mmap', side_effect=OSError('mapping failed')), \
                self.assertRaises(OSError):
            ShmRingWriter(name, 1, 8)
        self.assertTrue(opened[0].closed)

    def test_owner_replaces_wrong_size_without_truncating_existing_mappings(self) -> None:
        name = self._ring_name()
        writer = self._writer(name, capacity=8, slot_size=64)
        reader = self._reader(name)
        writer.push(b"old frame")
        path = Path("/dev/shm") / name[1:]
        old_inode = path.stat().st_ino

        replacement = ShmRingWriter(name, capacity=1, slot_size_bytes=8)
        self.addCleanup(self._safe_close, replacement)

        self.assertNotEqual(path.stat().st_ino, old_inode)
        self.assertEqual(reader.pop_payload(), b"old frame")
        self.assertTrue(replacement.push(b"new"))
        self.assertEqual(self._reader(name).pop_payload(), b"new")

    def test_waits_for_created_but_unallocated_ring(self) -> None:
        name = self._ring_name()
        path = Path("/dev/shm") / name[1:]
        path.touch()
        self.addCleanup(lambda: path.unlink(missing_ok=True))
        writers = []

        def initialize() -> None:
            time.sleep(0.02)
            writers.append(ShmRingWriter(name, capacity=1, slot_size_bytes=8))

        creator = threading.Thread(target=initialize)
        creator.start()

        def cleanup_creator() -> None:
            creator.join()
            for writer in writers:
                writer.close()

        self.addCleanup(cleanup_creator)
        reader = self._reader(name)
        creator.join()

        self.assertEqual(reader.capacity, 1)
        self.assertEqual(reader.slot_size_bytes, 8)

    def test_header_geometry_outside_mapping_is_rejected(self) -> None:
        name = self._ring_name()
        writer = self._writer(name, capacity=1, slot_size=8)
        writer._write_u32(12, 100)
        with self.assertRaisesRegex(RuntimeError, "geometry mismatch"):
            ShmRingReader(name)

    def test_missing_ring_fails_fast(self) -> None:
        with self.assertRaises(FileNotFoundError):
            ShmRingReader("/radar_test_definitely_missing", open_timeout_s=0.1, open_poll_s=0.02)

    def test_incompatible_header_rejected(self) -> None:
        name = "/radar_test_bad_header"
        path = Path("/dev/shm") / name[1:]
        path.write_bytes(b"\x00" * 128)  # right size, wrong magic
        self.addCleanup(lambda: path.unlink(missing_ok=True))
        with self.assertRaises(RuntimeError):
            ShmRingReader(name)


if __name__ == "__main__":
    unittest.main()