import os
import re
import random
import json
import requests
import gradio as gr
from gradio_client import Client, file
import numpy as np
from PIL import Image
from typing import Tuple
from translatepy import Translator
from io import BytesIO

MODEL = os.environ.get("MODEL")
MODEL2 = os.environ.get("MODEL2")
MODEL_TEST = os.environ.get("MODEL_TEST")

DESCRIPTION = """
# IMGEN🧚‍♀️
<center>🧚‍♂️输入描述生成图片，选择喜欢的风格~<br>🧚默认使用Turbo快速生成，高级选项可更改。</center>
"""
MODEL_ID="stabilityai/sdxl-turbo"

LICENSE = "采用 SDXL 模型"
translator = Translator()



def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

MAX_SEED = np.iinfo(np.int32).max

client = Client(MODEL)
client2 = Client(MODEL2)
client_test=Client(MODEL_TEST)

style_list = [
    {
        "name": "(无风格)",
        "prompt": "{prompt}",
    },
    {
        "name": "电影",
        "prompt": "cinematic, {prompt}, detailed, bokeh, moody, film grain",
    },
    {
        "name": "摄影",
        "prompt": "realistic photo, {prompt}, 35mm photograph, 4k",
    },
    {
        "name": "动画",
        "prompt": "anime style, {prompt}, studio anime",
    },
    {
        "name": "漫画",
        "prompt": "manga style, {prompt}, detailed, Japanese comic style",
    },
    {
        "name": "数绘",
        "prompt": "{prompt}, digital artwork, detailed",
    },
    {
        "name": "像素",
        "prompt": "pixel-art, {prompt}, 8-bit graphics",
    },
    {
        "name": "幻想",
        "prompt": "fantasy style, {prompt}, epic, magical, dreamy",
    },
    {
        "name": "朋克",
        "prompt": "cyberpunk style, {prompt}, stunningly beautiful, ultra detailed, realistic",
    },
    {
        "name": "三维",
        "prompt": "3d model, {prompt}, highly quality, dramatic lighting",
    },
]   
styles = {k["name"]: (k["prompt"]) for k in style_list}
STYLE_NAMES = list(styles.keys())
DEFAULT_STYLE_NAME = "(无风格)"

def apply_style(style_name: str, positive: str) -> Tuple[str]:
    p = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
    return p.replace("{prompt}", positive)

def list_dir(path):
    for root, dirs, files in os.walk(path):
        for file in files:
            print(os.path.join(root, file))

def generate(
    prompt: str,
    seed: float = 0,
    steps: float = 1,
    style: str = DEFAULT_STYLE_NAME,
    randomize_seed: bool = False,
    use_turbo: bool = True,
    use_flash: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    prompt = str(translator.translate(prompt, 'English'))
    print(f'prompt:{prompt}')
    seed = int(randomize_seed_fn(seed, randomize_seed))
    #print(client_test.view_api())

    prompt = apply_style(style, prompt)
    strength = 0.7
    images = []
    if use_turbo and use_flash is False:
        result = client.predict(
            None,
    		prompt,
            strength,
            steps,
            seed,
    		fn_index=1
        )
        print(result)
        image_url = MODEL+"file="+result
        images_request = requests.get(image_url)
        print(f'A:{images_request}')
        pil_image = Image.open(BytesIO(images_request.content))
        images.append(pil_image)
        print(f'B:{images}')
    elif use_flash is True:
        result = client_test.predict(
            prompt,
            0,
            True,
            api_name="/infer"
        )
        print(result)
        image_url = MODEL_TEST+"file="+result
        images_request = requests.get(image_url)
        print(f'A:{images_request}')
        pil_image = Image.open(BytesIO(images_request.content))
        images.append(pil_image)
        print(f'B:{images}')        
    else:
        seed = 0
        job = client2.submit(
            prompt,
            "low_quality",
            7.5,
            "(No style)",
            fn_index=3    
        )
        result = job.result()

        for root, dirs, files in os.walk(result):
            for file in files:
                print(file)
                if file.endswith('.jpg'):
                    images.append(os.path.join(root,file))
                    
    return images, seed

examples = [
    "坐在火箭上的秋田犬",
    "一只吃芝士片的猫",
    "太空中骑马的宇航员",
    "放学回家的孩子们",
    "一个可爱的机器人在照镜子",
    "钢铁侠在恐龙时代"
]

CSS = '''
.gradio-container{max-width: 560px !important}
h1{text-align:center}
footer {
    visibility: hidden
}
'''
with gr.Blocks(css=CSS, theme="xiaobaiyuan/theme_brief") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Group():
        with gr.Row():
            prompt = gr.Text(
                label="描述",
                show_label=False,
                max_lines=1,
                placeholder="画什么好呢",
                container=False,
                scale=3,
            )
            run_button = gr.Button("生成", scale=1)
        result = gr.Gallery(label="作品", columns = 1, preview=True)
    with gr.Accordion("高级选项", open=False):
        with gr.Row():
            steps = gr.Slider(
                label="步数",
                minimum=1,
                maximum=10,
                step=1,
                value=2,
            )
            seed = gr.Slider(
                label="种子",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
                visible=True
            )
            randomize_seed = gr.Checkbox(label="随机种子", value=True)
            use_turbo = gr.Checkbox(label="快速生成", value=True)
            use_flash = gr.Checkbox(label="Flash", value=False)
    with gr.Row(visible=True):
        style_selection = gr.Radio(
            show_label=True,
            container=True,
            interactive=True,
            choices=STYLE_NAMES,
            value=DEFAULT_STYLE_NAME,
            label="风格化",
        ) 
    with gr.Row():
        gr.Examples(
            examples=examples,
            inputs=prompt,
            outputs=[result, seed],
            fn=generate,
            cache_examples='lazy',
            label="示例",
        )
    with gr.Row():
        gr.Markdown(LICENSE)
    
    gr.on(
        triggers=[
            prompt.submit,
            run_button.click,
        ],
        fn=generate,
        inputs=[
            prompt,
            seed,
            steps,
            style_selection,
            randomize_seed,
            use_turbo,
            use_flash,
        ],
        outputs=[result, seed],
        api_name="run",
    )


        
if __name__ == "__main__":
    demo.queue(max_size=20).launch(show_api=False, debug=False)