import tempfile
import unittest
from pathlib import Path

from udp_probe import analyze_udp_records, save_udp_records


def mik_payload(flags, sequence, packet_number, value, data=b""):
    return bytes((0, flags, sequence, packet_number)) + value.to_bytes(4, "little") + data


def mik_frame(width=4, height=2):
    pixels = bytes(range(width * height))
    image_header = width.to_bytes(2, "little") + height.to_bytes(2, "little") + bytes((1, 0, 0, 0))
    return (0).to_bytes(4, "little") + image_header + pixels


def record(index, payload, source=("192.168.0.10", 40000)):
    return {
        "timestamp_ns": 1_000_000_000 + index * 1_000_000,
        "address": source,
        "payload": payload,
    }


class UdpProbeTests(unittest.TestCase):
    def test_detects_complete_mik_video_array(self):
        frame = mik_frame()
        split = 9
        payloads = [
            mik_payload(2, 7, 0, len(frame)),
            mik_payload(0, 7, 1, 0, frame[:split]),
            mik_payload(0, 7, 2, split, frame[split:]),
            mik_payload(1, 7, 3, len(frame)),
        ]

        result = analyze_udp_records([record(index, payload) for index, payload in enumerate(payloads)])

        self.assertEqual(result["detected"]["kind"], "mik_video")
        self.assertEqual(result["detected"]["confidence"], 100)
        self.assertEqual(result["detected"]["frame"]["width"], 4)
        self.assertEqual(result["detected"]["recommended"]["source_mode"], "udp_mik_live")

    def test_detects_selected_raw_frame_with_separator_packet(self):
        payloads = [b"\x01\x02\x03\x04", b"\x05\x06\x07\x08", b"\xff"]
        result = analyze_udp_records(
            [record(index, payload) for index, payload in enumerate(payloads)],
            width=4,
            height=2,
            separator=255,
            frame_encoding="gray8",
        )

        self.assertEqual(result["detected"]["kind"], "raw_delimited")
        self.assertEqual(result["detected"]["confidence"], 99)
        self.assertEqual(result["detected"]["recommended"]["frame_encoding"], "gray8")

    def test_saved_dump_preserves_every_payload_byte(self):
        records = [record(0, b"\x00\x01"), record(1, b"\xfe\xff")]
        with tempfile.TemporaryDirectory() as tmp:
            dump_path, report_path, report = save_udp_records(records, 40404, Path(tmp))
            expected = (
                (40404).to_bytes(2, "little") + (2).to_bytes(2, "little") + b"\x00\x01"
                + (40404).to_bytes(2, "little") + (2).to_bytes(2, "little") + b"\xfe\xff"
            )

            self.assertEqual(dump_path.read_bytes(), expected)
            self.assertTrue(report_path.is_file())
            self.assertEqual(report["dump_size"], len(expected))
            self.assertEqual(len(report["dump_sha256"]), 64)


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