#!/usr/bin/env python
# encoding: utf-8
"""
@author: simin. Gong
@license: (C) Copyright 2013-2023 Spacemit Corporation Limited.
@contact: simin.gong@spacemit.com
@file: app.py.py
@date: 2024/11/6 16:58 
@desc:
"""
import uuid
from flask import Flask, send_from_directory, request, send_file, Response
from flask_cors import CORS
from werkzeug.utils import secure_filename
import shutil
import  os
from urllib import request as url_request
from urllib.error import URLError
import subprocess
from  ai_tool import static
from ai_tool import aiLib
import json
import io
import sys
lib_path = os.path.dirname(aiLib.__file__)

sys.path.append(lib_path)

static_path = os.path.dirname(static.__file__)


home_path = os.path.expanduser("~")

save_path = os.path.join(home_path, '.ai-tmp')

try:
    import Classifi as cls
    import yolov8 as yolo
    import ObjectSuperpoint as Superpoint
    import PPOCRRec as ocr
    import ObjectMask as mask
    import cv2

    onnx_path = os.path.join(lib_path, 'mobilenet_v2.q.onnx')

    yolo_path = os.path.join(lib_path, 'yolo11n.onnx')

    super_point_path = os.path.join(lib_path, 'superpoint_512_480x640.onnx')

    ocr_path = os.path.join(lib_path, 'ppocrv4_rec_new.onnx')

    mask_path = os.path.join(lib_path, 'model_627_new.onnx')

    print('加载class 模型')
    cls_infer = cls.create_infer(onnx_path, 'onnx')

    print('加载yolo 模型')
    yolo_infer = yolo.create_infer(yolo_path, 'onnx')

    print('加载superpoint 模型')
    super_point_infer = Superpoint.create_infer(super_point_path)

    print('加载ocr模型')
    ocr_infer = ocr.create_infer(ocr_path)

    print('加载seaformer模型')
    mask_infer = mask.create_infer(mask_path)
except Exception:
    print("\033[31m\***************** not riscv env! ******************* /\033[0m")

if os.path.exists(save_path):
    shutil.rmtree(save_path)
    os.makedirs(save_path)

app = Flask(__name__)
CORS(app)
app.config['SECRET_KEY'] = 'secret!'
app.config['MAX_CONTENT_LENGTH'] = 100 * 1024 * 1024
app.static_folder = static_path

def get_default_images_path(image_name: str):
    if not image_name.startswith('default_'):
        return None
    for root, dirs, files in os.walk(os.path.join(app.static_folder, 'img')):
        for file_name in files:
            if file_name.split('.')[0].startswith(image_name.split('.')[0]):
                return os.path.join(root, file_name)
    return None

def get_image_path(image_name: str):
    default_path = get_default_images_path(image_name)
    if default_path:
        return default_path
    return os.path.join(save_path, image_name)


def run_command(command):
    try:
        result = subprocess.run(command, shell=True, capture_output=True, text=True, check=True)
        return result.stdout, result.stderr
    except subprocess.CalledProcessError as e:
        print(f"命令执行失败: {e}")
        return None, e.stderr


def run_yolo(image_name: str, backend: str = 'onnx'):
    image = cv2.imread(get_image_path(image_name))
    objects = yolo_infer.commit(image).get()
    data = None
    for Box in objects:
        data = json.loads(Box.json_info.text)
        print("json_info:", Box.json_info.text)
    if data is None:
        return Response('Error Cv infer', status=500)
    return data

def run_mobilenet(image_name: str, backend: str = 'onnx'):
    image = cv2.imread(get_image_path(image_name))
    objects = cls_infer.commit(image).get()
    data = None
    for text in objects:
        data = json.loads(text.text)
    if data is None:
        return Response('Error Cv infer', status=500)
    return data

def run_super_point(image_name: str, backend: str = 'onnx'):
    image = cv2.imread(get_image_path(image_name), 0)
    objects = super_point_infer.commit(image).get()
    data = None
    for obj in objects:
        data = json.loads(obj.json)
    if data is None:
        return Response('Error Cv infer', status=500)
    return data

def run_ocr(image_name: str, backend: str = 'onnx'):
    image = cv2.imread(get_image_path(image_name))
    objects = ocr_infer.commit(image).get()
    print("len:", len(objects))
    data = None
    for text_line in objects:
        data = json.loads(text_line.json)
    print("json_info:", data)
    if data is None:
        return Response('Error Cv infer', status=500)
    return data


def run_mask(image_name: str, backend: str = 'onnx'):
    image = cv2.imread(get_image_path(image_name))
    objects = mask_infer.commit(image).get()
    image_stream = None
    for idx, obj in enumerate(objects):
        data = obj.data
        _, buffer = cv2.imencode('.png', data)
        image_stream = io.BytesIO(buffer)
    if image_stream is None:
        return Response('Error Cv infer', status=500)
    return send_file(image_stream, mimetype='image/png')


run_ai_handle = {
    'yolo': run_yolo,
    'mobilenet': run_mobilenet,
    'superpoint': run_super_point,
    'ocr': run_ocr,
    'seaformer': run_mask
}

@app.route('/<path:path>')
def send_vue_app(path):
    # 直接将请求转发到Vue打包后的文件
    response = send_from_directory(app.static_folder, path)
    return response


@app.route('/api/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return 'No file part'
    file = request.files['file']
    if file.filename == '':
        return 'No selected file'
    fname = secure_filename(file.filename)
    ext = fname.rsplit('.')[-1]
    fileName = str(uuid.uuid4()) + "." + ext
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    file.save(os.path.join(save_path, fileName))
    return fileName

@app.route('/api/post-run', methods=['POST'])
def post_run():
    data = request.get_json()
    image_name = data['image']
    backend = data['backend']
    alg = data['alg']
    print(f"deal image: {image_name} ")
    if alg not in run_ai_handle:
        return Response(f'Not Support Alg {alg}', status=500)
    return run_ai_handle[alg](image_name, backend)

 # return  [
 #              { 'cls' :'ambulance', 'value': 0.5479885, 'pos': [1,2]},
 #              { 'cls':'streetcar, tram, tramcar, trolley, trolley car', 'value' : 0.1859075},
 #              { 'cls': 'police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria', 'value': 0.1277304},
 #              { 'cls': 'minibus', 'value': 0.0327211},
 #              { 'cls':'trolleybus, trolley coach, trackless trolley', 'value': 0.0268708}
 #            ]


@app.route('/api/get-support', methods=['GET'])
def get_support_ai_fun():
    return {'backend': ['onnx'], 'alg': ['yolo', 'mobilenet', 'superpoint', 'seaformer', 'ocr']}

@app.route('/proxy', methods=['GET'])
def proxy():
    image_url = request.args.get('url')
    try:
        with url_request.urlopen(image_url) as response:
            content_type = response.getheader('Content-Type')
            content = response.read()
            return Response(content, mimetype=content_type)
    except URLError as e:
        return Response(f'Error fetching the image: {e.reason}', status=500)

@app.route('/')
def index():
    response = send_from_directory(app.static_folder, 'index.html')
    return response

def run():
    app.run(host='0.0.0.0', port=5000, threaded=True)

if __name__ == '__main__':
    run()
