import os
import requests
import json
import gradio as gr
from PIL import Image
from io import BytesIO


def generate_figure(style, desc):
    url = os.environ["req_url"]

    requests_json = {
        "style": style,
        "desc": desc
    }

    headers = {
        "Content-Type": "application/json",
    }

    response = requests.post(url, json=requests_json, headers=headers, verify=False)

    status = response.status_code
    assert status == 201

    response = json.loads(response.text)

    url_dict = response["data"]["pictures"]
    image_list = []
    for k in url_dict:
        image_list.append(Image.open(BytesIO(requests.get(url_dict[k]).content)))

    return image_list


def read_content(file_path: str) -> str:
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()

    return content


examples_style = ["宫崎骏", "新海诚", "梵高", "赛博朋克", "水彩", "莫奈"]
examples_desc = ["城市夜景", "海滩 美景", "一只猫", "摩天大楼", "鸢尾花", "秋水共长天一色"]

css = """
.gradio-container {background-image: url('file=./background.jpg'); background-size:cover; background-repeat: no-repeat;}

#generate {
    background: linear-gradient(#D8C5EB, #C5E8EB，#90CCF6);
    border: 1px solid #C5E8EB;
    border-radius: 8px;
    color: #26498B
}
"""

# warm up
generate_figure("梵高", "一只猫")

with gr.Blocks(css=css) as demo:
    gr.HTML(read_content("./header.html"))

    gr.Markdown("# MindSpore Wukong-Huahua "
                "\nWukong-Huahua is a diffusion-based model that perfoms text-to-image task in Chinese, "
                "which was developed by the Huawei Noah's Ark Lab in cooperation with "
                "the Distributed & Parallel Software Lab and Ascend Product Develop Unit. "
                "It was trained on Wukong dataset and used MindSpore + Ascend,"
                " a software and hardware solution to implement. Welcome to try Wukong-Huahua by Our Online Platform.")

    with gr.Tab("图片生成 (Figure Generation)"):

        style_input = gr.Textbox(lines=1,
                                 placeholder="输入中文风格描述",
                                 label="Input the style of figure you want to generate. (use Chinese better)",
                                 elem_id="style-input")
        gr.Examples(
            examples=examples_style,
            inputs=style_input,
        )
        with gr.Row():
            gr.Markdown(" *** ")
        desc_input = gr.Textbox(lines=1,
                                placeholder="输入中文图片描述",
                                label="Input a sentence to describe the figure you want to generate. "
                                      "(use Chinese better)")
        gr.Examples(
            examples=examples_desc,
            inputs=desc_input,
        )
        generate_button = gr.Button("Generate", elem_id="generate")
        with gr.Row():
            img_output1 = gr.Image(type="pil")
            img_output2 = gr.Image(type="pil")
            img_output3 = gr.Image(type="pil")
            img_output4 = gr.Image(type="pil")

    with gr.Accordion("Open for More!"):
        gr.Markdown("- If you want to know more about the foundation models of MindSpore, please visit "
                    "[The Foundation Models Platform for Mindspore](https://xihe.mindspore.cn/)")
        gr.Markdown("- If you want to know more about MindSpore-related diffusion models, please visit "
                    "[minddiffusion](https://github.com/mindspore-lab/minddiffusion)")
        gr.Markdown("- Try [Wukong-Huahua model on the Foundation Models Platform for Mindspore]"
                    "(https://xihe.mindspore.cn/modelzoo/wukong)")

    generate_button.click(generate_figure,
                          inputs=[style_input, desc_input],
                          outputs=[img_output1, img_output2, img_output3, img_output4])

demo.queue(concurrency_count=5)
demo.launch(enable_queue=True)