from __future__ import annotations

from dataclasses import dataclass

from navel._shm import ShmClient, ShmFrame


HEAD_CAM_SHM_NAME = "bodyd_cam_head"
HEAD_CAM_WIDTH = 640
HEAD_CAM_HEIGHT = 480


@dataclass
class Frame:
    data: object
    width: int
    height: int
    timestamp_us: int


class HeadCamera:
    def __init__(self) -> None:
        self._client = ShmClient(
            name=HEAD_CAM_SHM_NAME,
            width=HEAD_CAM_WIDTH,
            height=HEAD_CAM_HEIGHT,
            max_n=4,
        )

    def __enter__(self) -> "HeadCamera":
        self.open()
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    def open(self) -> None:
        self._client.connect()

    def close(self) -> None:
        self._client.disconnect()

    def get_frame(self, timeout_ms: int = 200) -> Frame:
        frame: ShmFrame = self._client.get_frame(timeout_ms=timeout_ms)
        return Frame(
            data=frame.data,
            width=frame.width,
            height=frame.height,
            timestamp_us=frame.timestamp_us,
        )
