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

# Model information dictionary - maps model names to their descriptions
MODEL_INFO = {
    "None": "Default model (good for general purposes)",
    "stability": "Stability AI's Stable Diffusion (balanced quality)",
    "openjourney": "Midjourney-style images",
    "dreamlike": "Dreamlike photorealistic art",
    "protogen": "Sci-fi and futuristic styles",
    "sd14": "Stable Diffusion 1.4 (classic model)",
    "sd21": "Stable Diffusion 2.1 (improved details)",
    "kontext": "Context-aware generation",
    "gptimage": "GPT-assisted image generation",
    "anything": "Anime-style generation",
    "realistic": "Photorealistic outputs",
    "analog": "Film photography style",
    "arcanediffusion": "Fantasy art style",
    "dreamshaper": "Enhanced creative variations",
    "deliberate": "Precise prompt following",
    "revanimated": "Animated/cartoon styles",
    "meinamix": "Anime/portrait hybrid"
}

# ------------------------------------------------------------------
# Helper: fetch list of available models
# ------------------------------------------------------------------
def get_models():
    try:
        r = requests.get("https://image.pollinations.ai/models", timeout=10)
        models = [m.strip() for m in r.text.splitlines() if m.strip()]
    except:
        # Fallback if API is unavailable
        models = list(MODEL_INFO.keys())
    
    # Ensure our known models are included
    for m in MODEL_INFO:
        if m not in models and m != "None":
            models.append(m)
    
    # Sort alphabetically but keep "None" first
    models = sorted(models)
    if "None" in models:
        models.remove("None")
        models.insert(0, "None")
    return models

# ------------------------------------------------------------------
# Image-generation function
# ------------------------------------------------------------------
def generate_image(prompt, model):
    if model == "None":
        model = ""  # API expects empty string for default
        
    url = f"https://image.pollinations.ai/prompt/{prompt}?nologo=true&seed=random"
    if model:
        url += f"&model={model}"
    
    try:
        r = requests.get(url, timeout=30)
        r.raise_for_status()
        img = Image.open(io.BytesIO(r.content))
        return img
    except Exception as e:
        raise gr.Error(f"Failed to generate image: {str(e)}")

# ------------------------------------------------------------------
# Gradio UI with improved model selector
# ------------------------------------------------------------------
models = get_models()

with gr.Blocks(theme="soft") as demo:
    gr.Markdown("""# Openimage
    ### Select a model style that matches your desired output
    """)
    
    with gr.Row():
        with gr.Column(scale=3):
            prompt_box = gr.Textbox(
                label="Prompt",
                placeholder="A cat wearing sunglasses on the moon...",
                lines=2
            )
            
            # Improved model dropdown with descriptions
            model_dropdown = gr.Dropdown(
                choices=models,
                value="None",
                label="Model Style",
                info="Select a generation style",
                interactive=True,
                filterable=True,  # Allows searching through models
                elem_classes="model-selector"
            )
            
            # Add model description display
            model_description = gr.Markdown(MODEL_INFO["None"])
            
            # Update description when model changes
            def update_description(model):
                return MODEL_INFO.get(model, "No description available")
            
            model_dropdown.change(
                fn=update_description,
                inputs=model_dropdown,
                outputs=model_description
            )
            
            generate_btn = gr.Button("Generate", variant="primary")
        
        with gr.Column(scale=3):
            output_img = gr.Image(type="pil", label="Generated Image")
    
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt_box, model_dropdown],
        outputs=output_img
    )

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