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

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

def generate_image(prompt, aspect_ratio):
    """
    Calls the external API to generate an image based on the prompt and aspect ratio.
    """
    if not API_URL or not API_TOKEN:
        raise gr.Error("API_URL or API_TOKEN not set. Please add them to your Space secrets.")

    headers = {
        "Authorization": API_TOKEN
    }

    form_data = {
        "prompt": (None, prompt),
        "aspect_ratio": (None, aspect_ratio),
    }

    print(f"Generating Image... Aspect Ratio: {aspect_ratio}")

    try:
        response = requests.post(
            API_URL, 
            headers=headers, 
            files=form_data, 
            timeout=180
        )
        
        response.raise_for_status()

        content_type = response.headers.get('content-type', '')
        if 'image' not in content_type:
            raise gr.Error(f"API did not return an image. Response: {response.text}")
        image = Image.open(io.BytesIO(response.content))
        print("Image received successfully.")
        return image

    except requests.exceptions.Timeout:
        raise gr.Error("The request timed out. The model might be taking too long to respond.")
    except requests.exceptions.RequestException as e:
        raise gr.Error(f"An API error occurred: {e}")
    except Exception as e:
        raise gr.Error(f"An unexpected error occurred: {e}")

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🖼️ ImagineArt 1.0")
    gr.Markdown("A state-of-the-art Mixture of Experts (MoE) model for generating hyper-realistic images with unmatched detail, natural lighting, and photographic authenticity.")
    
    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(
                label="Prompt", 
                lines=5, 
                placeholder="A highly detailed photorealistic image of..."
            )
            aspect_ratio_input = gr.Dropdown(
                label="Aspect Ratio", 
                choices=["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"], 
                value="1:1"
            )
            submit_button = gr.Button("Generate", variant="primary")
        
        with gr.Column(scale=1):
            image_output = gr.Image(label="Generated Image", show_download_button=True)

    gr.Examples(
        examples=[
            [
                "Photorealistic editorial photograph, a vintage ornate pocket watch with intricate engravings, resting angled on a deep crimson velvet cushion, soft candlelight casting warm, elongated shadows. The composition is centered on the pocket watch, angled to showcase its intricate craftsmanship and timeless beauty. High resolution, crisp details.", 
                "1:1"
            ],
            [
                "Epic cinematic shot of a lone astronaut standing on a desolate Mars-like planet, watching two suns set against a dusty, orange sky. The landscape is rocky with vast canyons. The astronaut's suit has subtle reflections of the alien sunset. Style of Denis Villeneuve.",
                "16:9"
            ],
            [
                "An authentic, unposed candid photograph of a woman sitting in a cafe.",
                "4:3"
            ]
        ],
        inputs=[prompt_input, aspect_ratio_input]
    )

    submit_button.click(
        fn=generate_image,
        inputs=[prompt_input, aspect_ratio_input],
        outputs=image_output
    )

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