import base64
import hashlib
import json
import random
import uuid
from flask import Flask, request, jsonify
import sys
import openpose_gen as opg

sys.path.append('./')

app = Flask(__name__)

@app.route('/gen_image', methods=['POST'])
def gen_image():
    if request.is_json:
        data = request.get_json()
        coordinates = data['coordinates']
        canvas_size = data['canvas_size']
        pid = data['pid']

        if not coordinates or not canvas_size:
            return jsonify({"status": "error", "message": "Missing data"}), 422
        
        openpose_image_path = opg.save_bodypose(canvas_size[0], canvas_size[1], coordinates, pid)
        # gen_fencer_prompt(openpose_image_path, pid, opg.server_address)

        return jsonify({"status": "success", "message": "Data received"}), 201
    else:
        return jsonify({"status": "error", "message": "Request must be JSON"}), 415


@app.route('/gen_group_pic', methods=['POST'])
def init_gen_group_pic():
    if request.is_json:
        data = request.get_json()
        coordinates_list = data['coordinates_list']
        canvas_size = data['canvas_size']
        pid = data['pid']
        base_image = base64.b64decode(data['base_image'])

        # save base image to ./test.png
        with open("test.png", "wb") as f:
            f.write(base_image)

        if not coordinates_list or not canvas_size or not base_image or not pid:
            return jsonify({"status": "error", "message": "Missing data"}), 422
        
        for i in range(len(coordinates_list)):
            coordinates_list[i] = coordinates_list[i]['coordinates']
        
        openpose_image_path = opg.save_bodypose_mulit(canvas_size[0], canvas_size[1], coordinates_list, pid)
        # gen_group_pic(openpose_image_path, base_image, pid, opg.server_address)

        return jsonify({"status": "success", "message": "Data received"}), 201
    else:
        return jsonify({"status": "error", "message": "Request must be JSON"}), 415



def gen_fencer_prompt(openpose_image_path, pid, comfyUI_address):
    with open("fencerAPI.json", "r") as f:
        prompt_json = f.read()
        prompt = json.loads(prompt_json)

        openpose_image_name = opg.upload_image_circular_queue(openpose_image_path, 20, pid)
        opg.upload_image("ref_black.png", "ref_black.png")

        prompt["3"]["inputs"]["seed"] = random.randint(0, 10000000000)
        prompt["29"]["inputs"]['image'] = "ref_black.png"
        prompt["17"]["inputs"]['image'] = openpose_image_name

        opg.queue_prompt(prompt, comfyUI_address)

def gen_group_pic(openpose_image_path, base_image, pid, comfyUI_address):
    with open("groupAPI.json", "r") as f:
        prompt_json = f.read()
        prompt = json.loads(prompt_json)

        openpose_image_name = opg.upload_image_circular_queue(openpose_image_path, 30, pid)
        base_image_name = opg.upload_image_circular_queue(base_image, 30, pid)

        prompt["3"]["inputs"]["seed"] = random.randint(0, 10000000000)
        prompt["17"]["inputs"]['image'] = openpose_image_name
        prompt["17"]["inputs"]['image'] = base_image_name

        opg.queue_prompt(prompt, comfyUI_address)
    
if __name__ == '__main__':
    app.run(debug=True)