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

from PIL import Image  # pip install pillow

# ====== AJUSTE AQUI ======
SDK_DIR = os.path.join(os.path.dirname(__file__), "dlls")
DLL_NAME = "VT_SDK64.dll"
TIMEOUT_MS = 2000

OUT_DIR = os.path.join(os.path.dirname(__file__), "out")
os.makedirs(OUT_DIR, exist_ok=True)

# nomes de saída
OUT_RAW = os.path.join(OUT_DIR, "frame_000001.raw")         # RAW cru (do buffer)
OUT_BMP = os.path.join(OUT_DIR, "frame_000001_preview.bmp") # preview gerado pelo SDK
OUT_PNG = os.path.join(OUT_DIR, "frame_000001_preview.png") # preview final pra rotular

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

# ====== CONSTANTES ======
DEVICE_UDEF  = 0
DEVICE_INDEX = 0

DATA_RAW = 0

# Observação: no seu teste anterior, "filetype=0" gerou um BMP.
# Então aqui a gente assume que 0 = BMP (preview). O RAW vamos salvar manualmente do buffer.
FILE_BMP = 0


# ====== STRUCTS (compatível com o seu uso atual) ======
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),      # alinhamento
        ("pBufPtr",  C.POINTER(C.c_ubyte)),
        ("lFrameStatus", W.DWORD),
        ("lPixType", W.DWORD),
        ("lTimeStamp", W.DWORD),
        ("_reserve", W.DWORD * 8),
    ]


# ====== PROTÓTIPOS ======
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_SingleFrameSavefile.argtypes = [W.HANDLE, C.POINTER(VT_FRAMEINFO), C.c_int, C.c_char_p, C.c_int]
dll.VT_SingleFrameSavefile.restype  = C.c_int

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


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


def is_bmp(path: str) -> bool:
    try:
        with open(path, "rb") as f:
            return f.read(2) == b"BM"
    except Exception:
        return False


def main():
    # 1) scan
    n = C.c_ubyte(0)
    ret = dll.VT_DeviceScan(C.byref(n), DEVICE_UDEF)
    print("DeviceScan ret=", ret, "n=", n.value)
    ck(ret, "VT_DeviceScan")
    if n.value == 0:
        raise RuntimeError("Nenhum dispositivo encontrado")

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

    # 3) captura 1 frame (RAW no buffer)
    fi = VT_FRAMEINFO()
    ret = dll.VT_SingleFrameCapture(h, C.byref(fi), DATA_RAW, TIMEOUT_MS, True)
    print("SingleFrameCapture ret=", ret,
          "W,H=", fi.lWidth, fi.lHeight,
          "pixbits=", fi.lPixBits, "buf=", fi.lBufSize,
          "pixType=", fi.lPixType)
    ck(ret, "VT_SingleFrameCapture")

    # 4) salva RAW CRU (do buffer)
    raw_bytes = C.string_at(fi.pBufPtr, fi.lBufSize)
    with open(OUT_RAW, "wb") as f:
        f.write(raw_bytes)
    print(f"RAW cru salvo: {OUT_RAW} | {len(raw_bytes)/1024/1024:.2f} MB")

    # 5) salva preview via SDK (BMP) e converte pra PNG
    ret = dll.VT_SingleFrameSavefile(h, C.byref(fi), FILE_BMP, OUT_BMP.encode("utf-8"), 90)
    print("SingleFrameSavefile (preview BMP) ret=", ret, "->", OUT_BMP)
    ck(ret, "VT_SingleFrameSavefile")

    if not is_bmp(OUT_BMP):
        print("⚠️ Preview não parece BMP (não começa com 'BM'). Mesmo assim vou tentar abrir...")
    img = Image.open(OUT_BMP)
    img.save(OUT_PNG)
    print("Preview PNG salvo:", OUT_PNG)

    # 6) close
    ret = dll.VT_DeviceClose(C.byref(h))
    print("DeviceClose ret=", ret)
    ck(ret, "VT_DeviceClose")

    print("\nOK ✅")
    print(" - RAW cru (treino/inferência):", OUT_RAW)
    print(" - Preview (rotulagem):", OUT_PNG)


if __name__ == "__main__":
    main()