#!/usr/bin/env python3

import os
import uuid
import json
import time
import cv2
import queue
from pathlib import Path as P
from PIL import Image
from flask import Flask, request, send_from_directory
from flask_cors import CORS
from threading import Thread

app = Flask(__name__)
CORS(app)

QUEUE_STATE_IN_QUEUE = 'IN_QUEUE'
QUEUE_STATE_PROCESSING = 'PROCESSING'
QUEUE_STATE_READY = 'READY'

q = queue.Queue(maxsize=1000)
state = {}
root = P("./q").absolute()
root.mkdir(parents=True, exist_ok=True)


def process():
    print("importing ultralytics")
    from ultralytics import YOLO
    print("creating model")
    model = YOLO(os.environ.get("MODEL_PATH", "/best.pt"))
    print("created model")
    last = 0
    while True:
        # wait x seconds before processing another task
        if time.time() < last + 2:
            print("aa")
            time.sleep(0.5)
            continue

        print("BB")

        print("waiting for/getting job")
        job = q.get()
        filepath = root.joinpath(job)
        print("got job", flush=True)
        state[job] = {
            'state': QUEUE_STATE_PROCESSING,
        }
        print(f"processing job: {job=}")
        img = cv2.imread(str(filepath))
        results = model.predict(img)

        for r in results:
            im_array = r.plot()
            im = Image.fromarray(im_array[..., ::-1])
            im.save(str(filepath.with_suffix(".png")))
            state[job] = {
                'state': QUEUE_STATE_READY,
            }

        filepath.unlink()

        last = time.time()


p = Thread(target=process)
p.start()


@app.route("/submit", methods=['POST'])
def render():
    id = str(uuid.uuid4())
    request.files['image'].save(root.joinpath(id))
    q.put(id)
    state[id] = {'state': QUEUE_STATE_IN_QUEUE}
    return json.dumps({'id': id})


@app.route("/state/<id>")
def get_state(id):
    return json.dumps(state[id])


@app.route("/result/<i>")
def get_result(i):
    print(root)
    print(i + '.png')
    return send_from_directory(root, i + '.png')


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