from flask import Flask, request, jsonify
from api import mj_api
from api.error import unauthorized_error
from api.error import system_error

app = Flask(__name__)


@app.route('/')
def hello():
    return "Hello World"


# 查询历史所有任务列表
@app.route('/thomas_jobs', methods=['POST'])
def thomas_jobs():
    request_data = request.get_json()

    if request_data.get('cookie') is None:
        return jsonify({"error": "Cookie is missing!"}), 400

    if request_data.get('user_id') is None:
        return jsonify({"error": "page_size is required!"}), 400

    response_json_str = mj_api.thomas_jobs(cookie=request_data.get('cookie'), user_id=request_data.get('user_id'),
                                           page_size=request_data.get('page_size'))
    return response_json_str


# 查询指定任务状态
@app.route('/query_job_status', methods=['POST'])
def query_job_status():
    request_data = request.get_json()

    if request_data.get('cookie') is None:
        return jsonify({"error": "Cookie is missing!"}), 400

    if request_data.get('job_id') is None:
        return jsonify({"error": "job_id is required!"}), 400

    response_json_str = mj_api.query_job_status(cookie=request_data.get('cookie'), job_id=request_data.get('job_id'))
    return response_json_str


# 查询任务队列
@app.route('/query_job_queue', methods=['POST'])
def query_job_queue():
    request_data = request.get_json()

    if request_data.get('cookie') is None:
        return jsonify({"error": "Cookie is missing!"}), 400

    if request_data.get('user_id') is None:
        return jsonify({"error": "user_id is required!"}), 400

    response_json_str = mj_api.query_job_queue(cookie=request_data.get('cookie'), user_id=request_data.get('user_id'))
    return response_json_str


# 提交任务
@app.route('/submit_job', methods=['POST'])
def submit_job():
    request_data = request.get_json()

    if request_data.get('cookie') is None:
        return jsonify({"error": "Cookie is missing!"}), 400

    if request_data.get('prompt') is None:
        return jsonify({"error": "prompt is required!"}), 400

    if request_data.get('user_id') is None:
        return jsonify({"error": "user_id is required!"}), 400

    response_json_str = mj_api.submit_job(cookie=request_data.get('cookie'), prompt=request_data.get('prompt'),
                                          user_id=request_data.get('user_id'))
    return response_json_str


@app.errorhandler(unauthorized_error)
def handle_unauthorized_error(error):
    response = jsonify({"error": str(error)})
    response.status_code = 401
    return response


@app.errorhandler(system_error)
def handle_system_error(error):
    response = jsonify({"error": str(error)})
    response.status_code = 500
    return response


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