import os
import random

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


# 获取Model对象
def getModel(model_id):
    file_name = r'../JSON/' + str(model_id) + '.json'
    # 打开JSON文件
    if not os.path.exists(file_name):
        printf("文件：" + file_name + "不存在，跳过！")
        return
    with open(file_name, 'r', encoding='utf-8') as file:
        data = json.load(file)
    return data


# 整体的生成任务
def runWebUI(model_id, task_type_code, input_image, total_tasks, prompt):
    global data
    data = getModel(model_id)  # data是全局变量
    # 生成多少张
    data['batch_size'] = total_tasks
    # 图片
    args = data['alwayson_scripts']['ControlNet']['args']
    pose_image = data['pose_image']
    args[0]['input_image'] = encode_image(input_image)
    # 如果有特殊的姿态要求,那就用模板指定的姿态图
    if len(pose_image) > 0:
        if len(args) > 1:
            args[1]['input_image'] = encode_image(pose_image)
    else:  # 如果没有特殊的姿态要求，那么就统一用一张图片
        if len(args) > 1:
            args[1]['input_image'] = encode_image(input_image[0])

    # 提示词
    data['prompt'] = prompt

    # 开始生成图片
    printf("开始生成图片...")
    # 设置目标文件夹
    target_folder = "Out"

    # 创建以当前日期命名的目录
    folder = os.path.join(target_folder, 'Images/' + task_type_code + "/" + str(model_id) + '/')
    if not os.path.exists(folder):
        os.makedirs(folder)
    # 调用生图服务
    response = submit_post(txt2img_url, data)
    # 保存为图片
    for i in range(0, total_tasks):
        # 创建文件名
        png_file = str(uuid.uuid4()) + '.png'
        # 创建文件的完整路径
        file_path = os.path.join(folder, png_file)
        save_encoded_image(response.json()['images'][i], file_path)
    printf("本轮图片生成成功, 共完成" + str(total_tasks) + "张图片")


if __name__ == '__main__':
    # 打开配置文件
    config = ConfigUtil.getConfig()
    # 配置文件
    config = ConfigUtil.getConfig()
    # 处理机编号
    machine_id = config['system']['machine_id']
    # 生图服务地址
    txt2img_url = config['webui']['txt2img_url']
    # WEB服务器地址
    web_url = config['webui']['web_url']

    # 不使用代理
    set_http_proxy("")

    prompt = '(masterpiece, best quality:1.2),1boy,blue eyes,bodysuit,breasts,building,city,Smooth metal surface,Close fitting mecha,Part of the complex mecha structure,Streamlined mecha,lips,looking at viewer,mecha,medium breasts,realistic,science fiction,short hair,solo,<lora:cncnj95e878c73b9ogcg:0.8>,<lora:ip-adapter-faceid-plusv2_sdxl_lora:0.5>,'
    # 第20个少惠林学员
    for k in range(49, 50):
        total_count = 2
        runWebUI(15, 'User', '../ShlImages/shl_xs_' + str(k) + '.jpg', total_count,
                 prompt)  # 原来的input_image是单张图片，现在修改为支持多张图片