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


# 获取C26中有哪些样式
def C26_get_sdxl_styles():
    files = ['../JSON/C26/sdxl_styles_sai.json', '../JSON/C26/sdxl_styles_twri.json']
    SDXL_STYLES = []
    for json_file in files:
        with open(json_file, 'r', encoding="utf-8") as fi:
            data = json.load(fi)
            for x in data:
                SDXL_STYLES.append(x['name'])
    return SDXL_STYLES


def C26(prompt_data, img):
    # 原图
    prompt_data["13"]["inputs"]["image"] = img[0]

    ['sai-3d-model', 'sai-analog film', 'sai-anime', 'sai-cinematic', 'sai-comic book', 'sai-craft clay',
     'sai-digital art', 'sai-enhance', 'sai-fantasy art', 'sai-isometric', 'sai-line art', 'sai-lowpoly',
     'sai-neonpunk', 'sai-origami', 'sai-photographic', 'sai-pixel art', 'sai-texture', 'ads-advertising',
     'ads-automotive', 'ads-corporate', 'ads-fashion editorial', 'ads-food photography', 'ads-gourmet food photography',
     'ads-luxury', 'ads-real estate', 'ads-retail', 'artstyle-abstract', 'artstyle-abstract expressionism',
     'artstyle-art deco', 'artstyle-art nouveau', 'artstyle-constructivist', 'artstyle-cubist',
     'artstyle-expressionist', 'artstyle-graffiti', 'artstyle-hyperrealism', 'artstyle-impressionist',
     'artstyle-pointillism', 'artstyle-pop art', 'artstyle-psychedelic', 'artstyle-renaissance', 'artstyle-steampunk',
     'artstyle-surrealist', 'artstyle-typography', 'artstyle-watercolor', 'futuristic-biomechanical',
     'futuristic-biomechanical cyberpunk', 'futuristic-cybernetic', 'futuristic-cybernetic robot',
     'futuristic-cyberpunk cityscape', 'futuristic-futuristic', 'futuristic-retro cyberpunk',
     'futuristic-retro futurism', 'futuristic-sci-fi', 'futuristic-vaporwave', 'game-bubble bobble',
     'game-cyberpunk game', 'game-fighting game', 'game-gta', 'game-mario', 'game-minecraft', 'game-pokemon',
     'game-retro arcade', 'game-retro game', 'game-rpg fantasy game', 'game-strategy game', 'game-streetfighter',
     'game-zelda', 'misc-architectural', 'misc-disco', 'misc-dreamscape', 'misc-dystopian', 'misc-fairy tale',
     'misc-gothic', 'misc-grunge', 'misc-horror', 'misc-kawaii', 'misc-lovecraftian', 'misc-macabre', 'misc-manga',
     'misc-metropolis', 'misc-minimalist', 'misc-monochrome', 'misc-nautical', 'misc-space', 'misc-stained glass',
     'misc-techwear fashion', 'misc-tribal', 'misc-zentangle', 'papercraft-collage', 'papercraft-flat papercut',
     'papercraft-kirigami', 'papercraft-paper mache', 'papercraft-paper quilling', 'papercraft-papercut collage',
     'papercraft-papercut shadow box', 'papercraft-stacked papercut', 'papercraft-thick layered papercut',
     'photo-alien', 'photo-film noir', 'photo-glamour', 'photo-hdr', 'photo-iphone photographic', 'photo-long exposure',
     'photo-neon noir', 'photo-silhouette', 'photo-tilt-shift']

    # 风格图片
    SDXL_STYLES = C26_get_sdxl_styles()

    print(SDXL_STYLES)
    idx = random.randint(0, len(SDXL_STYLES)-1)
    prompt_data["79"]["inputs"]["style"] = SDXL_STYLES[idx]

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

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


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

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

    # 模板编号
    model_id = 26
    # 半身
    # https://www.esheep.com/app/1890?utm_source=app_tab
    input_images = [['../Image/wife.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)