"""
Architecture AI Enhancer - Simple Gradio Interface for HF Spaces
"""
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
from backend.services.upscaler import upscale_image
from backend.services.postprocess import postprocess_image

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

# Initialize pipeline manager
pipeline_manager = None

def initialize_models():
    """Initialize models on first use"""
    global pipeline_manager
    if pipeline_manager is None:
        logger.info("Loading Stable Diffusion pipeline...")
        pipeline_manager = get_pipeline_manager()
        
        # Check if there's a LoRA model to load
        from backend.config import get_lora_path
        lora_path = get_lora_path()
        if lora_path:
            logger.info(f"LoRA model found: {lora_path}")
            try:
                pipeline_manager.load_lora(lora_path)
                logger.info("✅ Custom LoRA loaded successfully!")
            except Exception as e:
                logger.warning(f"Failed to load LoRA: {e}")

def enhance_image_simple(
    input_image: Image.Image,
    strength: float,
    guidance_scale: float,
    custom_prompt: str,
    use_upscaler: bool,
    use_postprocess: bool,
    progress=gr.Progress()
) -> Image.Image:
    """Enhance architectural image"""
    try:
        if input_image is None:
            raise ValueError("Please upload an image first")
            
        progress(0, desc="Starting enhancement...")
        
        # Initialize models
        try:
            progress(0.1, desc="Initializing AI models...")
            initialize_models()
        except Exception as e:
            logger.error(f"Model initialization failed: {e}", exc_info=True)
            raise ValueError(f"Failed to load AI models: {str(e)}")
        
        # Build prompt
        base_prompt = "professional architectural photography, highly detailed, 8k, photorealistic"
        if custom_prompt and custom_prompt.strip():
            final_prompt = f"{base_prompt}, {custom_prompt.strip()}"
        else:
            final_prompt = base_prompt
        
        # Step 1: AI Enhancement
        progress(0.3, desc="Applying AI enhancement (this may take several minutes on CPU)...")
        try:
            # Create a callback to update progress during diffusion
            def update_progress(percent, description):
                # Map internal progress (30-85%) to our progress range
                mapped_progress = 0.3 + (percent / 100) * 0.4  # 30% to 70%
                progress(mapped_progress, desc=description)
            
            enhanced = pipeline_manager.enhance(
                image=input_image,
                prompt=final_prompt,
                negative_prompt="blurry, low quality, distorted, ugly, bad architecture",
                strength=strength,
                guidance_scale=guidance_scale,
                num_inference_steps=30,
                progress_callback=update_progress
            )
        except Exception as e:
            logger.error(f"AI enhancement failed: {e}", exc_info=True)
            raise ValueError(f"AI enhancement failed: {str(e)}")
        
        # Step 2: Upscaling
        if use_upscaler:
            progress(0.75, desc="Upscaling image...")
            try:
                enhanced = upscale_image(enhanced, scale=2)
            except Exception as e:
                logger.warning(f"Upscaling failed, skipping: {e}")
        
        # Step 3: Post-processing
        if use_postprocess:
            progress(0.85, desc="Applying final touches...")
            try:
                enhanced = postprocess_image(
                    enhanced,
                    apply_contrast=True,
                    apply_sharpening=True,
                    apply_color_grading=True,
                    apply_grain=False,
                    apply_vignette_effect=False
                )
            except Exception as e:
                logger.warning(f"Post-processing failed, skipping: {e}")
        
        progress(1.0, desc="Complete!")
        return enhanced
        
    except ValueError as e:
        # Re-raise ValueError with user-friendly message
        raise gr.Error(str(e))
    except Exception as e:
        logger.error(f"Enhancement failed: {e}", exc_info=True)
        raise gr.Error(f"Enhancement failed: {str(e)}")

# Create interface with tabs
with gr.Blocks(title="🏛️ Architecture AI Enhancer") as demo:
    gr.Markdown("""
    # 🏛️ Architecture AI Enhancer
    Transform your architectural renders with AI-powered enhancement using Stable Diffusion 1.5
    """)
    
    # Check for custom LoRA
    from backend.config import get_lora_path
    lora_path = get_lora_path()
    if lora_path and lora_path.exists():
        gr.Markdown(f"""
        <div style="background: linear-gradient(90deg, #4CAF50 0%, #45a049 100%); padding: 10px; border-radius: 8px; margin: 10px 0;">
            <p style="color: white; margin: 0; font-weight: bold;">
                ✅ Custom LoRA Active: {lora_path.name}
            </p>
            <p style="color: rgba(255,255,255,0.9); margin: 5px 0 0 0; font-size: 0.9em;">
                All enhancements will use your custom trained style
            </p>
        </div>
        """)
    else:
        gr.Markdown("""
        <div style="background: #f0f0f0; padding: 10px; border-radius: 8px; margin: 10px 0;">
            <p style="color: #666; margin: 0;">
                ℹ️ Using base Stable Diffusion 1.5 model. <a href="#" style="color: #2196F3;">Train a custom LoRA</a> for personalized style.
            </p>
        </div>
        """)
    
    with gr.Tabs():
        # Tab 1: Enhancement
        with gr.Tab("✨ Enhance Image"):
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(label="📤 Input Image", type="pil")
                    
                    with gr.Accordion("⚙️ Advanced Settings", open=False):
                        strength = gr.Slider(
                            0.1, 0.8, value=0.3, step=0.05,
                            label="Denoising Strength",
                            info="Lower = more faithful to input, Higher = more creative"
                        )
                        guidance_scale = gr.Slider(
                            1.0, 15.0, value=5.5, step=0.5,
                            label="Guidance Scale",
                            info="How closely to follow the prompt"
                        )
                        custom_prompt = gr.Textbox(
                            label="Additional Prompt (Optional)",
                            placeholder="e.g., modern minimalist, glass facade, sunset lighting...",
                            lines=2,
                            info="Add custom details to enhance specific aspects"
                        )
                        use_upscaler = gr.Checkbox(label="Enable Upscaling (2x)", value=True)
                        use_postprocess = gr.Checkbox(label="Enable Post-Processing", value=True)
                    
                    enhance_btn = gr.Button("✨ Enhance Image", variant="primary", size="lg")
                
                with gr.Column():
                    output_image = gr.Image(label="✅ Enhanced Result", type="pil")
                    
                    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
                    """)
            
            # Connect the button
            enhance_btn.click(
                fn=enhance_image_simple,
                inputs=[input_image, strength, guidance_scale, custom_prompt, use_upscaler, use_postprocess],
                outputs=output_image
            )
        
        # Tab 2: Training Guide
        with gr.Tab("🎓 Custom Training"):
            gr.Markdown("""
            # 🎓 Train Your Own Custom LoRA Model
            
            Want to teach the AI your specific architectural style? You can train a custom LoRA model!
            
            ## ⚠️ Important Note
            
            **Training execution is not available on this HF Space** due to computational requirements. However, you can:
            1. Upload and preview your training pairs below
            2. Download the guide to train locally
            3. Deploy your custom model
            """)
            
            gr.Markdown("---")
            
            # Training pair upload section
            gr.Markdown("""
            ## 📸 Preview Training Pairs
            
            Upload example pairs to visualize what you'll train with:
            """)
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Input Image (Before)")
                    training_input = gr.Image(label="Input Render", type="pil")
                
                with gr.Column():
                    gr.Markdown("### Target Image (After)")
                    training_target = gr.Image(label="Target Enhanced", type="pil")
            
            with gr.Row():
                preview_btn = gr.Button("👁️ Preview Training Pair", variant="secondary")
                clear_btn = gr.Button("🗑️ Clear", variant="secondary")
            
            training_preview = gr.Gallery(
                label="Training Pair Preview",
                columns=2,
                rows=1,
                height="auto"
            )
            
            def preview_training_pair(input_img, target_img):
                """Preview training pair side by side"""
                if input_img is None or target_img is None:
                    return []
                return [input_img, target_img]
            
            def clear_training():
                """Clear training inputs"""
                return None, None, []
            
            preview_btn.click(
                fn=preview_training_pair,
                inputs=[training_input, training_target],
                outputs=training_preview
            )
            
            clear_btn.click(
                fn=clear_training,
                inputs=[],
                outputs=[training_input, training_target, training_preview]
            )
            
            gr.Markdown("""
            ---
            
            ## 💡 Tips for Good Training Pairs
            
            - **Consistency**: Use similar compositions between input and target
            - **Quality**: High-resolution images (512x512 minimum)
            - **Variety**: Include different views and lighting conditions
            - **Alignment**: Input and target should show the same scene
            - **Target Quality**: Ensure targets represent your desired style accurately
            
            ---
            
            ## 🎯 Using a Custom LoRA Model
            
            Once you've trained a LoRA model locally, here's how to use it:
            
            ### Option 1: Deploy to This Space (Recommended)
            
            1. **Upload your LoRA to Hugging Face Hub**:
               ```bash
               # Install huggingface_hub
               pip install huggingface_hub
               
               # Upload your LoRA
               from huggingface_hub import upload_file
               upload_file(
                   path_or_fileobj="models/lora/your_lora.safetensors",
                   path_in_repo="your_lora.safetensors",
                   repo_id="your-username/your-lora-repo",
                   repo_type="model"
               )
               ```
            
            2. **Modify this Space to load your LoRA**:
               - Fork this Space or create a duplicate
               - Edit `backend/config.py`:
                 ```python
                 LORA_MODEL_NAME = "your_lora.safetensors"
                 ```
               - Add code in `app.py` to download from HF Hub:
                 ```python
                 from huggingface_hub import hf_hub_download
                 
                 lora_path = hf_hub_download(
                     repo_id="your-username/your-lora-repo",
                     filename="your_lora.safetensors",
                     local_dir="models/lora"
                 )
                 ```
            
            3. **Restart the Space** - Your custom LoRA will be loaded automatically!
            
            ### Option 2: Use Locally
            
            1. **Place your LoRA file** in `models/lora/` directory
            2. **Update config**: Set `LORA_MODEL_NAME` in `backend/config.py`
            3. **Run backend**: `uvicorn main:app --reload`
            4. The LoRA is automatically detected and loaded!
            
            ### How It Works
            
            When a LoRA is present:
            - ✅ Pipeline automatically loads it on startup
            - ✅ All enhancements use your custom style
            - ✅ No additional configuration needed
            - ✅ Can be combined with custom prompts for fine control
            
            ### Verify LoRA is Loaded
            
            Check the logs on startup:
            ```
            INFO: Loading Stable Diffusion pipeline...
            INFO: LoRA model found: models/lora/your_lora.safetensors
            INFO: ✅ Custom LoRA loaded successfully!
            ```
            
            ### Example: Architecture Firm Custom Style
            
            ```python
            # After training with your firm's rendering style
            # Place: models/lora/firm_style.safetensors
            
            # In config.py:
            LORA_MODEL_NAME = "firm_style.safetensors"
            
            # Now all enhancements will match your style!
            ```
            
            ---
            """)
            
            gr.Markdown("""
            ## 🚀 What is LoRA Training?
            
            LoRA (Low-Rank Adaptation) allows you to fine-tune the AI with just 10-50 image pairs to learn:
            - Your specific architectural rendering style
            - Preferred lighting and atmosphere
            - Consistent material treatments
            - Unique design aesthetics
            
            ### Benefits:
            - ✅ **Fast Training**: Only 1000 steps needed (~20-30 min on GPU)
            - ✅ **Small Models**: LoRA weights are only ~10-50 MB
            - ✅ **Style Consistency**: Perfect for architectural firms with specific styles
            - ✅ **Efficient**: Works on consumer GPUs
            
            ---
            
            ## 📋 Training Process Overview
            
            ### Step 1: Prepare Training Data
            
            Create **image pairs**:
            - **Input**: Your base architectural render (before)
            - **Target**: Your ideal enhanced result (after)
            
            **Requirements:**
            - Minimum: 10 pairs (recommended: 20-50)
            - Format: PNG or JPG
            - Resolution: 512x512 to 1024x1024
            
            **Example structure:**
            ```
            training_data/
              inputs/
                building_001_input.png
                building_002_input.png
              targets/
                building_001_target.png
                building_002_target.png
            ```
            
            ### Step 2: Setup Local Backend
            
            ```bash
            cd architecture-ai-enhancer/backend
            pip install -r requirements.txt
            uvicorn main:app --host 0.0.0.0 --port 8000
            ```
            
            ### Step 3: Upload Training Pairs
            
            Use the API at `http://localhost:8000/docs` to upload your image pairs:
            
            ```python
            import requests
            
            files = {
                'input_image': open('building_001_input.png', 'rb'),
                'target_image': open('building_001_target.png', 'rb')
            }
            response = requests.post(
                'http://localhost:8000/training/upload_pair',
                files=files
            )
            ```
            
            ### Step 4: Start Training
            
            ```python
            config = {
                "train_steps": 1000,
                "learning_rate": 1e-4,
                "lora_rank": 8,
                "batch_size": 1
            }
            
            response = requests.post(
                'http://localhost:8000/training/start',
                json=config
            )
            ```
            
            ### Step 5: Use Your Custom Model
            
            Once trained, your custom LoRA is automatically used for all enhancements!
            
            ---
            
            ## 🎨 Training Tips
            
            **For Subtle Enhancements:**
            - train_steps: 500
            - learning_rate: 5e-5
            - lora_rank: 4
            
            **For Dramatic Style Changes:**
            - train_steps: 1500
            - learning_rate: 1e-4
            - lora_rank: 12
            
            **Balanced (Recommended):**
            - train_steps: 1000
            - learning_rate: 1e-4
            - lora_rank: 8
            
            ---
            
            ## 📚 Complete Documentation
            
            For detailed step-by-step instructions, troubleshooting, and advanced techniques, see:
            
            **[📖 Complete LoRA Training Guide](https://huggingface.co/spaces/TransformacionDigitalAA/ArchEnhancer/blob/main/TRAINING_GUIDE.md)**
            
            This includes:
            - Detailed API usage examples
            - Complete Python training script
            - Troubleshooting common issues
            - Best practices for creating training data
            - How to deploy your custom LoRA
            
            ---
            
            ## 💡 Use Cases
            
            - **Architecture Firms**: Train on your signature rendering style
            - **Game Studios**: Consistent environmental concept art
            - **VFX Artists**: Specific lighting and atmosphere
            - **Real Estate**: Standardized visualization style
            
            ---
            
            ## 🔗 Resources
            
            - [Backend GitHub Repository](https://github.com/yourusername/architecture-ai-enhancer)
            - [LoRA Paper](https://arxiv.org/abs/2106.09685)
            - [Diffusers Documentation](https://huggingface.co/docs/diffusers/)
            """)
        
        # Tab 3: About
        with gr.Tab("ℹ️ About"):
            gr.Markdown("""
            # About Architecture AI Enhancer
            
            ## 🔧 Technical Details
            
            - **Model**: Stable Diffusion 1.5 (runwayml/stable-diffusion-v1-5)
            - **Upscaler**: ESRGAN with fallback to Lanczos
            - **Framework**: PyTorch + Diffusers
            - **Interface**: Gradio 4.20.0
            - **Version**: 1.0.0
            
            ## ⚙️ Settings Guide
            
            ### Denoising Strength (0.1-0.8)
            Controls how much the AI modifies your input image:
            - **0.2-0.3**: Subtle enhancements (recommended for most cases)
            - **0.4-0.5**: Moderate changes
            - **0.6-0.8**: Dramatic transformations
            
            ### Guidance Scale (1-15)
            How closely the AI follows the prompt:
            - **4-6**: Natural, balanced results (recommended)
            - **7-10**: More stylized output
            - **11-15**: Very strong prompt adherence
            
            ### Custom Prompt
            Add specific details to guide the enhancement:
            - Lighting: "sunset lighting", "dramatic shadows"
            - Style: "modern minimalist", "brutalist concrete"
            - Materials: "glass facade", "wooden accents"
            - Atmosphere: "foggy morning", "golden hour"
            
            ## 🚀 Features
            
            - ✅ AI-powered enhancement with Stable Diffusion
            - ✅ 2x image upscaling
            - ✅ Professional post-processing
            - ✅ Custom prompt support
            - ✅ Adjustable parameters
            - ✅ Custom LoRA training (local)
            
            ## 📊 Performance
            
            - **CPU (HF Spaces Free Tier)**: 2-5 minutes per image
            - **GPU (Local/Paid)**: ~30 seconds per image
            
            ## 🙏 Credits
            
            Built with:
            - [Stable Diffusion](https://github.com/CompVis/stable-diffusion)
            - [Diffusers](https://github.com/huggingface/diffusers)
            - [Gradio](https://gradio.app)
            - [PyTorch](https://pytorch.org)
            
            ## 📝 License
            
            MIT License - Free for commercial and personal use
            
            ---
            
            **Made with ❤️ for the architecture community**
            """)

if __name__ == "__main__":
    # Check if running on HF Spaces
    is_spaces = os.getenv("SPACE_ID") is not None
    
    if is_spaces:
        # HF Spaces specific configuration
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False
        )
    else:
        # Local development
        demo.launch(share=True)