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

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

# Parâmetros de entrada
json_data = None
output_folder = 'Python/Output/'
max_readings = int(sys.argv[1])
porta = sys.argv[2]
url = sys.argv[3]
arquivoSaida = sys.argv[4]
mostrar_linhas = sys.argv[5] == "1"
mqtt_topic = sys.argv[6]
model_folder = sys.argv[7]
script_version = sys.argv[8]

# Flask App
app = Flask(__name__)

# Caminho do modelo convertido
blob_path = os.path.join(model_folder, f'model-{script_version}.blob')

def detect_oak():
    width, height, model_resolution = 1280, 720, 640
    labelMap = ['chenopodio', 'grama-azul', 'tiririca', 'erva-daninha', 'videira', 'mostarda-preta', 'milho', 'amaranta', 'farinha-seca', 'guanxuma']

    pipeline = dai.Pipeline()

    # Criando um nó de câmera otimizado
    cam = pipeline.create(dai.node.ColorCamera)
    cam.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
    cam.setVideoSize(width, height)
    cam.setInterleaved(False)
    cam.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
    cam.setFps(35)  # FPS aumentado
    cam.setIspScale(2, 3)  # Reduz carga no pipeline

    # Criando um nó de manipulação de imagem otimizado
    manip = pipeline.create(dai.node.ImageManip)
    manip.initialConfig.setResize(model_resolution, model_resolution)
    manip.initialConfig.setKeepAspectRatio(True)
    manip.initialConfig.setFrameType(dai.RawImgFrame.Type.RGB888p)  # Força formato RGB correto
    manip.setMaxOutputFrameSize(1228800)
    cam.video.link(manip.inputImage)

    # Criando um nó de inferência YOLO otimizado
    nn = pipeline.create(dai.node.YoloDetectionNetwork)
    nn.setBlobPath(blob_path)
    nn.setConfidenceThreshold(0.25)
    nn.setNumClasses(len(labelMap))
    nn.setCoordinateSize(4)
    nn.setIouThreshold(0.45)
    # 🔹 Define âncoras e máscaras corretamente
    nn.setAnchors([
        12, 16, 19, 36, 40, 28,  
        36, 75, 76, 55, 72, 146, 
        142, 110, 192, 243, 459, 401 
    ])

    nn.setAnchorMasks({
        "side52": [0, 1, 2],  
        "side26": [3, 4, 5],  
        "side13": [6, 7, 8]   
    })
    nn.setNumInferenceThreads(2)  # Usa 2 threads para inferência
    nn.input.setBlocking(False)  # Não bloqueia a entrada de frames
    #nn.setReusePreviousInferenceResults(True)  # Reutiliza inferências para não travar o pipeline
    manip.out.link(nn.input)

    # Saída de vídeo otimizada
    xout_cam = pipeline.create(dai.node.XLinkOut)
    xout_cam.setStreamName("video")
    #xout_cam.setMetadataOnly(True)  # Reduz tráfego de vídeo
    cam.video.link(xout_cam.input)

    # Saída da inferência
    xout_nn = pipeline.create(dai.node.XLinkOut)
    xout_nn.setStreamName("detections")
    nn.out.link(xout_nn.input)

    with dai.Device(pipeline) as device:
        video_queue = device.getOutputQueue("video", maxSize=1, blocking=False)
        detections_queue = device.getOutputQueue("detections", maxSize=1, blocking=False)
        readings = []

        while True:
            frame = video_queue.get().getCvFrame()
            in_det = detections_queue.get()
            detections = in_det.detections

            current_readings = []
            for detection in detections:
                x1 = int(detection.xmin * width)
                y1 = int(detection.ymin * height)
                x2 = int(detection.xmax * width)
                y2 = int(detection.ymax * height)
                confidence = float(detection.confidence)
                class_id = int(detection.label)
                descricao = labelMap[class_id] if class_id < len(labelMap) else f"Classe {class_id}"

                detection_info = {
                    'id': class_id,
                    'descricao': descricao,
                    'x': x1,
                    'y': y1,
                    'largura': x2 - x1,
                    'altura': y2 - y1,
                    'confianca': confidence
                }
                current_readings.append(detection_info)

                if mostrar_linhas:
                    label = f'{descricao} {confidence:.2f}'
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

            global json_data
            timestamp = time.time()
            json_data = {'timestamp': timestamp, 'x_max': width, 'y_max': height, 'objetos': current_readings}
            print(json_data)
            readings.append(json_data)

            if len(readings) >= max_readings:
                mensagem_relevante = max(readings, key=lambda m: len(m['objetos']), default=None)
                if not mensagem_relevante:
                    mensagem_relevante = readings[0]
                mqtt_client.publish(mqtt_topic, json.dumps(mensagem_relevante).encode('utf-8'))
                readings.clear()

            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')


def send_script_ready():
    mqtt_client.publish(mqtt_topic, "OK")

def run_flask_server():
    app.run(host='0.0.0.0', port=porta, threaded=True, debug=False)

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

if __name__ == '__main__':
    mqtt_thread = threading.Thread(target=send_script_ready)
    mqtt_thread.start()
    run_flask_server()