import os
import base64
import io

from dotenv import load_dotenv
import requests
import gradio as gr
from PIL import Image

load_dotenv()

API_URL = os.getenv("IMAGE_API_URL", "https://llm.dev.maimaiag.com/v1/images/generations")
API_KEY = os.getenv("IMAGE_API_KEY", "")


def _call_remote_api(prompt: str, size: str, seed: int, num_inference_steps: int) -> Image.Image:
    headers = {
        "Content-Type": "application/json",
    }
    if API_KEY:
        headers["Authorization"] = f"Bearer {API_KEY}"

    payload = {
        "model": "openai/z-image-turbo",
        "prompt": prompt,
        "size": size,
        "seed": seed,
        "num_inference_steps": num_inference_steps,
    }

    response = requests.post(API_URL, headers=headers, json=payload, timeout=300)
    response.raise_for_status()
    data = response.json()

    b64_image = data["data"][0]["b64_json"]
    image_bytes = base64.b64decode(b64_image)
    return Image.open(io.BytesIO(image_bytes))


def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
    """Generate an image from the given prompt via remote API."""
    if randomize_seed:
        import random
        seed = random.randint(0, 2**32 - 1)

    size = f"{int(width)}x{int(height)}"
    image = _call_remote_api(
        prompt=prompt,
        size=size,
        seed=int(seed),
        num_inference_steps=int(num_inference_steps),
    )
    return image, seed

# 示例提示词
examples = [
    ["一张黑白照片（约1950年代），捕捉了一场欢乐的跨代家庭感恩节或节日晚餐。中心人物是一位身穿白衬衫、系深色领带的微笑男子，他正在餐桌主位上兴致勃勃地切割一只大烤火鸡。他被一大群各年龄段的家庭成员包围，所有人都迫不及待地伸出盘子来接食物。许多不同年龄的孩子围在周围，眼睛睁得大大的，充满期待，有的站着，有的坐着。还有几位女性和青少年，有的在帮忙上菜，有的抱着婴儿。整个场景充满了自然的互动、笑声和热闹温馨的氛围。餐桌上摆满了传统的节日菜肴。背景是一面简单的、可能贴着壁纸的墙壁，反映了真实家庭住宅的质朴感。"],
    ["一幅全景场景，背景从石器时代过渡到未来（洞穴→金字塔→城堡→工厂→摩天大楼→漂浮城市），主体是前景中同一张面孔/同一个人，从左到右佩戴着符合时代特征的头盔：骨/兽皮头饰、青铜古代头盔、中世纪板甲头盔、一战钢盔、现代太空头盔，以及未来能量/全息头盔。"],
    ["一张现代会议室里的商务团队照片。在会议桌的首席位置，一位自信的老板站着，充满热情地介绍一个雄心勃勃的新产品创意。围坐在桌旁的雇员们反应各异，有的好奇、有的挑眉、有的若有所思，有的在记笔记，有的在提问。他们身后的大窗外可以看到摩天大楼和城市灯火。氛围专业但充满紧张和引人入胜的气息。"],
    ["一幅文艺复兴风格的绘画，描绘了一座现代城市景观。"],
    ["节日期间繁忙的城市街道，彩旗飘扬，人群熙攘，街头艺人在表演。"],
    ["科幻电影的电影海报，宇航员站在火星上凝望地球，戏剧性的灯光，黑暗的氛围，顶部有大号金属质感字体标题\"THE VOYAGE\""],
    ["兴奋的年轻内容创作者指着漂浮的机器人全息图，充满活力的紫色和蓝色霓虹背景，高对比度，YouTube缩略图风格，粗体黄色文字\"AI UPDATE\""],
    ["现代工作空间，木质桌面上放着笔记本电脑和咖啡，晨光透过窗户洒入，温馨的氛围，照片级真实感，高质量，便利贴上写着\"Productivity\""],
]

# 自定义主题，深色科技感 (Gradio 6)
custom_theme = gr.themes.Soft(
    primary_hue="cyan",
    secondary_hue="violet",
    neutral_hue="zinc",
    font=gr.themes.GoogleFont("Inter"),
    text_size="lg",
    spacing_size="md",
    radius_size="lg"
).set(
    button_primary_background_fill="*primary_500",
    button_primary_background_fill_hover="*primary_400",
    button_primary_text_color="*neutral_950",
    block_title_text_weight="600",
    block_background_fill="*neutral_900",
    body_background_fill="*neutral_950",
    body_text_color="*neutral_200",
    block_label_text_color="*neutral_300",
    input_background_fill="*neutral_800",
    border_color_primary="*neutral_700",
)

# 构建 Gradio 界面
with gr.Blocks(fill_height=True) as demo:
    # 页眉
    gr.Markdown(
        """
        # Z-Image-Turbo
        极速 AI 图像生成 • 仅需 8 步即可生成精美图像
        """,
        elem_classes="header-text"
    )
    
    with gr.Row(equal_height=False):
        # 左栏 - 输入控制
        with gr.Column(scale=1, min_width=320):
            prompt = gr.Textbox(
                label="提示词",
                placeholder="描述你想要生成的图像...",
                lines=5,
                max_lines=10,
                autofocus=True,
            )
            
            with gr.Accordion("高级设置", open=False):
                with gr.Row():
                    height = gr.Slider(
                        minimum=512,
                        maximum=2048,
                        value=1024,
                        step=64,
                        label="高度",
                        info="图像高度（像素）"
                    )
                    width = gr.Slider(
                        minimum=512,
                        maximum=2048,
                        value=1024,
                        step=64,
                        label="宽度",
                        info="图像宽度（像素）"
                    )
                
                num_inference_steps = gr.Slider(
                    minimum=1,
                    maximum=20,
                    value=9,
                    step=1,
                    label="推理步数",
                    info="9 步 = 8 次 DiT 前向传播（推荐）"
                )
                
                with gr.Row():
                    randomize_seed = gr.Checkbox(
                        label="随机种子",
                        value=True,
                    )
                    seed = gr.Number(
                        label="种子",
                        value=42,
                        precision=0,
                        visible=False,
                    )
                
                def toggle_seed(randomize):
                    return gr.Number(visible=not randomize)
                
                randomize_seed.change(
                    toggle_seed,
                    inputs=[randomize_seed],
                    outputs=[seed]
                )
            
            generate_btn = gr.Button(
                "生成图像",
                variant="primary",
                size="lg",
                scale=1
            )
            
            # 示例提示词
            gr.Examples(
                examples=examples,
                inputs=[prompt],
                label="试试这些提示词",
                examples_per_page=5,
            )
        
        # 右栏 - 输出
        with gr.Column(scale=1, min_width=320):
            output_image = gr.Image(
                label="生成的图像",
                type="pil",
                format="png",
                show_label=False,
                height=600,
                buttons=["download", "share"],
            )
            
            used_seed = gr.Number(
                label="使用的种子",
                interactive=False,
                container=True,
            )
    
    # 连接生成按钮
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
        outputs=[output_image, used_seed],
    )
    
    # 也支持在提示词框中按回车键生成
    prompt.submit(
        fn=generate_image,
        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
        outputs=[output_image, used_seed],
    )

if __name__ == "__main__":
    demo.launch(
        theme=custom_theme,
        css="""
        .header-text h1 {
            font-size: 2.5rem !important;
            font-weight: 700 !important;
            margin-bottom: 0.5rem !important;
            background: linear-gradient(135deg, #22d3ee 0%, #a78bfa 100%);
            -webkit-background-clip: text;
            -webkit-text-fill-color: transparent;
            background-clip: text;
            letter-spacing: -0.02em;
        }

        .header-text p {
            font-size: 1.1rem !important;
            color: #94a3b8 !important;
            margin-top: 0 !important;
        }

        .footer-text {
            padding: 1rem 0;
        }

        .footer-text a {
            color: #22d3ee !important;
            text-decoration: none !important;
            font-weight: 500;
        }

        .footer-text a:hover {
            text-decoration: underline !important;
        }

        /* Mobile optimizations */
        @media (max-width: 768px) {
            .header-text h1 {
                font-size: 1.8rem !important;
            }

            .header-text p {
                font-size: 1rem !important;
            }
        }

        /* Smooth transitions */
        button, .gr-button {
            transition: all 0.2s ease !important;
        }

        button:hover, .gr-button:hover {
            transform: translateY(-1px);
            box-shadow: 0 0 20px rgba(34, 211, 238, 0.35) !important;
        }

        /* Better spacing */
        .gradio-container {
            max-width: 1400px !important;
            margin: 0 auto !important;
        }

        /* Glassmorphism panels */
        .gradio-container .gr-block,
        .gradio-container .gr-form,
        .gradio-container .gr-box {
            backdrop-filter: blur(8px) !important;
            border: 1px solid rgba(148, 163, 184, 0.12) !important;
        }

        /* Input focus glow */
        input:focus, textarea:focus {
            box-shadow: 0 0 0 2px rgba(34, 211, 238, 0.25) !important;
        }
        """,
        footer_links=[
            "api",
            "gradio"
        ],
        mcp_server=True
    )