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

from utils import get_token


def generate_figure(style, desc):
    url = "https://a2f051d4cabf45f885d7b0108edc9b9c.infer.ovaijisuan.com/" \
          "v1/infers/975eedfd-6e15-4571-8ca9-b945da0da24b/wukong_hf"

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

    headers = {
        "Content-Type": "application/json",
        "X-Auth-Token": token
    }

    response = requests.post(url, json=requests_json, headers=headers, verify=False)
    response = json.loads(response.text)
    status = response["status"]
    assert status == 200
    url_list = response["output_image_url"]

    image1 = Image.open(BytesIO(requests.get(url_list[0]).content))
    image2 = Image.open(BytesIO(requests.get(url_list[1]).content))
    image3 = Image.open(BytesIO(requests.get(url_list[2]).content))
    image4 = Image.open(BytesIO(requests.get(url_list[3]).content))

    return image1, image2, image3, image4


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

    return content


token = get_token()


examples = [
    ["宫崎骏", "城市夜景"],
    ["新海诚", "海滩 美景"],
    ["赛博朋克", "一只猫"],

]

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: #407BEA
}

"""

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

    gr.Markdown("# MindSpore Wukong•Huahua "
                " \n With the help of Wukong dataset, the largest Chinese open-source multi-modal dataset for training,"
                " the Wukong-Huahua model has excellent Chinese text-image generation capabilities."
                " Wukong-Huahua can identify various scene descriptions and painting styles,"
                " bringing users a good experience.")

    with gr.Tab("Figure Generation"):
        style_input = gr.Textbox(lines=1,
                                 placeholder="Input the style of figure you want to generate.",
                                 label="例如：新海诚 宫崎骏 梵高 莫奈 赛博朋克",
                                 elem_id="style-input")
        desc_input = gr.Textbox(lines=1,
                                placeholder="Input a sentence to describe the figure you want to generate.",
                                label="")
        gr.Examples(
            examples=examples,
            inputs=[style_input, 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("- [The Foundation models Platform for Mindspore](https://xihe.mindspore.cn/)")

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

demo.launch()