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

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

model_config = 'C:\\train\\models\\ervas\\ervas.cfg'
model_weights = 'C:\\train\\models\\ervas\\backup\\ervas_final.weights'
labels_path = 'C:\\train\\models\\ervas\\labels.txt'

with open(labels_path, 'rt') as f:
    classes = f.read().rstrip('\n').split('\n')

# Utilizando a GPU com CUDA
net = cv2.dnn.readNetFromDarknet(model_config, model_weights)
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)

# Configurações adicionais
output_folder = 'Python/Output/'
max_readings = int(sys.argv[1])

app = Flask(__name__)

def detect_objects(conf_threshold, nms_threshold, _camera_index):
    cap = cv2.VideoCapture(_camera_index, cv2.CAP_DSHOW)
    width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
    height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
    readings = []

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        blob = cv2.dnn.blobFromImage(frame, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
        net.setInput(blob)
        outs = net.forward(net.getUnconnectedOutLayersNames())

        current_readings = process_detections(outs, width, height, conf_threshold, frame)  # Passando 'frame' como argumento

        global json_data
        timestamp = time.time()
        json_data = {'timestamp': timestamp, 'x_max': width, 'y_max': height, 'objetos': current_readings}
        readings.append(json_data)
        
        mqtt_client.publish(sys.argv[6], json.dumps(json_data).encode('utf-8'))

        manage_output(readings, frame, max_readings)

        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 process_detections(outs, width, height, conf_threshold, frame):  # Recebendo 'frame' como argumento
    current_readings = []
    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > conf_threshold:
                x, y, w, h = calculate_box(detection, width, height)
                detection_info = create_detection_info(class_id, x, y, w, h, confidence)
                current_readings.append(detection_info)
                if sys.argv[5] == "1":
                    draw_predictions(frame, class_id, confidence, x, y, w, h)  # 'frame' já está disponível aqui
    return current_readings

def draw_predictions(frame, class_id, confidence, x, y, w, h):  # 'frame' é utilizado aqui
    cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
    cv2.putText(frame, f'{classes[class_id]} {confidence:.2f}', (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

def calculate_box(detection, width, height):
    center_x, center_y, w, h = detection[0:4] * np.array([width, height, width, height])
    x, y = int(center_x - w / 2), int(center_y - h / 2)
    return x, y, int(w), int(h)

def create_detection_info(class_id, x, y, w, h, confidence):
    return {
        'id': int(class_id),
        'descricao': classes[class_id],
        'x': x,
        'y': y,
        'largura': w,
        'altura': h,
        'confianca': float(confidence)
    }

def manage_output(readings, frame, max_readings):
    if len(readings) > max_readings:
        readings.pop(0)
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    with open(output_folder + sys.argv[4], 'w') as file:
        json.dump(readings, file, indent=4)

@app.route('/' + sys.argv[3], methods=['GET'])
def video_feed():
    conf_threshold = float(request.args.get('conf_threshold'))
    nms_threshold = float(request.args.get('nms_threshold'))
    camera_index = int(request.args.get('camera_index'))
    return Response(detect_objects(conf_threshold, nms_threshold, camera_index), mimetype='multipart/x-mixed-replace; boundary=frame')

if __name__ == '__main__':
    threading.Thread(target=lambda: mqtt_client.publish(sys.argv[6], "OK")).start()
    app.run(host='0.0.0.0', port=sys.argv[2], threaded=True, debug=False)