import base64
import datetime
import json
import os

import requests

# 文档
# https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/API

# 参考
# https://www.bilibili.com/read/cv28519209/

# 主机地址
HOST = "http://10.10.21.21:7860"


def submit_post(url: str, data: dict):
    return requests.post(url, data=json.dumps(data))


def save_encoded_image(b64_image: str, output_path: str):
    # 判断当前目录下是否存在 output 文件夹，如果不存在则创建
    if not os.path.exists("output"):
        os.mkdir("output")
    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    output_path = f"{output_path}_{timestamp}" + ".png"
    # 将文件放入当前目录下的 output 文件夹中
    output_path = f"output/{output_path}"
    with open(output_path, "wb") as f:
        f.write(base64.b64decode(b64_image))


def save_json_file(data: dict, output_path: str):
    # 忽略 data 中的 images 字段
    data.pop('images')
    # 将 data 中的 info 字段转为 json 字符串，info 当前数据需要转义
    data['info'] = json.loads(data['info'])

    # 输出 data.info.infotexts
    info_texts = data['info']['infotexts']
    for info_text in info_texts:
        print(info_text)

    # 判断当前目录下是否存在 output 文件夹，如果不存在则创建
    if not os.path.exists("output"):
        os.mkdir("output")
    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    output_path = f"{output_path}_{timestamp}" + ".json"
    # 将文件放入当前目录下的 output 文件夹中
    output_path = f"output/{output_path}"
    with open(output_path, "w") as f:
        json.dump(data, f, indent=4, ensure_ascii=False)


def getQueue():
    # 探测队列中任务数量
    url = HOST + "/queue/status"

    # 发送GET请求
    response = requests.get(url)

    # 检查请求是否成功
    if response.status_code == 200:
        # 获取响应内容
        data = response.json()  # 假设返回的数据是JSON格式
        print(data)
    else:
        print('请求失败，状态码：', response.status_code)


# 文生图
def txt2img():
    txt2img_url = HOST + "/sdapi/v1/txt2img"  # 服务器地址
    prompt = input("请输入提示词：")
    negative_prompt = input("请输入反面提示词：")
    data = {'prompt': prompt, 'negative_prompt': negative_prompt, "cfg_scale": 7, "height": 1024, "width": 1024,
            "sampler_name": "DPM++ 2M Karras", "samples": 1, "steps": 30}

    # 将 data.prompt 中的文本，删除文件名非法字符，已下划线分隔，作为文件名
    output_path = data['prompt'].replace(" ", "_").replace("/", "_").replace("\\", "_").replace(":", "_").replace("\"",
                                                                                                                  "_").replace(
        "<", "_").replace(">", "_").replace("|", "_")
    response = submit_post(txt2img_url, data)
    save_encoded_image(response.json()['images'][0], output_path)
    save_json_file(response.json(), output_path)


if __name__ == '__main__':
    # 查询运行队列情况
    getQueue()
    # 文生图
    txt2img()