"""
Byte Dream - Gradio Web Interface
Interactive web UI for image generation
"""

import gradio as gr
from bytedream.generator import ByteDreamGenerator
import torch
import os


# Initialize generator
print("Loading Byte Dream model...")

# Check if we should load from Hugging Face
HF_REPO_ID = os.getenv("HF_REPO_ID", None)
MODEL_PATH = os.getenv("MODEL_PATH", "./models/bytedream")

try:
    if HF_REPO_ID:
        print(f"Loading model from Hugging Face: {HF_REPO_ID}...")
        generator = ByteDreamGenerator(
            hf_repo_id=HF_REPO_ID,
            config_path="config.yaml",
            device="cpu",
        )
    else:
        print(f"Loading model from local path: {MODEL_PATH}...")
        generator = ByteDreamGenerator(
            model_path=MODEL_PATH,
            config_path="config.yaml",
            device="cpu",
        )
    print("✓ Model loaded successfully!")
except Exception as e:
    print(f"⚠ Warning: Could not load model: {e}")
    print("  Please train the model first using: python train.py")
    print("  Or download pretrained weights from Hugging Face.")
    print("")
    print("  To use a model from Hugging Face, set environment variable:")
    print("    HF_REPO_ID=username/repo_name")
    print("")
    print("Starting in demo mode...")
    generator = None


def generate_image(
    prompt,
    negative_prompt,
    width,
    height,
    num_steps,
    guidance_scale,
    seed,
):
    """Generate image from prompt"""
    
    if generator is None:
        return None, "Error: Model not loaded. Please train or download model weights."
    
    # Convert seed to None if -1
    seed_value = None if seed == -1 else seed
    
    try:
        # Generate image
        image = generator.generate(
            prompt=prompt,
            negative_prompt=negative_prompt if negative_prompt else None,
            width=int(width),
            height=int(height),
            num_inference_steps=int(num_steps),
            guidance_scale=float(guidance_scale),
            seed=seed_value,
        )
        
        return image, "Success! ✓"
    
    except Exception as e:
        print(f"Error generating image: {e}")
        import traceback
        traceback.print_exc()
        return None, f"Error: {str(e)}"


# Create Gradio interface
with gr.Blocks(
    title="Byte Dream - AI Image Generator",
    css="""
    .gradio-container {
        max-width: 1400px !important;
    }
    #main-heading {
        text-align: center;
        margin-bottom: 20px;
    }
    .description {
        text-align: center;
        margin-bottom: 30px;
    }
    """
) as demo:
    
    gr.Markdown("""
    # 🎨 Byte Dream - AI Image Generator
    
    ### Transform your imagination into reality with advanced AI
    
    Powered by state-of-the-art latent diffusion models, optimized for CPU inference.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📝 Create Your Prompt")
            
            prompt_input = gr.Textbox(
                label="Positive Prompt",
                placeholder="Describe the image you want to create...",
                lines=3,
                value="A beautiful sunset over mountains, digital art, highly detailed, vibrant colors",
            )
            
            negative_prompt_input = gr.Textbox(
                label="Negative Prompt (Optional)",
                placeholder="What to avoid in the image...",
                lines=2,
                value="ugly, blurry, low quality, distorted, deformed",
            )
            
            gr.Markdown("### ⚙️ Settings")
            
            with gr.Row():
                width_slider = gr.Slider(
                    minimum=256,
                    maximum=1024,
                    step=64,
                    value=512,
                    label="Width (px)",
                    info="Image width - multiples of 64"
                )
                height_slider = gr.Slider(
                    minimum=256,
                    maximum=1024,
                    step=64,
                    value=512,
                    label="Height (px)",
                    info="Image height - multiples of 64"
                )
            
            with gr.Row():
                steps_slider = gr.Slider(
                    minimum=10,
                    maximum=150,
                    step=5,
                    value=50,
                    label="Inference Steps",
                    info="More steps = better quality but slower"
                )
                guidance_slider = gr.Slider(
                    minimum=1.0,
                    maximum=20.0,
                    step=0.5,
                    value=7.5,
                    label="Guidance Scale",
                    info="Higher = closer to prompt, Lower = more creative"
                )
            
            seed_input = gr.Number(
                label="Seed",
                value=-1,
                precision=0,
                info="-1 for random, any number for reproducibility",
            )
            
            generate_btn = gr.Button(
                "🎨 Generate Image",
                variant="primary",
                size="lg",
            )
        
        with gr.Column(scale=1):
            gr.Markdown("### 🖼️ Result")
            
            output_image = gr.Image(
                label="Generated Image",
                type="pil",
                height=512,
            )
            status_text = gr.Textbox(
                label="Status",
                interactive=False,
            )
            
            download_btn = gr.File(
                label="Download",
                visible=True,
            )
    
    # Tips section
    with gr.Accordion("💡 Tips for Better Results", open=False):
        gr.Markdown("""
        **Writing Effective Prompts:**
        - Be specific and descriptive
        - Include art style references (e.g., "digital art", "oil painting", "watercolor")
        - Mention lighting ("dramatic lighting", "soft sunlight", "neon lights")
        - Add quality modifiers ("highly detailed", "4K", "masterpiece")
        - Specify mood and atmosphere ("peaceful", "dramatic", "mysterious")
        
        **Using Negative Prompts:**
        - Remove unwanted elements ("no people", "no text")
        - Avoid quality issues ("no blur", "no distortion")
        - Fix common problems ("bad anatomy", "extra limbs")
        
        **Parameter Guide:**
        - **Steps**: 20-30 for quick previews, 50-75 for final images, 100+ for best quality
        - **Guidance**: 5-7 for creative freedom, 7-9 for balanced, 9-12 for strict prompt following
        - **Resolution**: Higher = more detail but slower. Start with 512x512, increase if needed
        """)
    
    # Examples section
    gr.Markdown("### 💡 Example Prompts")
    
    with gr.Row():
        example_btn1 = gr.Button(
            "🌆 Cyberpunk City",
            size="sm",
            elem_id="example_btn1",
        )
        example_btn2 = gr.Button(
            "🐉 Fantasy Dragon",
            size="sm",
            elem_id="example_btn2",
        )
        example_btn3 = gr.Button(
            "🏔️ Peaceful Landscape",
            size="sm",
            elem_id="example_btn3",
        )
    
    with gr.Row():
        example_btn4 = gr.Button(
            "👤 Character Portrait",
            size="sm",
            elem_id="example_btn4",
        )
        example_btn5 = gr.Button(
            "🌊 Underwater Scene",
            size="sm",
            elem_id="example_btn5",
        )
        example_btn6 = gr.Button(
            "🎨 Abstract Art",
            size="sm",
            elem_id="example_btn6",
        )
    
    # Example prompt values
    example_prompts = {
        "example_btn1": (
            "A cyberpunk city at night with neon lights, futuristic architecture, flying cars, rain-slicked streets, highly detailed, digital art, cinematic lighting",
            "ugly, blurry, low quality, distorted, dark, gloomy"
        ),
        "example_btn2": (
            "A majestic dragon breathing fire, fantasy art, dramatic lighting, epic scene, scales gleaming, powerful wings, mountain landscape background",
            "ugly, deformed, blurry, low quality, cartoonish"
        ),
        "example_btn3": (
            "A peaceful cottage in a meadow, wildflowers, sunny day, blue sky, studio ghibli style, serene atmosphere, pastoral landscape",
            "people, animals, buildings, urban, dark, stormy"
        ),
        "example_btn4": (
            "Portrait of a warrior princess, ornate armor, fantasy setting, intricate details, character design, dramatic lighting, confident expression, long flowing hair",
            "ugly, deformed, asymmetrical, blurry, low quality, bad anatomy"
        ),
        "example_btn5": (
            "Underwater coral reef, tropical fish, sunlight filtering through water, photorealistic, vibrant colors, marine life, crystal clear water",
            "polluted, murky, dark, blurry, low quality"
        ),
        "example_btn6": (
            "Abstract geometric art, colorful shapes, dynamic composition, modern art, bold patterns, artistic expression, vivid colors",
            "representational, realistic, boring, dull colors, simple"
        ),
    }
    
    # Connect example buttons
    def set_example(prompt, negative):
        return prompt, negative, "Click Generate to create!"
    
    # Map button names to their variables
    button_map = {
        "example_btn1": example_btn1,
        "example_btn2": example_btn2,
        "example_btn3": example_btn3,
        "example_btn4": example_btn4,
        "example_btn5": example_btn5,
        "example_btn6": example_btn6,
    }
    
    for btn_name, (prompt, negative) in example_prompts.items():
        button_map[btn_name].click(
            fn=set_example,
            inputs=[gr.State(prompt), gr.State(negative)],
            outputs=[prompt_input, negative_prompt_input, status_text],
        )
    
    # Connect generate button
    generate_btn.click(
        fn=generate_image,
        inputs=[
            prompt_input,
            negative_prompt_input,
            width_slider,
            height_slider,
            steps_slider,
            guidance_slider,
            seed_input,
        ],
        outputs=[output_image, status_text],
    )
    
    # Footer
    gr.Markdown("""
    ---
    **Byte Dream** v1.0.0 | Powered by Latent Diffusion Models | Optimized for CPU Inference
    
    Created with ❤️ using PyTorch and Hugging Face Diffusers
    """)


if __name__ == "__main__":
    print("\n" + "="*60)
    print("Starting Byte Dream Web Interface")
    print("="*60)
    print("\nOpening browser...")
    print("Press Ctrl+C to close\n")
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        theme=gr.themes.Soft(),
    )