from PIL import Image
from flask import Flask, json, request, abort
import argparse

import base64
import re
from io import BytesIO
import uuid
import os

from yolo_e_bike import YOLO

api = Flask(__name__)
yolo = None


@api.route('/yolo/api/v1/detect', methods=['POST'])
def detect():
    if not request.json or not 'image' in request.json:
        abort(400)
    img = request.json.get("image")
    threshold = request.json.get("threshold", 0.5)
    show_pic = request.json.get("show", False)
    base64_data = re.sub('^data:image/.+;base64,', '', img)
    byte_data = base64.b64decode(base64_data)
    image_data = BytesIO(byte_data)
    img = Image.open(image_data)
    results = yolo.detect_json(img, threshold, show_pic=show_pic)
    log_id = uuid.uuid3(uuid.NAMESPACE_DNS, 'yolov3')
    detect_result = {
        "log_id": log_id,
        "results": results
    }
    return json.dumps(detect_result)


@api.route('/yolo/api/v1/detect_mock', methods=['POST'])
def detect_mock():
    dict = {}
    results = []
    dict['log_id'] = uuid.uuid3(uuid.NAMESPACE_DNS, 'yolov3')
    result = {}
    location = {}
    location['height'] = 256
    location['width'] = 350
    location['top'] = 260
    location['left'] = 300
    result['name'] = "ddzxc"
    result['score'] = 0.9
    result['location'] = location
    results.append(result)
    dict['results'] = results
    return json.dumps(dict)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS)
    parser.add_argument(
        '--model_path', type=str,
        help='path to model weight file, default ' + YOLO.get_defaults("model_path")
    )

    parser.add_argument(
        '--anchors_path', type=str,
        help='path to anchor definitions, default ' + YOLO.get_defaults("anchors_path")
    )

    parser.add_argument(
        '--classes_path', type=str,
        help='path to class definitions, default ' + YOLO.get_defaults("classes_path")
    )

    parser.add_argument(
        '--gpu_num', type=int,
        help='Number of GPU to use, default ' + str(YOLO.get_defaults("gpu_num"))
    )

    parser.add_argument(
        '--cpu', default=False, action="store_true",
        help='Cpu mode'
    )

    FLAGS = parser.parse_args()
    if FLAGS.cpu:
        os.environ["CUDA_VISIBLE_DEVICES"] = " "
    yolo = YOLO(**vars(FLAGS))
    api.run(host="0.0.0.0", port=8085)