from __future__ import annotations

import ctypes
import ctypes.util
import os
from dataclasses import dataclass

import numpy as np


KLIB_TYPE_LIFO = 2
KLIB_MEM_SHM = 2
KLIB_CFMT_RGB = 0
KLIB_DFMT_HWC_8U = 0


class ShmError(RuntimeError):
    pass


class _KlibAttr(ctypes.Structure):
    _fields_ = [
        ("type", ctypes.c_uint32),
        ("mem", ctypes.c_uint32),
        ("cfmt", ctypes.c_uint32),
        ("dfmt", ctypes.c_uint32),
        ("w", ctypes.c_uint32),
        ("h", ctypes.c_uint32),
        ("max_n", ctypes.c_uint32),
    ]


class _ShmCtx(ctypes.Structure):
    _fields_ = [
        ("shm", ctypes.c_void_p),
        ("prev", ctypes.c_uint64),
        ("size", ctypes.c_size_t),
    ]


@dataclass
class ShmFrame:
    data: np.ndarray
    width: int
    height: int
    timestamp_us: int


class ShmClient:
    def __init__(
        self,
        name: str,
        width: int,
        height: int,
        max_n: int = 4,
    ) -> None:
        self._name = name
        self._attr = _KlibAttr(
            type=KLIB_TYPE_LIFO,
            mem=KLIB_MEM_SHM,
            cfmt=KLIB_CFMT_RGB,
            dfmt=KLIB_DFMT_HWC_8U,
            w=width,
            h=height,
            max_n=max_n,
        )
        self._ctx = _ShmCtx()
        self._lib = _load_lib()
        self._connected = False

    def connect(self) -> None:
        if self._connected:
            return

        rv = self._lib.navel_shm_connect(
            ctypes.byref(self._ctx),
            self._name.encode("utf-8"),
            ctypes.byref(self._attr),
        )
        if rv != 0:
            raise ShmError(f"navel_shm_connect failed (err={rv})")

        self._connected = True

    def disconnect(self) -> None:
        if not self._connected:
            return

        rv = self._lib.navel_shm_disconnect(
            ctypes.byref(self._ctx), self._name.encode("utf-8")
        )
        if rv != 0:
            raise ShmError(f"navel_shm_disconnect failed (err={rv})")

        self._connected = False

    def get_frame(self, timeout_ms: int = 200) -> ShmFrame:
        if not self._connected:
            raise ShmError("SHM client not connected")

        buf = ctypes.create_string_buffer(self._ctx.size)
        width = ctypes.c_uint32()
        height = ctypes.c_uint32()
        timestamp = ctypes.c_uint64()

        rv = self._lib.navel_shm_get_frame(
            ctypes.byref(self._ctx),
            int(timeout_ms * 1000),
            buf,
            len(buf),
            ctypes.byref(width),
            ctypes.byref(height),
            ctypes.byref(timestamp),
        )
        if rv != 0:
            raise ShmError(f"navel_shm_get_frame failed (err={rv})")

        data = np.frombuffer(buf, dtype=np.uint8)
        expected = int(width.value) * int(height.value) * 3
        if data.size < expected:
            raise ShmError(
                f"frame buffer too small: {data.size} < {expected}"
            )

        data = data[:expected].reshape((height.value, width.value, 3))

        return ShmFrame(
            data=data,
            width=int(width.value),
            height=int(height.value),
            timestamp_us=int(timestamp.value),
        )


def _init_ctype_signatures(lib: ctypes.CDLL) -> None:
    lib.navel_shm_connect.argtypes = [
        ctypes.POINTER(_ShmCtx),
        ctypes.c_char_p,
        ctypes.POINTER(_KlibAttr),
    ]
    lib.navel_shm_connect.restype = ctypes.c_int

    lib.navel_shm_disconnect.argtypes = [
        ctypes.POINTER(_ShmCtx),
        ctypes.c_char_p,
    ]
    lib.navel_shm_disconnect.restype = ctypes.c_int

    lib.navel_shm_get_frame.argtypes = [
        ctypes.POINTER(_ShmCtx),
        ctypes.c_int,
        ctypes.c_void_p,
        ctypes.c_size_t,
        ctypes.POINTER(ctypes.c_uint32),
        ctypes.POINTER(ctypes.c_uint32),
        ctypes.POINTER(ctypes.c_uint64),
    ]
    lib.navel_shm_get_frame.restype = ctypes.c_int


_load_lib_cached: ctypes.CDLL | None = None


def _load_lib() -> ctypes.CDLL:
    global _load_lib_cached
    if _load_lib_cached is not None:
        return _load_lib_cached

    env_path = os.getenv("NAVEL_SHM_LIB")
    if env_path:
        lib = ctypes.CDLL(env_path)
        _init_ctype_signatures(lib)
        _load_lib_cached = lib
        return lib

    here = os.path.abspath(os.path.dirname(__file__))
    local_path = os.path.join(here, "libnavel_shm.so")
    if os.path.exists(local_path):
        lib = ctypes.CDLL(local_path)
        _init_ctype_signatures(lib)
        _load_lib_cached = lib
        return lib

    lib_name = ctypes.util.find_library("navel_shm")
    if lib_name:
        lib = ctypes.CDLL(lib_name)
        _init_ctype_signatures(lib)
        _load_lib_cached = lib
        return lib

    raise ShmError(
        "Could not load libnavel_shm.so. "
        "Build it with `make build-shm` in krill/jetson/sdk/python3 "
        "or set NAVEL_SHM_LIB to its path."
    )
