import depthai as dai
import cv2
import os
from datetime import datetime

# === CONFIGURAÇÕES ===
PASTA_SAIDA = os.path.join("dataset", "oak-d-lite")
DIMENSAO_FINAL = (512, 512)
QUALIDADE_JPEG = 95  # 0 a 100
TECLA_SALVAR = ord('s')  # pressione 's' para salvar

# Cria pasta de saída se não existir
os.makedirs(PASTA_SAIDA, exist_ok=True)

# === PIPELINE OAK ===
pipeline = dai.Pipeline()
cam_rgb = pipeline.createColorCamera()
cam_rgb.setPreviewSize(640, 640)
cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
cam_rgb.setInterleaved(False)
cam_rgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)

xout = pipeline.createXLinkOut()
xout.setStreamName("rgb")
cam_rgb.preview.link(xout.input)

# === EXECUÇÃO ===
with dai.Device(pipeline) as device:
    print("[INFO] Pressione 's' para salvar imagem, 'q' para sair.")
    queue = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)

    while True:
        frame = queue.get().getCvFrame()
        frame_resized = cv2.resize(frame, DIMENSAO_FINAL, interpolation=cv2.INTER_AREA)

        cv2.imshow("OAK-D RGB", frame_resized)
        key = cv2.waitKey(1) & 0xFF

        if key == TECLA_SALVAR:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            caminho = os.path.join(PASTA_SAIDA, f"img_{timestamp}.jpg")
            cv2.imwrite(caminho, frame_resized, [cv2.IMWRITE_JPEG_QUALITY, QUALIDADE_JPEG])
            print(f"[✔] Imagem salva: {caminho}")

        elif key == ord('q'):
            print("[INFO] Encerrando...")
            break

    cv2.destroyAllWindows()