import socket
import tempfile
import unittest
from pathlib import Path

from delimited_frame_capture import DelimitedFrameCapture


class DelimitedFrameCaptureTests(unittest.TestCase):
    WIDTH = 4
    HEIGHT = 2
    SEPARATOR = 255

    def frame_bytes(self, offset=0):
        size = self.WIDTH * self.HEIGHT * 3
        return bytes((offset + index) % 200 for index in range(size))

    def test_reads_raw_frames_from_delimited_log(self):
        first = self.frame_bytes()
        second = self.frame_bytes(20)
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "frames.dump"
            path.write_bytes(first + bytes((self.SEPARATOR,)) + second)
            cap = DelimitedFrameCapture(
                source=path,
                separator=self.SEPARATOR,
                encoding="bgr24",
                width=self.WIDTH,
                height=self.HEIGHT,
            )
            try:
                ok1, frame1 = cap.read()
                ok2, frame2 = cap.read()
                self.assertTrue(ok1)
                self.assertTrue(ok2)
                self.assertEqual(frame1.shape, (self.HEIGHT, self.WIDTH, 3))
                self.assertEqual(int(frame2[0, 0, 0]), 20)
            finally:
                cap.release()

    def test_reads_raw_frame_from_live_udp_stream(self):
        frame_data = self.frame_bytes()
        cap = DelimitedFrameCapture(
            host="127.0.0.1",
            port=0,
            separator=self.SEPARATOR,
            encoding="bgr24",
            width=self.WIDTH,
            height=self.HEIGHT,
        )
        sender = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            sender.sendto(
                frame_data + bytes((self.SEPARATOR,)),
                ("127.0.0.1", cap.port),
            )
            ok, frame = cap.read()
            self.assertTrue(ok)
            self.assertEqual(frame.shape, (self.HEIGHT, self.WIDTH, 3))
        finally:
            sender.close()
            cap.release()

    def test_live_raw_frame_may_contain_separator_byte(self):
        frame_data = bytes(range(self.WIDTH * self.HEIGHT))
        cap = DelimitedFrameCapture(
            host="127.0.0.1",
            port=0,
            separator=0,
            encoding="gray8",
            width=self.WIDTH,
            height=self.HEIGHT,
        )
        sender = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            sender.sendto(frame_data[:3], ("127.0.0.1", cap.port))
            sender.sendto(frame_data[3:], ("127.0.0.1", cap.port))
            sender.sendto(b"\x00", ("127.0.0.1", cap.port))
            ok, frame = cap.read()
            self.assertTrue(ok)
            self.assertEqual(frame.shape, (self.HEIGHT, self.WIDTH, 3))
            self.assertEqual(int(frame[0, 0, 0]), 0)
            self.assertEqual(int(frame[-1, -1, 0]), 7)
        finally:
            sender.close()
            cap.release()


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