import json
import json
from flask import Flask, request, jsonify
import logging
from ultralytics import YOLO


model = YOLO("best.pt")  # load a pretrained model (recommended for training)


app = Flask(__name__)
log = logging.getLogger("werkzeug")
log.setLevel(logging.ERROR)

class_names = {0: 'apple', 1: 'banana', 2: 'banana'}



def predict_single_image(image_url):
    results = model(image_url)
    processed_resutls = []

    for result in results:
        boxes = result.boxes  # Boxes object for bbox outputs
        for box in boxes:
            confidence = box.conf.item()
            name = class_names[int(box.cls.item())]
            processed_resutls.append([name, confidence])

    return processed_resutls


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

    if not image_url:
        return jsonify({"result": "No image_url"}), 400
    token = request.headers.get("Authorization")
    if token != "8Jw1kj5Woa4SDHtD%hH!7v4--":
        return jsonify({"result": "Invalid token"}), 401
    print(image_url)
    result = {
        "predictions": predict_single_image(image_url),
    }

    if data.get("return_version"):
        result["version"] = "1.0.4"

    return jsonify(result), 200


@app.route("/", methods=["GET"])
def home_get():
    return "Hello World!"


if __name__ == "__main__":
    app.run(debug=True, host="0.0.0.0", port=80)
    # app.run(debug=True, host="0.0.0.0", port=80, threaded=False, processes=64)