import os
import random

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


def C16(prompt_data, img, source_code, target_code):
    # 图生图
    prompt_data["22"]["inputs"]["image"] = img

    # 替换关键词
    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):
    # 配置文件
    config = ConfigUtil.getConfig()
    server_address = config.get('comfyui', 'server_address')

    # 创建目标目录
    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())

    # 上传图片
    with open(input_image, "rb") as f:
        img = upload_file(server_address, f, "", True)

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

    #  分支判断
    prompt_data, myfilter = C16(prompt_data, img, source_code, target_code)

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


if __name__ == '__main__':
    # 打开配置文件
    config = ConfigUtil.getConfig()
    # 不使用代理
    set_http_proxy("")

    directory = r'D:\backup\三级读书营照片'
    for root, dirs, files in os.walk(directory):
        for name in files:
            # print(os.path.join(root, name))
            fi = os.path.join(root, name)
            if '男_' in fi:
                replace_content = '1 boy'
            else:
                replace_content = '1 girl'

            runComfyUI(json_file='../JSON/16/16.json',
                       model_id=16,
                       task_type_code='User',
                       input_image=fi,
                       source_code=['1 girl'],
                       target_code=[replace_content]
                       )