import os
import requests
import numpy as np
from PIL import Image
import gradio as gr
from io import BytesIO


def gen_image(desc: str):
    """generate the image from the wukong huahua model of ascend server in Wuhan AICC
    Args:
        desc(str): the input description text
    """
    if not desc:
        return
    access_token = os.environ['token']
    headers = {'content-type': "application/json", 'X-Subject-Token': access_token}

    # url = f"https://a2f051d4cabf45f885d7b0108edc9b9c.infer.ovaijisuan.com/v1/infers/dce9ad51-7cde-4eeb-8291-ae29f267ed2c/wukong_hf"
    url = f"https://a2f051d4cabf45f885d7b0108edc9b9c.infer.ovaijisuan.com/v1/infers/975eedfd-6e15-4571-8ca9-b945da0da24b/wukong_hf"
    body = {
        "desc": desc,
        "style": ""
    }

    resp_data = requests.post(url, json=body, headers=headers)
    if resp_data['status'] != 200:
        return []
    # img_np = np.array(resp_data['output_image_list'][0])
    # image = Image.fromarray(np.uint8(img_np))
    img_rep = requests.get(resp_data['output_image_url'][0])
    image = Image.open(BytesIO(img_rep.content))

    return [image]


examples = [
    '天空之城 赛博朋克 动漫',
    '秋水共长天一色',
    '海滩 蓝天 美景',
    '教堂 巴洛克风格',
    '落日 莫奈',
    '来自深渊 雪山飞狐'
]

block = gr.Blocks()


with block:
    with gr.Group():
        with gr.Box():
            with gr.Row().style(mobile_collapse=False, equal_height=True):
                text = gr.Text(
                    placeholder="输入中文，生成图片",
                ).style(
                    border=(True, False, True, True),
                    rounded=(True, False, False, True),
                    container=False,
                )

                btn = gr.Button("Generate image").style(
                    margin=False,
                    rounded=(False, True, True, False),
                )

        gallery = gr.Gallery(
            label="Generated images", show_label=False, elem_id="gallery"
        ).style(grid=[1, 1], height="auto")

        gr.Examples(examples=examples, fn=gen_image, inputs=text, outputs=gallery)

block.queue(concurrency_count=3).launch(debug=True)