import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModel
from PIL import Image
import numpy as np
import random
from rudalle import get_rudalle_model, get_tokenizer, get_vae, get_realesrgan
from rudalle.pipelines import generate_images, super_resolution
from rudalle.utils import seed_everything

# Load model components (these would be loaded in a real deployment)
def load_model():
    """Load the ruDALL-E Malevich model components"""
    try:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        model = get_rudalle_model('Malevich', pretrained=True, fp16=True, device=device)
        tokenizer = get_tokenizer()
        vae = get_vae().to(device)
        realesrgan = get_realesrgan('x4', device=device)
        return model, tokenizer, vae, realesrgan, device
    except Exception as e:
        print(f"Model loading failed: {e}")
        return None, None, None, None, 'cpu'

# Initialize model
model, tokenizer, vae, realesrgan, device = load_model()

def generate_images(prompt, negative_prompt=""):
    """Generate 4 images using ruDALL-E Malevich model"""
    if model is None:
        # Fallback: generate placeholder images if model fails to load
        images = []
        for i in range(4):
            # Create colorful placeholder with different patterns
            img_array = np.random.randint(50, 255, (512, 512, 3), dtype=np.uint8)
            # Add some structure to make it look more like generated art
            for _ in range(5):
                x, y = random.randint(0, 512), random.randint(0, 512)
                radius = random.randint(20, 100)
                color = [random.randint(0, 255) for _ in range(3)]
                for i in range(max(0, x-radius), min(512, x+radius)):
                    for j in range(max(0, y-radius), min(512, y+radius)):
                        if (i-x)**2 + (j-y)**2 <= radius**2:
                            img_array[i, j] = color
            img = Image.fromarray(img_array)
            images.append(img)
        return images
    
    # Use empty prompt if none provided
    if not prompt.strip():
        prompt = "абстрактное искусство"  # "abstract art" in Russian
    
    # Handle negative prompt
    if negative_prompt:
        # In a real implementation, you would incorporate negative prompting
        # For now, we'll just pass it through
        pass
    
    try:
        # Generate 4 images
        images = []
        for _ in range(4):
            _pil_images = generate_images(prompt, tokenizer, model, vae, top_k=512, top_p=0.99, images_num=1)
            # Upscale images
            _pil_images = super_resolution(_pil_images, realesrgan)
            images.extend(_pil_images)
        return images[:4]
    except Exception as e:
        print(f"Generation failed: {e}")
        # Fallback to placeholder images
        return generate_images("", "")

def select_image(gallery, select_data: gr.SelectData):
    """Handle image selection from gallery"""
    if select_data is not None and gallery and len(gallery) > select_data.index:
        return gallery[select_data.index]
    return None

# Create the Gradio 6 application
with gr.Blocks() as demo:
    gr.Markdown("# 🎨 ruDALL-E Malevich Image Generation")
    gr.Markdown("Generate images using the ai-forever/rudalle-Malevich model")
    gr.Markdown("[Built with anycoder](https://huggingface.co/spaces/akhaliq/anycoder)")
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="Enter your prompt here (can be empty)",
                lines=3,
                value="",
                info="Leave empty for random art generation"
            )
            
            negative_prompt_input = gr.Textbox(
                label="Negative Prompt",
                placeholder="Enter what to avoid in the image",
                lines=2,
                value="",
                info="Describe elements you don't want in the image"
            )
            
            with gr.Row():
                generate_btn = gr.Button("🎲 Generate 4 Images", variant="primary", size="lg")
                clear_btn = gr.Button("Clear", variant="secondary")
        
        with gr.Column(scale=2):
            gr.Markdown("### Generated Gallery")
            gallery = gr.Gallery(
                label="Click any image to view larger",
                show_label=False,
                elem_id="gallery",
                columns=2,
                rows=2,
                height="auto",
                allow_preview=True,
                object_fit="cover",
                show_download_button=True,
                show_share_button=True
            )
    
    with gr.Row():
        selected_image = gr.Image(
            label="Selected Image (Click gallery image to view)",
            type="pil",
            height=600,
            width=600,
            interactive=False,
            show_label=True
        )
    
    # Event handlers
    generate_btn.click(
        fn=generate_images,
        inputs=[prompt_input, negative_prompt_input],
        outputs=gallery,
        api_visibility="public"
    )
    
    gallery.select(
        fn=select_image,
        inputs=[gallery],
        outputs=selected_image,
        api_visibility="private"
    )
    
    clear_btn.click(
        fn=lambda: ([], None),
        inputs=[],
        outputs=[gallery, selected_image],
        api_visibility="private"
    )
    
    # Add example prompts
    gr.Examples(
        examples=[
            ["космический пейзаж", "люди, здания"],
            ["абстрактная композиция", ""],
            ["", "текст, надписи"],
            ["магический лес", "город, техника"],
            ["портрет в стиле импрессионизма", "современные элементы"]
        ],
        inputs=[prompt_input, negative_prompt_input],
        label="Example Prompts",
        examples_per_page=5
    )

# Launch the demo with Gradio 6 theme
demo.launch(
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="indigo",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
        text_size="lg",
        spacing_size="lg",
        radius_size="md"
    ).set(
        button_primary_background_fill="*primary_600",
        button_primary_background_fill_hover="*primary_700",
        block_title_text_weight="600",
        block_border_width="1px",
        block_border_color="*neutral_200"
    ),
    footer_links=[
        {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}
    ]
)