"""
Architecture AI Enhancer - Gradio App for Hugging Face Spaces

This file creates a Gradio interface for the enhancement pipeline.
"""
import gradio as gr
import torch
from PIL import Image
import logging
from pathlib import Path
import sys
import os

# Add backend to path
sys.path.insert(0, str(Path(__file__).parent))

from backend.services.diffusion_pipeline import get_pipeline_manager, enhance_image as enhance_with_diffusion
from backend.services.upscaler import upscale_image
from backend.services.postprocess import postprocess_image
from backend.config import settings

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize pipeline manager (lazy loading)
pipeline_manager = None

def initialize_models():
    """Initialize all models on first use"""
    global pipeline_manager
    
    if pipeline_manager is None:
        logger.info("Loading Stable Diffusion pipeline...")
        pipeline_manager = get_pipeline_manager()

def enhance_image_gradio(
    input_image: Image.Image,
    strength: float = 0.3,
    guidance_scale: float = 5.5,
    custom_prompt: str = "",
    use_upscaler: bool = True,
    use_postprocess: bool = True
) -> Image.Image:
    """
    Enhance architectural image using AI pipeline
    
    Args:
        input_image: Input PIL Image
        strength: Denoising strength (0.1-0.8)
        guidance_scale: CFG scale (1.0-15.0)
        custom_prompt: Optional custom prompt
        use_upscaler: Apply upscaling
        use_postprocess: Apply post-processing
        
    Returns:
        Enhanced PIL Image
    """
    try:
        # Initialize models
        logger.info("Loading models...")
        initialize_models()
        
        # Step 1: AI Enhancement
        logger.info("Preprocessing image...")
        
        prompt = custom_prompt if custom_prompt else settings.DEFAULT_PROMPT
        
        logger.info("Enhancing with AI (this may take a few minutes)...")
        
        enhanced = pipeline_manager.run_inference(
            image=input_image,
            prompt=prompt,
            negative_prompt=settings.NEGATIVE_PROMPT,
            strength=strength,
            guidance_scale=guidance_scale,
            num_inference_steps=30
        )
        
        logger.info("AI enhancement complete!")
        
        # Step 2: Optional Upscaling
        if use_upscaler:
            logger.info("Upscaling image...")
            enhanced = upscale_image(enhanced, scale=2)
            logger.info("Upscaling complete!")
        
        # Step 3: Optional Post-processing
        if use_postprocess:
            logger.info("Applying final touches...")
            enhanced = postprocess_image(
                enhanced,
                clahe=True,
                sharpen=1.0,
                color_enhance=1.1
            )
            logger.info("Post-processing complete!")
        
        logger.info("Enhancement process complete!")
        return enhanced
        
    except Exception as e:
        logger.error(f"Enhancement failed: {e}", exc_info=True)
        raise gr.Error(f"Enhancement failed: {str(e)}")

# Create Gradio interface
with gr.Blocks(title="Architecture AI Enhancer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🏗️ Architecture AI Enhancer
    
    Transform your architectural renders with AI-powered enhancement using Stable Diffusion 1.5.
    
    **Upload an image** and adjust the settings below to enhance your architectural visualization.
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="📤 Input Image",
                type="pil",
                height=400
            )
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                strength = gr.Slider(
                    minimum=0.1,
                    maximum=0.8,
                    value=0.3,
                    step=0.05,
                    label="Denoising Strength",
                    info="Lower = more faithful to input, Higher = more creative"
                )
                
                guidance_scale = gr.Slider(
                    minimum=1.0,
                    maximum=15.0,
                    value=5.5,
                    step=0.5,
                    label="Guidance Scale",
                    info="How closely to follow the prompt"
                )
                
                custom_prompt = gr.Textbox(
                    label="Custom Prompt (optional)",
                    placeholder="professional architectural photography, detailed, high quality...",
                    lines=3
                )
                
                use_upscaler = gr.Checkbox(
                    label="Enable Upscaling (2x)",
                    value=True
                )
                
                use_postprocess = gr.Checkbox(
                    label="Enable Post-Processing",
                    value=True,
                    info="Adds photographic enhancements"
                )
            
            enhance_btn = gr.Button("✨ Enhance Image", variant="primary", size="lg")
            
        with gr.Column():
            output_image = gr.Image(
                label="✅ Enhanced Result",
                type="pil",
                height=400
            )
            
            gr.Markdown("""
            ### 📝 Tips for Best Results:
            - Use high-quality architectural renders as input
            - Start with default settings and adjust if needed
            - Lower strength for subtle enhancements
            - Higher strength for more dramatic changes
            - Processing takes 2-5 minutes on CPU, ~30 seconds on GPU
            """)
    
    gr.Markdown("""
    ---
    ### 🔧 Technical Details:
    - **Model**: Stable Diffusion 1.5
    - **Upscaler**: ESRGAN (optional)
    - **Processing**: CPU/GPU automatic detection
    - **Version**: 1.0.0
    
    ### 📚 Resources:
    - [GitHub Repository](#)
    - [Documentation](#)
    - [Report Issues](#)
    """)
    
    # Connect the enhance button
    enhance_btn.click(
        fn=enhance_image_gradio,
        inputs=[
            input_image,
            strength,
            guidance_scale,
            custom_prompt,
            use_upscaler,
            use_postprocess
        ],
        outputs=output_image
    )

# Launch configuration
if __name__ == "__main__":
    demo.queue(max_size=10)  # Enable queue for multiple users
    
    # Detect if running on HF Spaces
    is_hf_space = os.getenv("SPACE_ID") is not None
    
    if is_hf_space:
        # HF Spaces configuration
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            show_api=False
        )
    else:
        # Local development
        demo.launch(share=True)