import os

from Util import ConfigUtil
from Util.SDUtil import *
import random


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 C25(prompt_data, img, source_code, target_code):
    # 原图
    prompt_data["111"]["inputs"]["image"] = img[0]

    # 风格图片
    style_id = random.randint(1, 3)
    style_id = 1
    with open('../Image/C25/style_' + str(style_id) + ".png", "rb") as f:
        y = upload_file(server_address, f, "", True)

    # 风格(随机)
    prompt_data["113"]["inputs"]["image"] = y

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


def runComfyUI(json_file, model_id, task_type_code, input_image, source_code, target_code):
    # 创建目标目录
    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 = C25(prompt_data, img, source_code, target_code)

    print(prompt_data)
    # 生成
    files = generate_clip(server_address, prompt_data, client_id, output_path, myfilter)
    print(files)


if __name__ == '__main__':
    # 配置文件
    config = ConfigUtil.getConfig()
    server_address = config.get('comfyui', 'server_address')

    # 需要替换的提示词
    source_code = ['']
    # 替换后的提示词
    target_code = ['']

    # 模板编号
    model_id = 25
    # 半身
    # https://www.esheep.com/app/1890?utm_source=app_tab
    # input_images = [['../Image/wife.jpg'], ['../Image/xiao2.jpg'], ['../Image/xiaoqiao.jpg']]
    input_images = [['../Image/xiao2.png']]
    # input_images = [['../Image/wife.jpg']]
    # input_image = ['../ShlImages/shl_xs_1.jpg']
    # input_image = ['../ShlImages/shl_xs_2.jpg']

    # 使用哪个文件
    json_file = '../JSON/' + str(model_id) + '.json'
    for input_image in input_images:
        #  请求
        runComfyUI(json_file=json_file,
                   model_id=model_id,
                   task_type_code='User',
                   input_image=input_image,
                   source_code=source_code,
                   target_code=target_code)