from flask import Flask, request,send_file
from tool import runtool
import yaml
import os
import json

import threading
app = Flask(__name__)

with open('config.yml', 'r') as file:
    config = yaml.load(file, Loader=yaml.FullLoader)

# env = os.getenv('FLASK_ENV', 'development')
app.config.update(config['default'])
# if env in config:
# app.config.update(config[env])


base_dir = os.path.dirname(os.path.abspath(__file__))
task_output_dir = os.path.join(base_dir, "output")

def get_task_list():
    '''
    获取任务列表,
    任务是task_output_dir下的第一层文件夹
    
    同时检查是否超过了最大任务数，如果超过了，删除最早的任务直到任务数小于最大任务数
    '''
    
    task_list = []
    task_list = os.listdir(task_output_dir)
    # 文件夹的名称是按照task_<task_num>_...来的, 按照task_num排序
    
    task_list_dict = {}
    
    for task_dir in task_list:
        task_num = int(task_dir.split('_')[1])
        task_list_dict[task_num] = task_dir
    
    task_num_list = list(task_list_dict.keys())
    task_num_list.sort()
    
    return_list = []
    for task_num in task_num_list:
        return_list.append(task_list_dict[task_num])
    return return_list

def get_max_task_id():
    
    task_list = []
    task_list = os.listdir(task_output_dir)
    if len(task_list) == 0:
        return 0
    # 文件夹的名称是按照task_<task_id>_...来的，取出最大的task_id
    task_id_to_task_dir = {}
    task_id_list = []
    for task_dir in task_list:
        task_id = int(task_dir.split('_')[1])
        task_id_list.append(task_id)
        task_id_to_task_dir[task_id] = task_dir
    
    max_task_id =  max(task_id_list)
    
    # 检查是否超过了最大任务数，删除task_num最小的任务直到任务数小于最大任务数
    print(len(task_list))
    print(app.config['max_task_num'])
    if len(task_list) > app.config['max_task_num']:
        print("超过最大任务数，删除最早的任务直到任务数小于最大任务数")
        task_id_list.sort()
        over_num = len(task_list) - app.config['max_task_num']
        for i in range(over_num):
            task_id = task_id_list[i]
            task_dir = task_id_to_task_dir[task_id]
            task_dir_path = os.path.join(task_output_dir, task_dir)
            print("删除任务:", task_dir_path)
            os.system("rm -rf " + task_dir_path)
        
    
    return max_task_id
    
def get_result(task_id:str):
    result_dir = os.path.join(task_output_dir, task_id)
    # 查看result_dir的下一级文件夹
    
    result_model_name = os.listdir(result_dir)[0]
    
    result_model_dir = os.path.join(result_dir,result_model_name)
    print(result_model_dir)
    
    mxnet_json = os.path.join(result_model_dir, "mxnet.json")
    with open(
        mxnet_json, "r"
    ) as file1:  # tensorflow.json
        data1 = json.load(file1)
    
    tensorflow_json = os.path.join(result_model_dir, "tensorflow.json")
    with open(tensorflow_json, "r") as file2:
        data2 = json.load(file2)
    
    img_root_relative_path = "/model_accuracy_api/image/" + task_id + '/' + result_model_name + '/'
    
    
    combined_data = {
        "mxnet": data1,
        "tensorflow": data2,
        "img_path":{
            "mxnet_train":img_root_relative_path + 'mxnet_train.jpg',
            "tensorflow_train": img_root_relative_path + 'tensorflow_train.jpg',
            "accuracy": img_root_relative_path + 'accuracy.jpg',
            "losses": img_root_relative_path + 'losses.jpg',
            "memory": img_root_relative_path + 'memory.jpg'
        }
    }
    return combined_data
    
    
@app.route('/')
def hello_world():  # put application's code here
    return 'Hello World!'

@app.route('/models', methods=['GET'])
def getModels():
    models = ['lexnet-cifar10', 'lenet5-fashion-mnist', 'fashion2', 'svhn', 'lenet5-mnist',
            'alexnet-cifar10', 'mobilenet.1.00.224-imagenet', 'vgg16-imagenet']
    return models

import time

@app.route('/run', methods=['POST'])
def run():
    data = request.get_json()
    exp = data['exp']
    mutate_num = data['mutate_num']
    
    #创建task_id: task_ task_num exp mutate_num
    # task_id = 'task_' + time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime()) + '_' + str(len(get_task_list())+1)
    task_id = 'task_' + str(get_max_task_id()+1) + '_' + exp + '_' + str(mutate_num)
    # 开启一个线程执行，直接返回请求
    # data = runtool(exp, mutate_num,task_id)
    t = threading.Thread(target=runtool, args=(exp, mutate_num,task_id))
    t.start()
    
    return {'task_id': task_id}

@app.route('/get_task_list', methods=['GET'])
def get_task_list_api():
    return {'task_list': get_task_list()}

@app.route('/get_task_result', methods=['POST'])
def get_task_result():
    data = request.get_json()
    task_id = data['task_id']
    
    try:
        res = get_result(task_id)
        return res
    except Exception as e:
        return {'error:':'错误的task_id或任务未完成'}
    
@app.route('/image/<path:filename>')
def serve_image(filename):
    # 指定图片所在的根目录
    
    # 构建完整的图片路径
    image_path = os.path.join(task_output_dir, filename)
    
    # 使用send_file()函数返回图片
    return send_file(image_path, mimetype='image/jpeg')

if __name__ == '__main__':
    # app.run()
    from waitress import serve
    serve(app, host="0.0.0.0", port=5000)