"""
Omni Image Editor - Redesigned
A modern, mobile-first Gradio 6 interface for image editing and generation.
"""

import gradio as gr
from pipeline import OmniMMDitV2Pipeline
from nfsw_filter import NSFWDetector

# Initialize Pipeline and Safety Filter
# Note: Ensure you have the required models downloaded or accessible
try:
    pipeline = OmniMMDitV2Pipeline(
        model=None, # Placeholder: Load actual model weights here
        vae=None,   # Placeholder: Load actual VAE here
        text_encoder=None, # Placeholder
        tokenizer=None, # Placeholder
        scheduler=None, # Placeholder
    )
    safety_checker = NSFWDetector()
except Exception as e:
    print(f"Warning: Models could not be initialized: {e}")
    # Create dummy instances for UI demonstration if models fail to load
    pipeline = None
    safety_checker = None

def check_safety(image):
    """Run NSFW check on the image."""
    if safety_checker is None:
        return image, "Safety check unavailable."
    
    try:
        # Assuming PIL image input
        label, _ = safety_checker.predict_from_pil(image)
        if label == "nsfw":
            return None, "⚠️ NSFW content detected. Image blocked."
        return image, "✅ Image is safe."
    except Exception as e:
        return image, f"Error checking safety: {e}"

def edit_image(input_image, prompt, negative_prompt, guidance_scale, num_steps):
    """Process image editing request."""
    if input_image is None:
        return None, "Please upload an image first."
    
    if pipeline is None:
        return None, "Model not loaded. Check server logs."

    # 1. Safety Check
    safe_image, safety_msg = check_safety(input_image)
    if safe_image is None:
        return None, safety_msg

    # 2. Run Inference
    try:
        # Note: This is a mockup of the pipeline call. 
        # Adjust parameters based on actual OmniMMDitV2Pipeline implementation
        result = pipeline(
            prompt=prompt,
            input_images=[input_image], # Image-to-Image
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=int(num_steps),
            height=1024,
            width=1024
        )
        
        # Extract image from result (assuming pipeline returns dict or object)
        if hasattr(result, 'images'):
            output_image = result.images[0]
        else:
            output_image = result[0]
            
        return output_image, f"Success! {safety_msg}"
        
    except Exception as e:
        return None, f"Error during generation: {str(e)}"

# --- Gradio 6 Interface ---

# Custom Theme: Modern, Soft, Mobile-Friendly
custom_theme = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="cyan",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
    text_size="lg",
    spacing_size="lg",
    radius_size="lg"
).set(
    # Button styling
    button_primary_background_fill="*primary_500",
    button_primary_background_fill_hover="*primary_600",
    button_primary_text_color="white",
    button_secondary_background_fill="*neutral_100",
    button_secondary_background_fill_hover="*neutral_200",
    # Container styling
    block_background_fill="*background_fill_primary",
    block_border_width="0px",
    block_shadow="*shadow_drop_lg",
    # Typography
    block_title_text_weight="600",
    block_title_text_size="*text_xl",
)

with gr.Blocks() as demo:
    
    # Header / Navbar
    gr.HTML("""
    <div style="display: flex; align-items: center; justify-content: space-between; padding: 1rem 0;">
        <div>
            <h1 style="margin: 0; font-size: 1.5rem; font-weight: 700; color: #4F46E5;">Omni Editor</h1>
            <p style="margin: 0; color: #64748B;">Next-Gen Image Editing</p>
        </div>
        <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="text-decoration: none; color: #4F46E5; font-weight: 500; font-size: 0.9rem;">
            Built with anycoder ↗
        </a>
    </div>
    """)

    with gr.Row(equal_height=True):
        # --- Sidebar (Controls) ---
        with gr.Column(scale=1, min_width=300):
            gr.Markdown("### 🎨 Controls")
            
            with gr.Accordion("⚙️ Generation Settings", open=True):
                prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Describe the changes you want...",
                    lines=3,
                    show_label=False
                )
                negative_prompt = gr.Textbox(
                    label="Negative Prompt",
                    placeholder="Things to avoid...",
                    lines=2,
                    show_label=False
                )
                
                with gr.Row():
                    guidance = gr.Slider(
                        minimum=1.0,
                        maximum=20.0,
                        value=7.5,
                        step=0.5,
                        label="Guidance Scale"
                    )
                    steps = gr.Slider(
                        minimum=10,
                        maximum=50,
                        value=25,
                        step=1,
                        label="Steps"
                    )
            
            generate_btn = gr.Button("✨ Generate Edit", variant="primary", size="lg")
            status_msg = gr.Textbox(label="Status", interactive=False, visible=False)

            gr.Markdown("""
            ---
            ### 💡 Tips
            - Upload a clear image for best results.
            - Be descriptive in your prompt.
            - Use negative prompts to remove artifacts.
            """)

        # --- Main Content (Canvas) ---
        with gr.Column(scale=2):
            gr.Markdown("### 🖼️ Workspace")
            
            with gr.Row():
                input_image = gr.Image(
                    label="Input",
                    type="pil",
                    sources=["upload", "clipboard", "webcam"],
                    interactive=True
                )
                output_image = gr.Image(
                    label="Result",
                    type="pil",
                    interactive=False
                )

    # --- Event Listeners ---
    generate_btn.click(
        fn=edit_image,
        inputs=[input_image, prompt, negative_prompt, guidance, steps],
        outputs=[output_image, status_msg],
        api_visibility="public"
    ).then(
        # Hide status after success if desired, or keep it
        lambda: gr.Textbox(visible=True),
        outputs=[status_msg]
    )

# --- Launch Configuration ---
demo.launch(
    theme=custom_theme,
    footer_links=[{"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}]
)