import gradio as gr
import torch
import spaces
from diffusers import OvisImagePipeline

# Load the model at startup
pipe = OvisImagePipeline.from_pretrained(
    "AIDC-AI/Ovis-Image-7B", 
    torch_dtype=torch.bfloat16
)
pipe.to("cuda")

@spaces.GPU(duration=120)
def generate_image(prompt: str, progress=gr.Progress(track_tqdm=True)):
    if not prompt.strip():
        raise gr.Error("Please enter a prompt to generate an image.")
    
    image = pipe(
        prompt=prompt,
        negative_prompt="",
        num_inference_steps=50,
        true_cfg_scale=5.0
    ).images[0]
    
    return image

# Modern minimal CSS with mobile-first approach
modern_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');

:root {
    --primary: #6366f1;
    --primary-hover: #4f46e5;
    --bg-light: #ffffff;
    --bg-dark: #0f172a;
    --text-light: #1e293b;
    --text-dark: #f1f5f9;
    --border-light: #e2e8f0;
    --border-dark: #334155;
    --card-light: #ffffff;
    --card-dark: #1e293b;
}

* {
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
}

body {
    background: var(--bg-light);
    transition: background 0.3s ease;
}

.dark body {
    background: var(--bg-dark);
}

.gradio-container {
    max-width: 680px !important;
    margin: 0 auto !important;
    padding: 0 20px 40px !important;
}

/* Hero Section */
.hero {
    text-align: center;
    padding: 48px 0 40px;
    margin-bottom: 32px;
}

.hero-title {
    font-size: clamp(32px, 5vw, 48px);
    font-weight: 700;
    margin: 0 0 12px 0;
    background: linear-gradient(135deg, var(--primary) 0%, #8b5cf6 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    letter-spacing: -0.02em;
}

.hero-subtitle {
    font-size: clamp(16px, 2.5vw, 20px);
    color: #64748b;
    font-weight: 400;
    margin: 0;
    line-height: 1.6;
}

.dark .hero-subtitle {
    color: #94a3b8;
}

/* Main Card */
.main-card {
    background: var(--card-light);
    border-radius: 24px;
    padding: 32px;
    box-shadow: 0 1px 3px rgba(0, 0, 0, 0.05), 0 20px 40px rgba(0, 0, 0, 0.03);
    border: 1px solid var(--border-light);
    margin-bottom: 24px;
    transition: all 0.3s ease;
}

.dark .main-card {
    background: var(--card-dark);
    border-color: var(--border-dark);
    box-shadow: 0 1px 3px rgba(0, 0, 0, 0.3), 0 20px 40px rgba(0, 0, 0, 0.2);
}

/* Input Section */
.input-section {
    margin-bottom: 24px;
}

label {
    display: block;
    margin-bottom: 12px;
}

label span {
    font-size: 14px !important;
    font-weight: 600 !important;
    color: var(--text-light) !important;
    letter-spacing: -0.01em !important;
}

.dark label span {
    color: var(--text-dark) !important;
}

textarea {
    width: 100% !important;
    background: #f8fafc !important;
    border: 2px solid transparent !important;
    border-radius: 16px !important;
    padding: 16px 18px !important;
    font-size: 15px !important;
    line-height: 1.6 !important;
    color: var(--text-light) !important;
    transition: all 0.2s ease !important;
    resize: none !important;
    box-sizing: border-box !important;
}

.dark textarea {
    background: #0f172a !important;
    color: var(--text-dark) !important;
    border-color: transparent !important;
}

textarea:focus {
    background: #ffffff !important;
    border-color: var(--primary) !important;
    box-shadow: 0 0 0 4px rgba(99, 102, 241, 0.1) !important;
    outline: none !important;
}

.dark textarea:focus {
    background: #1e293b !important;
    border-color: var(--primary) !important;
}

textarea::placeholder {
    color: #94a3b8 !important;
    opacity: 1 !important;
}

/* Button */
button.primary {
    width: 100% !important;
    background: linear-gradient(135deg, var(--primary) 0%, #8b5cf6 100%) !important;
    color: white !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 16px 24px !important;
    font-size: 16px !important;
    font-weight: 600 !important;
    cursor: pointer !important;
    transition: all 0.2s ease !important;
    box-shadow: 0 4px 12px rgba(99, 102, 241, 0.25) !important;
    min-height: 52px !important;
    letter-spacing: -0.01em !important;
}

button.primary:hover:not(:disabled) {
    transform: translateY(-1px) !important;
    box-shadow: 0 6px 20px rgba(99, 102, 241, 0.35) !important;
}

button.primary:active:not(:disabled) {
    transform: translateY(0) !important;
}

button.primary:disabled {
    opacity: 0.6 !important;
    cursor: not-allowed !important;
}

/* Output Image */
.output-card {
    background: var(--card-light);
    border-radius: 24px;
    padding: 24px;
    border: 1px solid var(--border-light);
    box-shadow: 0 1px 3px rgba(0, 0, 0, 0.05), 0 20px 40px rgba(0, 0, 0, 0.03);
    margin-bottom: 24px;
    min-height: 400px;
    display: flex;
    align-items: center;
    justify-content: center;
    transition: all 0.3s ease;
}

.dark .output-card {
    background: var(--card-dark);
    border-color: var(--border-dark);
    box-shadow: 0 1px 3px rgba(0, 0, 0, 0.3), 0 20px 40px rgba(0, 0, 0, 0.2);
}

.output-card img {
    border-radius: 16px !important;
    max-width: 100% !important;
    height: auto !important;
    display: block !important;
}

/* Empty state */
.output-card:empty::before {
    content: 'Your generated image will appear here';
    color: #94a3b8;
    font-size: 15px;
    font-weight: 500;
}

/* Footer */
.footer {
    text-align: center;
    padding: 32px 0 0;
}

.footer a {
    display: inline-flex;
    align-items: center;
    gap: 8px;
    color: var(--primary);
    text-decoration: none;
    font-size: 14px;
    font-weight: 500;
    transition: opacity 0.2s ease;
    padding: 8px 16px;
    border-radius: 8px;
    background: rgba(99, 102, 241, 0.05);
}

.dark .footer a {
    background: rgba(99, 102, 241, 0.1);
}

.footer a:hover {
    opacity: 0.8;
    background: rgba(99, 102, 241, 0.1);
}

.dark .footer a:hover {
    background: rgba(99, 102, 241, 0.15);
}

/* Progress bar */
.progress-bar-wrap {
    border-radius: 8px !important;
    overflow: hidden !important;
}

.progress-bar {
    background: linear-gradient(90deg, var(--primary), #8b5cf6) !important;
    border-radius: 8px !important;
}

/* Remove default Gradio padding */
.block {
    border: none !important;
    box-shadow: none !important;
    background: transparent !important;
}

.block.padded {
    padding: 0 !important;
}

/* Mobile optimizations */
@media (max-width: 640px) {
    .gradio-container {
        padding: 0 16px 32px !important;
    }
    
    .hero {
        padding: 32px 0 32px;
        margin-bottom: 24px;
    }
    
    .main-card {
        padding: 24px 20px;
        border-radius: 20px;
    }
    
    .output-card {
        padding: 20px;
        border-radius: 20px;
        min-height: 320px;
    }
    
    textarea {
        font-size: 16px !important;
        padding: 14px 16px !important;
    }
    
    button.primary {
        padding: 14px 20px !important;
        font-size: 15px !important;
    }
}

/* Smooth transitions */
* {
    -webkit-tap-highlight-color: transparent;
}

.gradio-container,
.main-card,
.output-card,
textarea,
button {
    transition: background-color 0.3s ease, border-color 0.3s ease, box-shadow 0.3s ease, transform 0.2s ease !important;
}

/* Loading state */
@keyframes pulse {
    0%, 100% { opacity: 1; }
    50% { opacity: 0.5; }
}

.generating {
    animation: pulse 2s cubic-bezier(0.4, 0, 0.6, 1) infinite;
}

/* Accessibility */
button:focus-visible,
textarea:focus-visible {
    outline: 2px solid var(--primary) !important;
    outline-offset: 2px !important;
}
"""

# Create the Gradio interface - Gradio 6 syntax: NO parameters in gr.Blocks()!
with gr.Blocks() as demo:
    
    # Hero Section
    gr.HTML("""
        <div class="hero">
            <h1 class="hero-title">OVIS Image Generator</h1>
            <p class="hero-subtitle">Transform your imagination into stunning visuals with AI</p>
        </div>
    """)
    
    # Main Input Card
    with gr.Column(elem_classes="main-card"):
        with gr.Column(elem_classes="input-section"):
            prompt = gr.Textbox(
                label="What would you like to create?",
                placeholder="A mystical forest with glowing mushrooms under moonlight, digital art style...",
                lines=4,
                max_lines=8,
                show_label=True
            )
        
        generate_btn = gr.Button(
            "Generate Image",
            variant="primary",
            size="lg"
        )
    
    # Output Card - Gradio 6: Use 'buttons' parameter instead of show_download_button/show_share_button
    with gr.Column(elem_classes="output-card"):
        output_image = gr.Image(
            label="Generated Image",
            type="pil",
            show_label=False,
            buttons=["download", "share"],  # Gradio 6 syntax
            container=False
        )
    
    # Footer
    gr.HTML("""
        <div class="footer">
            <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank">
                ✨ Built with anycoder
            </a>
        </div>
    """)
    
    # Event handlers - Gradio 6: use api_visibility parameter
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt],
        outputs=output_image,
        api_visibility="public"
    )
    
    prompt.submit(
        fn=generate_image,
        inputs=[prompt],
        outputs=output_image,
        api_visibility="public"
    )

# Gradio 6: ALL app parameters go in demo.launch()!
if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Soft(
            primary_hue="indigo",
            secondary_hue="purple",
            neutral_hue="slate",
            font=gr.themes.GoogleFont("Inter"),
            text_size="md",
            spacing_size="md",
            radius_size="lg"
        ),
        css=modern_css,
        footer_links=[
            {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}
        ]
    )