import json

import numpy as np
import cv2
import onnxruntime
import urllib.request
import json
from flask import Flask, request, jsonify
import logging


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


model = "model_fp16.onnx"
session = onnxruntime.InferenceSession(model, None)
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name


def predict_single_image(image_url, session=session):
    req = urllib.request.Request(image_url, headers={"User-Agent": "Mozilla/5.0"})
    with urllib.request.urlopen(req) as url:
        image_contents = url.read()

    arr = np.asarray(bytearray(image_contents), dtype=np.uint8)
    img = cv2.imdecode(arr, -1)

    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    resized_image = cv2.resize(img, (300, 300))
    image_for_prediction = resized_image.transpose(2, 0, 1)
    image_for_prediction = image_for_prediction.astype("float16") / 255.0
    image_for_prediction = image_for_prediction[np.newaxis, :]

    data = json.dumps({"data": image_for_prediction.tolist()})
    data = np.array(json.loads(data)["data"]).astype("float16")

    result = session.run([output_name], {input_name: data})
    prediction = int(np.argmax(np.array(result).squeeze(), axis=0))

    probabilities = np.exp(result) / np.sum(np.exp(result))
    score = np.max(probabilities)
    score = float(score)

    return {
        "prediction": prediction,
        "score": score,
    }


def test_image(image_url):
    return predict_single_image(image_url)


@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!7v%NBox0!47^WL4-1;:":
        return jsonify({"result": "Invalid token"}), 401
    print(image_url)
    result = test_image(image_url)

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

    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)