# 文档： https://www.youtube.com/watch?v=va8Jkc7o9d4

# ComfyUI需要打开开发者模式，然后才能导出API JSON
# pip install websocket
# pip install websocket-client
# 如果不安装websocket-client会报错 wekscocket构造参数不正确
# 参考：https://blog.csdn.net/it_is_arlon/article/details/130359239

import os
from Util.SDUtil import *
from Util import ConfigUtil


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, input_img, source_code, target_code):
    # 配置文件
    config = ConfigUtil.getConfig()
    server_address = config.get('comfyui', 'server_address')

    # 创建目标目录
    output_path = "../Out"
    if not os.path.exists(output_path):
        os.makedirs(output_path)

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

    # 上传图片
    with open(input_img, "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)

    # 调用不同的实现函数
    number = json_file[json_file.rfind('/') + 1:-5]  # 从最后一个'/'之后开始，到倒数第五个字符结束

    #  分支判断，如果是16，走C16
    if number == '16':
        prompt_data, myfilter = C16(prompt_data, img, source_code, target_code)

    # 添加随机值
    prompt_data['guid'] = str(uuid.uuid4())

    # 生成
    files = generate_clip(server_address, prompt_data, client_id, output_path, myfilter)
    # 上传到云存储，回写数据接口
    for file in files:
        print('./Out/' + file + '.png')


if __name__ == "__main__":
    # 上传的图片
    input_img = "./Image/xiaoqiao.jpg"
    # 需要替换的提示词
    source_code = ['1 girl']
    # 替换后的提示词
    # target_code = ['1 boy']
    target_code = ['1 girl']
    # 使用哪个文件
    json_file = '../JSON/16.json'
    #  请求
    runComfyUI(json_file=json_file, input_img=input_img, source_code=source_code, target_code=target_code)