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 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'])

        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_prompt(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

@app.route('/gen_latest_pic', methods=['POST'])
def gen_latest_pic():
    if request.is_json:
        data = request.get_json()
        base_image = base64.b64decode(data['base_image'])
        pid = data.get('pid', 'default_user')  # Use provided pid or default

        if not base_image:
            return jsonify({"status": "error", "message": "Missing data"}), 422

        # Get the latest created openpose image from output directory
        import os
        import glob
        
        output_dir = 'output'
        # Search for all openpose output images
        image_patterns = [
            os.path.join(output_dir, 'body_pose_output*.png'),
            os.path.join(output_dir, 'body_pose_output_multi*.png')
        ]
        
        all_images = []
        for pattern in image_patterns:
            all_images.extend(glob.glob(pattern))
        
        if not all_images:
            return jsonify({"status": "error", "message": "No openpose images found in output directory"}), 404
        
        # Get the latest file by modification time
        openpose_image_path = max(all_images, key=os.path.getmtime)

        gen_group_pic_prompt(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, comfyUI_address)
        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_prompt(openpose_image_path, base_image, pid, comfyUI_address):
    with open("group_pic.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, comfyUI_address)
        base_image_name = opg.upload_image_circular_queue(base_image, 30, pid, comfyUI_address)

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

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