from typing import NamedTuple, Generator, Any

import numpy as np

CAMERA_NODE_TYPES = ('unknown', 'depth', 'ir', 'rgb', 'rgb_luminance', 'rgb_mapped')

HEADER_BYTES = 32
BYTES_PER_SHORT = 2
TIMESTAMP_BYTES = 8
RAW_UNRELIABLE_PIXEL_VALUE = -11

ENDIANNESS = 'little'

# Pixels are stored as signed 16-bit integers with little endianness.
# However, the only valid negative value is `-11` (unreliable pixel), which should just be converted into the
# corresponding unsigned value, without any loss of information
SOURCE_NUMPY_DATATYPE = np.dtype('<i2')
TARGET_NUMPY_DATATYPE = np.dtype('=u2')  # TODO: is this cross-platform?
MIN_PIXEL_VALUE = np.iinfo(TARGET_NUMPY_DATATYPE).min
MAX_PIXEL_VALUE = np.iinfo(TARGET_NUMPY_DATATYPE).max
UNRELIABLE_PIXEL_VALUE = np.cast[TARGET_NUMPY_DATATYPE](RAW_UNRELIABLE_PIXEL_VALUE).min()

XHDFSHeader = NamedTuple(typename='XHDFSHeader',
                         fields=[('n_frames', int),
                                 ('width', int),
                                 ('height', int),
                                 ('version_number', int),
                                 ('camera_node_type', str)])


def read_unsigned_integer(fileobj, n_bytes: int = 4, signed=False) -> int:
    return int.from_bytes(fileobj.read(n_bytes), ENDIANNESS, signed=signed)


def bytes_to_image(frame_bytes: bytes, width: int, height: int) -> np.ndarray:
    frame_bytes = frame_bytes[8:]  # Cut off the timestamp

    # Height & width in that order match what OpenCV expects
    frame_image = np.frombuffer(frame_bytes, SOURCE_NUMPY_DATATYPE).reshape(height, width)
    frame_image = frame_image.astype(TARGET_NUMPY_DATATYPE, casting='unsafe', copy=False)
    frame_image[frame_image == UNRELIABLE_PIXEL_VALUE] = 0

    return frame_image


class XHDFS:
    def __init__(self, fileobj):
        self._fileobj = fileobj

        self._header = None
        self._pixels_per_frame = -1

    def __enter__(self):
        header = self._read_header()

        self._header = header
        self._pixels_per_frame = header.width * header.height * BYTES_PER_SHORT + TIMESTAMP_BYTES

        return self

    def _read_header(self) -> XHDFSHeader:
        n_frames = read_unsigned_integer(self._fileobj)
        width = read_unsigned_integer(self._fileobj)
        height = read_unsigned_integer(self._fileobj)

        # Skip `CameraNodeId` & `WithBodies`
        _ = self._fileobj.read(4 + 1)

        version = read_unsigned_integer(self._fileobj)
        node_type = read_unsigned_integer(self._fileobj, n_bytes=2)
        node_type = CAMERA_NODE_TYPES[node_type]

        # Skip `CameraNodeSerial` & 3 "spare" bytes
        _ = self._fileobj.read(6 + 3)

        header = XHDFSHeader(n_frames, width, height, version, node_type)
        return header

    def _read_frame_bytes(self) -> bytes:
        return self._fileobj.read(self._pixels_per_frame)

    def frame_sequence(self) -> Generator[np.ndarray, None, None]:
        w, h = self._header.width, self._header.height
        for _ in range(self._header.n_frames):
            frame_bytes = self._read_frame_bytes()
            frame = bytes_to_image(frame_bytes, w, h)
            del frame_bytes
            yield frame
        pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._fileobj.close()