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

# === CONFIGURAÇÕES ===
PASTA_SAIDA = "dataset"
DIMENSAO_FINAL = (1920, 1080)
QUALIDADE_JPEG = 95
TECLA_SALVAR = ord('s')

# Área de exibição
ALTURA_JANELA = 1080
LARGURA_JANELA = 1920

EXPOSICAO_ATUAL = [6000]
ISO_ATUAL = [400]

os.makedirs(PASTA_SAIDA, exist_ok=True)

# === PIPELINE ===
pipeline = dai.Pipeline()

cam_rgb = pipeline.createColorCamera()
cam_rgb.setPreviewSize(DIMENSAO_FINAL)
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)

control_in = pipeline.createXLinkIn()
control_in.setStreamName("control")
control_in.out.link(cam_rgb.inputControl)

# === EXECUÇÃO ===
with dai.Device(pipeline) as device:
    print("[INFO] Pressione 's' para salvar | '+' e '-' para ISO | 'q' para sair.")

    controlQueue = device.getInputQueue("control")
    rgbQueue = device.getOutputQueue("rgb", maxSize=4, blocking=False)

    def aplicar_controle_manual(exposicao, iso):
        ctrl = dai.CameraControl()
        ctrl.setManualExposure(exposicao, iso)
        ctrl.setManualFocus(130)
        controlQueue.send(ctrl)
        print(f"[⚙️] Exposicao: {exposicao} µs | ISO: {iso}")

    aplicar_controle_manual(EXPOSICAO_ATUAL[0], ISO_ATUAL[0])
    ultimo_frame_salvo = np.zeros((DIMENSAO_FINAL[1], DIMENSAO_FINAL[0], 3), dtype=np.uint8)

    # Fullscreen
    cv2.namedWindow("Coletor de Dataset", cv2.WINDOW_NORMAL)
    cv2.setWindowProperty("Coletor de Dataset", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)

    while True:
        in_rgb = rgbQueue.get()
        frame = in_rgb.getCvFrame()

        # Redimensiona o frame ao vivo
        frame_vivo = cv2.resize(frame, (1280, 960))

        # Redimensiona a última captura
        captura_redimensionada = cv2.resize(ultimo_frame_salvo, (512, 512))

        # Cria o canvas vazio
        canvas = np.zeros((ALTURA_JANELA, LARGURA_JANELA, 3), dtype=np.uint8)

        # Posiciona o ao vivo à esquerda
        canvas[60:1020, 30:1310] = frame_vivo

        # Posiciona a última imagem salva à direita
        canvas[60:572, 1400:1912] = captura_redimensionada

        # Adiciona textos no lado direito
        cv2.putText(canvas, f"ISO: {ISO_ATUAL[0]} | Exposicao: {EXPOSICAO_ATUAL[0]}", (1350, 620), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 255, 0), 2)
        cv2.putText(canvas, "Pressione 's' para salvar", (1350, 680), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 1)
        cv2.putText(canvas, "'+' ou '-' para ajustar ISO", (1350, 720), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 1)
        cv2.putText(canvas, "'m' ou 'n' para ajustar a exposicao", (1350, 760), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 1)
        cv2.putText(canvas, "'q' para sair", (1350, 800), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 1)

        cv2.imshow("Coletor de Dataset", canvas)

        key = cv2.waitKey(1) & 0xFF

        if key == TECLA_SALVAR:
            now = datetime.now()
            timestamp = now.strftime("%Y%m%d_%H%M%S_%f")[:-3]
            caminho = os.path.join(PASTA_SAIDA, f"img_{timestamp}.jpg")
            
            # Gera timestamp e nome-base
            now = datetime.now()
            timestamp = now.strftime("%Y%m%d_%H%M%S_%f")[:-3]
            nome_base = f"img_{timestamp}.jpg"

            # Original (sem resize)
            os.makedirs(os.path.join("dataset", "original", "images"), exist_ok=True)
            cv2.imwrite(os.path.join("dataset", "original", "images", nome_base), frame, [cv2.IMWRITE_JPEG_QUALITY, QUALIDADE_JPEG])

            # Demais resoluções
            #resolucoes = {
            #    "512x512": (512, 512),
            #    "768x768": (768, 768),
            #    "384x384": (384, 384),
            #}
            #for nome_pasta, dim in resolucoes.items():
            #    path_pasta = os.path.join("dataset", nome_pasta)
            #    os.makedirs(path_pasta, exist_ok=True)
            #    imagem = cv2.resize(frame, dim, interpolation=cv2.INTER_AREA)
            #    cv2.imwrite(os.path.join(path_pasta, nome_base), imagem, [cv2.IMWRITE_JPEG_QUALITY, QUALIDADE_JPEG])
            
            print(f"[✔] Imagem salva: {caminho}")
            ultimo_frame_salvo = frame.copy()

        elif key == ord('+') or key == ord('='):
            ISO_ATUAL[0] = min(1600, ISO_ATUAL[0] + 50)
            aplicar_controle_manual(EXPOSICAO_ATUAL[0], ISO_ATUAL[0])

        elif key == ord('-'):
            ISO_ATUAL[0] = max(100, ISO_ATUAL[0] - 50)
            aplicar_controle_manual(EXPOSICAO_ATUAL[0], ISO_ATUAL[0])

        elif key == ord('m'):
            EXPOSICAO_ATUAL[0] = min(10000, EXPOSICAO_ATUAL[0] + 500)
            aplicar_controle_manual(EXPOSICAO_ATUAL[0], ISO_ATUAL[0])

        elif key == ord('n'):
            EXPOSICAO_ATUAL[0] = max(1000, EXPOSICAO_ATUAL[0] - 500)
            aplicar_controle_manual(EXPOSICAO_ATUAL[0], ISO_ATUAL[0])

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

    cv2.destroyAllWindows()