import os
import random

from Util import ConfigUtil
from Util.SDUtil import *
from Util.CommonUtil import *


def getMyFilter(prompt_data):
    myfilter = []
    # 遍历prompt_data中的所有节点
    for key, value in prompt_data.items():
        if 'inputs' in value:
            if 'filename_prefix' in value['inputs']:
                if value['inputs']['filename_prefix'] == 'ComfyUI':
                    myfilter.append(key)
    return myfilter
def C17(prompt_data, img, model_id,  prompt):
    # 上传LOGO
    input_image_logo = '../Image/C17/logo.png'
    with open(input_image_logo, "rb") as f:
        imgLogo = upload_file(server_address, f, "", True)

    # 上传姿势
    # random_number = random.randint(1, 9)  # 9张姿势图
    # input_image_pose = './Image/C17/' + str(random_number) + '.png'
    input_image_pose = '../Image/C17/1.jpg'
    with open(input_image_pose, "rb") as f:
        imgPose = upload_file(server_address, f, "", True)

    # 图生图
    prompt_data["331"]["inputs"]["image"] = imgPose
    #prompt_data["614"]["inputs"]["image"] = imgLogo
    prompt_data["432"]["inputs"]["image"] = img[0]

    # 指定数字长度，例如16位
    number_length = 16
    random_number = generate_random_number(number_length)
    prompt_data["354"]["inputs"]["seed"] = random_number
    random_number = generate_random_number(number_length)
    prompt_data["425"]["inputs"]["seed"] = random_number

    # 替换关键词
    prompt_data["331"]["inputs"]["text"] = prompt

    # 过滤器,哪些返回节点中获取到的图片是有用的
    myfilter = getMyFilter(prompt_data)

    return prompt_data, myfilter


def generate_random_number(length):
    return ''.join(random.choices('0123456789', k=length))


def runComfyUI(json_file, model_id, task_type_code, input_image,  prompt):
    # 创建目标目录
    output_path = "Out/Images/" + task_type_code + "/" + str(model_id) + "/"
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    # 生成一个唯一的客户端ID
    client_id = str(uuid.uuid4())

    # 上传图片
    img = []
    for x in input_image:
        with open(x, "rb") as f:
            y = upload_file(server_address, f, "", True)
            img.append(y)

    with open(json_file, 'r', encoding="utf-8") as fi:
        prompt_data = json.load(fi)

    #  分支判断
    prompt_data, myfilter = C17(prompt_data, img, model_id,  prompt)
    # 生成
    generate_clip(server_address, prompt_data, client_id, output_path, myfilter)

if __name__ == '__main__':
    # 配置文件
    config = ConfigUtil.getConfig()
    # 处理机编号
    machine_id = config['system']['machine_id']
    # 生图服务地址
    txt2img_url = config['webui']['txt2img_url']
    # WEB服务器地址
    web_url = config['webui']['web_url']
    # 配置文件
    server_address = config.get('comfyui', 'server_address')

    # 第20个少惠林学员
    for k in range(49, 50):
        prompt='(Cute cartoon character sitting on chair),(looking at viewer：2),(feet on floor：2), (wooden floor: 2), wearing simple beautiful clothes, portrait photography, (stage lighting from front:1.2), large aperture, empty background, empty studio, blue clean background :2, blue curtain: 2, blue empty background: 2, clean floor, blue floor: 2, solo, blue background: 2, happy smile: 2, solo'
        runComfyUI(json_file='../JSON/17.json',
                   model_id=17,
                   task_type_code='User',
                   input_image=['../ShlImages/shl_xs_' + str(k) + '.jpg'],
                   prompt=prompt)