from flask import Flask, Response, request, render_template, jsonify
import paho.mqtt.client as mqtt
from ultralytics import YOLO
import cv2
import numpy as np

app = Flask(__name__)

# Load a model
model = YOLO("yolo11n.pt")  # pretrained YOLO11n model

kobuki_message = ""
latest_image = None
yolo_results = []

def on_message(client, userdata, message):
    global kobuki_message, latest_image, yolo_results
    if message.topic == "kobuki/data":
        kobuki_message = str(message.payload.decode("utf-8"))
    elif message.topic == "kobuki/cam":
        latest_image = np.frombuffer(message.payload, np.uint8)
        latest_image = cv2.imdecode(latest_image, cv2.IMREAD_COLOR)
        # Process the image with YOLO
        results = model(latest_image)
        yolo_results = []
        for result in results:
            for box in result.boxes:
                yolo_results.append({
                    "class": box.cls,
                    "confidence": box.conf,
                    "bbox": box.xyxy.tolist()
                })
                # Draw bounding box on the image
                x1, y1, x2, y2 = map(int, box.xyxy)
                cv2.rectangle(latest_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(latest_image, f"{box.cls} {box.conf:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

# Create an MQTT client instance
mqtt_client = mqtt.Client()
mqtt_client.username_pw_set("server", "serverwachtwoordofzo")
mqtt_client.connect("localhost", 1884, 60)
mqtt_client.loop_start()
mqtt_client.subscribe("kobuki/data")
mqtt_client.subscribe("kobuki/cam")

mqtt_client.on_message = on_message # this lines needs to be under the function definition otherwise it cant find which function it needs to use

@app.route('/')
def index():
    return render_template('index.html')


@app.route('/control', methods=["GET","POST"])
def control():
    if request.authorization and request.authorization.username == 'ishak' and request.authorization.password == 'kobuki':
        return render_template('control.html')
    else:
        return ('Unauthorized', 401, {'WWW-Authenticate': 'Basic realm="Login Required"'})

@app.route('/move', methods=['POST'])
def move():
    data = request.get_json()
    direction = data.get("direction")

    # Verstuur de richting via MQTT
    if direction:
        mqtt_client.publish("home/commands", direction)  # Het topic kan aangepast worden

    return jsonify({"status": "success", "direction": direction})


@app.route('/data', methods=['GET'])
def data():
    return kobuki_message

@app.route('/image')
def image():
    global latest_image
    if latest_image is not None:
        _, buffer = cv2.imencode('.jpg', latest_image)
        return Response(buffer.tobytes(), mimetype='image/jpeg')
    else:
        return "No image available", 404

@app.route('/yolo_results', methods=['GET'])
def yolo_results_endpoint():
    global yolo_results
    return jsonify(yolo_results)

if __name__ == '__main__':
    app.run(debug=True, port=5000)