import base64
import io
import os
import requests
import gradio as gr
from PIL import Image

API_URL = os.getenv("API_URL")
API_KEY = os.getenv("API_KEY")

SIZE_OPTIONS = [
    "1024x1024", "848x1264", "1264x848",
    "1376x768", "1200x896", "896x1200", "768x1376",
]


def generate_image(prompt: str, size: str, seed: int):
    if not prompt.strip():
        raise gr.Error("Please enter a prompt to describe your image.")

    headers = {
        "Authorization": f"bearer {API_KEY}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": "ernie-image-turbo",
        "prompt": prompt,
        "n": 1,
        "response_format": "b64_json",
        "size": size,
        "seed": int(seed),
        "use_pe": True,
        "num_inference_steps": 8,
        "guidance_scale": 1.0,
    }

    try:
        response = requests.post(API_URL, headers=headers, json=payload, timeout=300)
    except requests.exceptions.Timeout:
        raise gr.Error("Request timed out. Please try again.")
    except requests.exceptions.RequestException as e:
        raise gr.Error(f"Network error: {e}")

    try:
        data = response.json()
    except Exception:
        raise gr.Error(f"Failed to parse response (HTTP {response.status_code}).")

    if not response.ok:
        err = data.get("error", {})
        msg = err.get("message") or err.get("msg") or response.text[:200]
        raise gr.Error(f"API error ({response.status_code}): {msg}")

    if "error" in data:
        raise gr.Error(f"API error: {data['error'].get('message', data['error'])}")

    items = data.get("data", [])
    if not items:
        raise gr.Error("No image returned from the API.")

    images, revised_prompts = [], []
    for item in items:
        b64 = item.get("b64_json")
        if b64:
            images.append(Image.open(io.BytesIO(base64.b64decode(b64))))
        rp = item.get("revised_prompt")
        if rp:
            revised_prompts.append(rp)

    if not images:
        raise gr.Error("Could not decode the image data.")

    return images[0], "\n\n".join(revised_prompts)


CUSTOM_CSS = """
.gradio-container { max-width: 1180px !important; margin: 0 auto !important; }
#title-block h1 { text-align: center; font-size: 2rem; margin-bottom: 0.25rem; }
#title-block p { text-align: center; color: var(--body-text-color-subdued); margin-top: 0; }
#generate-btn { min-height: 48px; font-size: 1rem; font-weight: 600; }
#output-image { border-radius: 12px; overflow: hidden; }
footer { display: none !important; }
.brand-footer { text-align: center; padding: 16px 0 8px; font-size: 0.85rem;
                color: var(--body-text-color-subdued); }
.brand-footer a { color: inherit; text-decoration: underline; }
"""

with gr.Blocks(
    title="ERNIE-Image - Free AI Image Generator",
    theme=gr.themes.Soft(
        primary_hue="orange",
        radius_size=gr.themes.sizes.radius_lg,
        font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
    ),
    css=CUSTOM_CSS,
) as demo:
    gr.Markdown(
        """
        # ERNIE-Image
        Free AI image generator powered by Baidu ERNIE-Image-Turbo 8B. Describe your scene in English or Chinese.
        """,
        elem_id="title-block",
    )

    with gr.Row(equal_height=False):
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="e.g. A black cat wearing a blue hat, photorealistic",
                lines=5,
                max_lines=10,
            )
            with gr.Row():
                size_dropdown = gr.Dropdown(
                    label="Image size",
                    choices=SIZE_OPTIONS,
                    value="1024x1024",
                )
                seed_number = gr.Number(
                    label="Seed (-1 = random)",
                    value=-1,
                    precision=0,
                )
            generate_btn = gr.Button(
                "Generate",
                variant="primary",
                elem_id="generate-btn",
            )

        with gr.Column(scale=3):
            output_gallery = gr.Image(
                label="Output",
                type="pil",
                height=480,
                elem_id="output-image",
                show_download_button=True,
            )
            with gr.Accordion("Revised prompt (model rewrite)", open=False):
                revised_prompt_output = gr.Textbox(
                    label="",
                    lines=4,
                    interactive=False,
                    show_label=False,
                )

    gr.Examples(
        examples=[
            ["一只黑猫，蓝帽子，写实风格", "1024x1024", -1],
            ["雄伟的雪山，日出时分，超写实摄影", "1024x1024", 42],
            ["A vintage travel poster of Kyoto in autumn, bold typography", "848x1264", -1],
        ],
        inputs=[prompt_input, size_dropdown, seed_number],
        label="Examples",
    )

    gr.HTML(
        '<div class="brand-footer">'
        'Powered by <a href="https://huggingface.co/baidu" target="_blank">Baidu ERNIE</a> · '
        'Visit <a href="https://ernie-image.com" target="_blank">ERNIE-Image</a> for the full experience'
        '</div>'
    )

    generate_btn.click(
        fn=generate_image,
        inputs=[prompt_input, size_dropdown, seed_number],
        outputs=[output_gallery, revised_prompt_output],
        concurrency_limit=10,
    )

if __name__ == "__main__":
    demo.queue().launch()