import os

from Util import ConfigUtil
from Util.SDUtil import *


def C19(prompt_data, img, source_code, target_code):
    #
    prompt_data["12"]["inputs"]["image"] = img[0]
    prompt_data["27"]["inputs"]["image"] = img[1]

    # 替换关键词
    # for i in range(len(source_code)):
    #     s = source_code[i]
    #     t = target_code[i]
    #     prompt_data["331"]["inputs"]["text"] = prompt_data["331"]["inputs"]["text"].replace(s, t)

    # 过滤器,哪些返回节点中获取到的图片是有用的
    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 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 = C19(prompt_data, img, source_code, target_code)

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

if __name__ == '__main__':
    # 需要替换的提示词
    source_code = ['']
    # 替换后的提示词
    target_code = ['']
    # 使用哪个文件
    json_file = '../JSON/19.json'
    input_image = ['./Image/huanghai.png', './Image/littlehb.png']
    #  请求
    runComfyUI(json_file=json_file,
               model_id=19,
               task_type_code='User',
               input_image=input_image,
               source_code=source_code,
               target_code=target_code)