"""Tests for the LAT stream extractor and replayer."""
import io
import os
import struct
import sys
import unittest

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import extract_stream as ex   # noqa: E402
import replay_stream as rp    # noqa: E402

PCAP = os.path.join(HERE, "..", "local", "filter-test2.pcap")


def lat_frame(src, slots):
    lat = bytearray()
    lat += bytes([0x02, len(slots)])          # flags/type, nbr_slots
    lat += struct.pack("<H", 1) + struct.pack("<H", 0x5460)
    lat += bytes([0, 0])
    for (_ds, _ss, _bc, _cr, st, data) in slots:
        lat += bytes([1, 1, len(data), st])
        lat += data
    f = bytearray()
    f += b"\x00" * 6 + ex.mac_bytes(src)
    f += struct.pack(">H", 0x6004)
    f += lat
    return bytes(f)


def write_pcap(path, frames):
    with open(path, "wb") as fh:
        fh.write(b"\xd4\xc3\xb2\xa1" + b"\x00" * 20)
        for f in frames:
            fh.write(struct.pack("<IIII", 0, 0, len(f), len(f)) + f)


class ExtractTest(unittest.TestCase):
    def test_picks_only_host_to_terminal_data(self):
        import tempfile
        frames = [
            lat_frame(ex.HOST_MAC, [(1, 1, 3, 0, 0, b"abc")]),          # ours
            lat_frame(ex.VT1200_MAC, [(1, 1, 3, 0, 0, b"XXX")]),        # terminal -> skip
            lat_frame(ex.HOST_MAC, [(1, 1, 2, 0, 9, b"ST")]),           # Start slot -> skip
            lat_frame(ex.HOST_MAC, [(1, 1, 2, 0, 0, b"de")]),           # ours
        ]
        with tempfile.TemporaryDirectory() as d:
            p = os.path.join(d, "t.pcap")
            write_pcap(p, frames)
            stream, events = ex.extract(p)
        self.assertEqual(stream, b"abcde")
        self.assertEqual([n for _, n in events], [3, 2])

    def test_real_capture(self):
        if not os.path.exists(PCAP):
            self.skipTest("no capture")
        stream, events = ex.extract(PCAP)
        self.assertGreater(len(stream), 1000)
        # the handoff says the stream is escape-free apart from one colour prompt
        self.assertLessEqual(stream.count(0x1b if isinstance(stream, bytes) else b"\x1b"), 2)


class ReplayTest(unittest.TestCase):
    def test_replays_exactly_and_honours_limit(self):
        out = io.BytesIO()
        sent = rp.replay(b"0123456789" * 30, out, chunk=85, pace=0)
        self.assertEqual(sent, 300)
        self.assertEqual(out.getvalue(), b"0123456789" * 30)

        out2 = io.BytesIO()
        data = b"ABCDEFGHIJ"
        rp.replay(data[:7], out2, chunk=3, pace=0)
        self.assertEqual(out2.getvalue(), b"ABCDEFG")

    def test_raw_mode_is_safe_on_non_tty(self):
        rp.make_raw(999999)      # bogus fd: must not raise


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