import depthai as dai
import cv2
import numpy as np
import time
import json
import sys
import threading
import paho.mqtt.client as mqtt
import uuid
from flask import Flask, Response

# 🔹 Configurações MQTT
mqtt_client = mqtt.Client(f"client_oak_d_lite_{uuid.uuid4()}")
mqtt_client.connect("localhost", port=1883)

# 🔹 Parâmetros de entrada
output_folder = 'Python/Output/'
mqtt_topic = sys.argv[1]  # Tópico MQTT para envio dos dados de profundidade
porta = int(sys.argv[2])  # Porta do Flask
url = sys.argv[3]  # URL do vídeo
mostrar_linhas = sys.argv[4] == "1"  # Se deve exibir sobreposições na imagem

# 🔹 Flask App
app = Flask(__name__)

# 🔹 Criando o pipeline
pipeline = dai.Pipeline()

# 📷 Câmera RGB
cam_rgb = pipeline.create(dai.node.ColorCamera)
cam_rgb.setPreviewSize(1280, 720)
cam_rgb.setBoardSocket(dai.CameraBoardSocket.RGB)
cam_rgb.setInterleaved(False)

# Criar XLinkOut para saída RGB
xout_rgb = pipeline.create(dai.node.XLinkOut)
xout_rgb.setStreamName("rgb")
cam_rgb.preview.link(xout_rgb.input)

# 🔹 Câmera de Profundidade
mono_left = pipeline.create(dai.node.MonoCamera)
mono_right = pipeline.create(dai.node.MonoCamera)
stereo = pipeline.create(dai.node.StereoDepth)

mono_left.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)
mono_right.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)

mono_left.setBoardSocket(dai.CameraBoardSocket.LEFT)
mono_right.setBoardSocket(dai.CameraBoardSocket.RIGHT)

# Configuração do StereoDepth
stereo.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.HIGH_DENSITY)
stereo.setLeftRightCheck(True)
stereo.setSubpixel(True)

mono_left.out.link(stereo.left)
mono_right.out.link(stereo.right)

# Criar XLinkOut para profundidade
xout_depth = pipeline.create(dai.node.XLinkOut)
xout_depth.setStreamName("depth")
stereo.depth.link(xout_depth.input)

# 🔹 Função para capturar e processar os frames
def process_frames():
    with dai.Device(pipeline) as device:
        rgb_queue = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
        depth_queue = device.getOutputQueue(name="depth", maxSize=4, blocking=False)

        while True:
            timestamp = time.time()

            # 📷 Captura RGB
            rgb_frame = rgb_queue.get().getCvFrame()

            # 📏 Captura Profundidade
            depth_frame = depth_queue.get().getFrame()

            # Normaliza a profundidade para visualização (mapa de calor)
            depth_visual = cv2.normalize(depth_frame, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
            depth_visual = cv2.applyColorMap(depth_visual, cv2.COLORMAP_JET)

            # 📡 Enviar dados via MQTT (Apenas a matriz de profundidade reduzida)
            depth_data = depth_frame.tolist()  # Converte a matriz para lista JSON
            json_data = {
                'timestamp': timestamp,
                'depth_data': depth_data
            }
            mqtt_client.publish(mqtt_topic, json.dumps(json_data).encode('utf-8'))


# 🔹 Função para enviar vídeo via Flask
def generate_video():
    with dai.Device(pipeline) as device:
        rgb_queue = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
        depth_queue = device.getOutputQueue(name="depth", maxSize=4, blocking=False)

        while True:
            frame_data = rgb_queue.get()
            if frame_data is None:
                continue

            timestamp = time.time()

            # 📏 Captura Profundidade
            depth_frame = depth_queue.get().getFrame()

            # Normaliza a profundidade para visualização (mapa de calor)
            depth_visual = cv2.normalize(depth_frame, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
            depth_visual = cv2.applyColorMap(depth_visual, cv2.COLORMAP_JET)

            # 📡 Enviar dados via MQTT (Apenas a matriz de profundidade reduzida)
            depth_data = depth_frame.tolist()  # Converte a matriz para lista JSON
            json_data = {
                'timestamp': timestamp,
                'depth_data': depth_data
            }
            mqtt_client.publish(mqtt_topic, json.dumps(json_data).encode('utf-8'))

            frame = frame_data.getCvFrame()
            ret, buffer = cv2.imencode('.jpg', frame)
            frame = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')

@app.route('/' + url, methods=['GET'])
def video_feed():
    return Response(generate_video(), mimetype='multipart/x-mixed-replace; boundary=frame')

# 🔹 Rodar Flask e processamento em threads separadas
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=porta, threaded=True, debug=False)