import math
import os
import time
import ctypes as C
from ctypes import wintypes as W

import numpy as np
import cv2

# =========================
# CONFIG
# =========================
SDK_DIR = os.path.join(os.path.dirname(__file__), "dlls")
DLL_NAME = "VT_SDK64.dll"

TIMEOUT_MS = 2000
WINDOW_NAME = "GAL5000 4CH Preview (E=AEC toggle, G=AGC toggle, Q=quit)"

# Camera scan/open
DEVICE_UDEF  = 0
DEVICE_INDEX = 0
DATA_RAW     = 0

# RAW geometry (como você já capturou)
RAW_W = 2592
RAW_H = 2056

# Bayer+NIR pattern (2x2):
#   R  G
#   IR B
# => R = [0::2,0::2], G=[0::2,1::2], IR=[1::2,0::2], B=[1::2,1::2]

# =========================
# PARAM IDs (VT_Param.h)
# =========================
BUF_SIZE = 256

PARAM_ID_SENSOR_EXPOSURETIMERAW          = 0x00003010
PARAM_ID_SENSOR_GAINANALOGRAW            = 0x00003020
PARAM_ID_SENSOR_EXPOSUREAUTOENABLE       = 0x00003011
PARAM_ID_SENSOR_GAINANALOGAUTOENABLE     = 0x00003021
PARAM_ID_SENSOR_GAINANALOGAGCMAX         = 0x00003022
PARAM_ID_COMMON_DISPLAYFPS               = 0x00000203  # float R (média efetiva)  :contentReference[oaicite:9]{index=9}


# PARAM_VALUETYPE
VALUE_INT   = 0
VALUE_FLOAT = 1
VALUE_STR   = 2


# =========================
# STRUCTS (mínimo necessário)
# =========================
class VT_FRAMEINFO(C.Structure):
    _fields_ = [
        ("lFrameID", W.DWORD),
        ("lBufSize", W.DWORD),
        ("lWidth",   W.DWORD),
        ("lHeight",  W.DWORD),
        ("lPixBits", C.c_ubyte),
        ("_pad0",    C.c_ubyte * 3),
        ("pBufPtr",  C.POINTER(C.c_ubyte)),
        ("lFrameStatus", W.DWORD),
        ("lPixType", W.DWORD),
        ("lTimeStamp", W.DWORD),
        ("_reserve", W.DWORD * 8),
    ]

class VT_DEVPARAM(C.Structure):
    _fields_ = [
        ("bUseName", W.BOOL),
        ("lParamByID", W.DWORD),
        ("lParamByName", C.c_char * BUF_SIZE),
    ]

def devparam_by_id(pid: int) -> VT_DEVPARAM:
    p = VT_DEVPARAM()
    p.bUseName = False
    p.lParamByID = pid
    p.lParamByName = b""  # <- CORRETO: bytes (fica zerado / string vazia)
    return p

# =========================
# DLL LOAD + prototypes
# =========================
os.add_dll_directory(SDK_DIR)
dll = C.WinDLL(os.path.join(SDK_DIR, DLL_NAME))
print("DLL carregada OK:", dll)

dll.VT_DeviceScan.argtypes = [C.POINTER(C.c_ubyte), C.c_int]
dll.VT_DeviceScan.restype  = C.c_int

dll.VT_DeviceOpen.argtypes = [C.c_void_p, C.POINTER(W.HANDLE), C.c_int, C.c_int]
dll.VT_DeviceOpen.restype  = C.c_int

dll.VT_SingleFrameCapture.argtypes = [W.HANDLE, C.POINTER(VT_FRAMEINFO), C.c_int, C.c_int, W.BOOL]
dll.VT_SingleFrameCapture.restype  = C.c_int

dll.VT_DeviceClose.argtypes = [C.POINTER(W.HANDLE)]
dll.VT_DeviceClose.restype  = C.c_int

# Param API
dll.VT_ParamGetValue.argtypes = [W.HANDLE, VT_DEVPARAM, C.c_void_p, C.c_int]
dll.VT_ParamGetValue.restype  = C.c_int

dll.VT_ParamSetValue.argtypes = [W.HANDLE, VT_DEVPARAM, C.c_void_p, C.c_int]
dll.VT_ParamSetValue.restype  = C.c_int


def ck(ret: int, name: str):
    if ret != 0:
        print(f"{name} falhou, ret={ret}")
        raise RuntimeError(f"{name} falhou, ret={ret}")

def param_set_int(h: W.HANDLE, pid: int, value: int):
    p = devparam_by_id(pid)
    v = C.c_int(value)
    ret = dll.VT_ParamSetValue(h, p, C.byref(v), VALUE_INT)
    ck(ret, f"VT_ParamSetValue({hex(pid)})")

def param_get_int(h: W.HANDLE, pid: int) -> int:
    p = devparam_by_id(pid)
    v = C.c_int(0)
    ret = dll.VT_ParamGetValue(h, p, C.byref(v), VALUE_INT)
    ck(ret, f"VT_ParamGetValue({hex(pid)})")
    return int(v.value)

def param_get_float(h: W.HANDLE, pid: int) -> float:
    p = devparam_by_id(pid)
    v = C.c_float(0)
    ret = dll.VT_ParamGetValue(h, p, C.byref(v), VALUE_FLOAT)
    ck(ret, f"VT_ParamGetValue({hex(pid)})")
    return float(v.value)

def set_bool(h: W.HANDLE, pid: int, enabled: bool):
    param_set_int(h, pid, 1 if enabled else 0)   # boolean no SDK é int 0/1

def param_supported(h, pid, vtype):
    p = devparam_by_id(pid)
    minv = C.c_int()
    maxv = C.c_int()
    inc  = C.c_int()
    ret = dll.VT_ParamGetRange(h, p,
                              C.byref(minv),
                              C.byref(maxv),
                              C.byref(inc),
                              vtype)
    return ret == 0

def capture_raw8(h: W.HANDLE) -> np.ndarray:
    fi = VT_FRAMEINFO()
    ret = dll.VT_SingleFrameCapture(h, C.byref(fi), DATA_RAW, TIMEOUT_MS, True)
    ck(ret, "VT_SingleFrameCapture")

    w, hh = int(fi.lWidth), int(fi.lHeight)
    if w != RAW_W or hh != RAW_H:
        # Se em algum momento você mudar resolução/ROI, aqui te avisa.
        print(f"[WARN] Res mudou: {w}x{hh} (esperado {RAW_W}x{RAW_H})")

    buf = C.string_at(fi.pBufPtr, fi.lBufSize)
    arr = np.frombuffer(buf, dtype=np.uint8)

    # garante reshape correto
    needed = w * hh
    if arr.size < needed:
        arr = np.pad(arr, (0, needed - arr.size), mode="constant", constant_values=0)
    arr = arr[:needed].reshape(hh, w)
    return arr

def make_rgb_preview(raw: np.ndarray, upscale=2):
    R  = raw[0::2, 0::2]
    G  = raw[0::2, 1::2]
    B  = raw[1::2, 1::2]

    # normalização leve só para display (p2-p98)
    def norm8(x):
        lo = np.percentile(x, 2)
        hi = np.percentile(x, 98)
        if hi <= lo + 1:
            return x
        y = (x.astype(np.float32) - lo) * (255.0 / (hi - lo))
        return np.clip(y, 0, 255).astype(np.uint8)

    Rn, Gn, Bn = map(norm8, [R, G, B])

    rgb = np.dstack([Bn, Gn, Rn])  # OpenCV = BGR
    if upscale and upscale != 1:
        rgb = cv2.resize(rgb, (rgb.shape[1]*upscale, rgb.shape[0]*upscale), interpolation=cv2.INTER_NEAREST)
    return rgb


def main():
    # scan
    n = C.c_ubyte(0)
    ret = dll.VT_DeviceScan(C.byref(n), DEVICE_UDEF)
    ck(ret, "VT_DeviceScan")
    if n.value == 0:
        raise RuntimeError("Nenhuma câmera encontrada.")

    # open
    idx = C.c_ubyte(0)
    h = W.HANDLE()
    ret = dll.VT_DeviceOpen(C.byref(idx), C.byref(h), DEVICE_INDEX, DEVICE_UDEF)
    ck(ret, "VT_DeviceOpen")
    print("DeviceOpen OK, handle=", h.value)

    cv2.namedWindow("RGB", cv2.WINDOW_NORMAL)

    aec_on = False
    agc_on = False
    exp_raw = 1500
    gain_a = 0
    gain_d = 0

    has_hw_aec = param_supported(h, PARAM_ID_SENSOR_EXPOSUREAUTOENABLE, VALUE_INT)
    has_hw_agc = param_supported(h, PARAM_ID_SENSOR_GAINANALOGAUTOENABLE, VALUE_INT)
    has_exp_raw = param_supported(h, PARAM_ID_SENSOR_EXPOSURETIMERAW, VALUE_INT)
    has_gain_a = param_supported(h, PARAM_ID_SENSOR_GAINANALOGRAW, VALUE_INT)

    if (has_hw_aec):
        aec_on = bool(param_get_int(h, PARAM_ID_SENSOR_EXPOSUREAUTOENABLE))
    if (has_hw_agc):
        agc_on = bool(param_get_int(h, PARAM_ID_SENSOR_GAINANALOGAUTOENABLE))
    if (has_exp_raw):
        exp_raw = param_get_int(h, PARAM_ID_SENSOR_EXPOSURETIMERAW)
    if (has_gain_a):
        gain_a  = param_get_int(h, PARAM_ID_SENSOR_GAINANALOGRAW)

    print(f"has_hw_aec: {has_hw_aec}, aec_on: {aec_on}\r\nhas_hw_agc: {has_hw_agc}, agc_on: {agc_on}\r\nhas_exp_raw: {has_exp_raw}, exp_raw: {exp_raw}\r\nhas_gain_a: {has_gain_a}, gain_a: {gain_a}")

    while True:
        raw = capture_raw8(h)
        rgb = make_rgb_preview(raw, upscale=2)
        cv2.imshow("RGB", rgb)

        k = cv2.waitKey(1) & 0xFF
        if k in (ord('q'), ord('Q'), 27):
            break
        elif k in (ord('e'), ord('E')):   # exemplo: E alterna AEC do hardware
            aec_on = not aec_on
            set_bool(h, PARAM_ID_SENSOR_EXPOSUREAUTOENABLE, aec_on)
            print("AEC(hw) =", aec_on)
        elif k in (ord('g'), ord('G')):   # G alterna AGC do hardware
            agc_on = not agc_on
            set_bool(h, PARAM_ID_SENSOR_GAINANALOGAUTOENABLE, agc_on)
            print("AGC(hw) =", agc_on)

if __name__ == "__main__":
    main()