import os
import gradio as gr
import torch
import requests
from PIL import Image
import numpy as np
import io
import json
from models.ledits_model import LEDITSModel
from utils.image_processing import preprocess_image, postprocess_image
from utils.feature_detection import detect_features, create_mask

# Global variables
FEATURE_TYPES = ["Eyes", "Nose", "Lips", "Face Shape", "Hair", "Body"]
MODIFICATION_PRESETS = {
    "Eyes": ["Larger", "Smaller", "Change Color", "Change Shape"],
    "Nose": ["Refine", "Reshape", "Resize"],
    "Lips": ["Fuller", "Thinner", "Change Color"],
    "Face Shape": ["Slim", "Round", "Define Jawline", "Soften Features"],
    "Hair": ["Change Color", "Change Style", "Add Volume"],
    "Body": ["Slim", "Athletic", "Curvy", "Muscular"]
}

# Initialize models
def initialize_models():
    ledits_model = LEDITSModel()
    return ledits_model

# Local CPU processing function
def edit_image_cpu(image, feature_type, modification_type, intensity, 
                  num_inference_steps, guidance_scale, resolution,
                  custom_prompt="", use_custom_prompt=False):
    if image is None:
        return None, "Please upload an image first."
    
    try:
        # Convert to numpy array if needed
        if isinstance(image, Image.Image):
            image_np = np.array(image)
        else:
            image_np = image
        
        # Resize image based on resolution setting
        if resolution != "Original":
            max_dim = int(resolution.split("x")[0])
            height, width = image_np.shape[:2]
            
            if height > width:
                new_height = min(max_dim, height)
                new_width = int(width * (new_height / height))
            else:
                new_width = min(max_dim, width)
                new_height = int(height * (new_width / width))
            
            image_np = Image.fromarray(image_np).resize((new_width, new_height), Image.LANCZOS)
            image_np = np.array(image_np)
            
        # Preprocess image
        processed_image = preprocess_image(image_np)
        
        # Detect features and create mask
        features = detect_features(processed_image)
        mask = create_mask(processed_image, feature_type, features)
        
        # Get model
        ledits_model = initialize_models()
        
        # Prepare prompt
        if use_custom_prompt and custom_prompt:
            prompt = custom_prompt
        else:
            prompt = f"{feature_type} {modification_type}"
        
        # Apply edit with custom parameters
        edited_image = ledits_model.edit_image(
            processed_image, 
            mask, 
            prompt, 
            intensity=intensity,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps
        )
        
        # Postprocess
        final_image = postprocess_image(edited_image, processed_image, mask)
        
        return final_image, "Edit completed successfully."
    
    except Exception as e:
        import traceback
        traceback.print_exc()
        return image, f"Error during editing: {str(e)}"

# Remote GPU processing function
def process_with_local_server(image, feature_type, modification_type, intensity, 
                             num_inference_steps, guidance_scale, resolution,
                             custom_prompt="", use_custom_prompt=False, server_url=None):
    if image is None:
        return None, "Please upload an image first."
    
    if server_url is None or server_url == "":
        return image, "Error: Local server URL not provided. Please enter your local server URL."
    
    try:
        # Ensure server URL ends with /api/predict/
        if not server_url.endswith("/"):
            server_url += "/"
        if not server_url.endswith("api/predict/"):
            server_url += "api/predict/"
        
        # Convert image to bytes
        if isinstance(image, np.ndarray):
            image_pil = Image.fromarray(image)
        else:
            image_pil = image
            
        img_byte_arr = io.BytesIO()
        image_pil.save(img_byte_arr, format='PNG')
        img_byte_arr.seek(0)
        
        # Prepare the request data
        files = {
            'input_image': ('image.png', img_byte_arr, 'image/png')
        }
        
        data = {
            'feature_type': feature_type,
            'modification_type': modification_type,
            'intensity': str(intensity),
            'num_inference_steps': str(num_inference_steps),
            'guidance_scale': str(guidance_scale),
            'resolution': resolution,
            'custom_prompt': custom_prompt,
            'use_custom_prompt': str(use_custom_prompt).lower()
        }
        
        # Send request to local server
        response = requests.post(server_url, files=files, data=data)
        
        if response.status_code == 200:
            # Parse the response
            result = response.json()
            
            # Get the output image
            if 'data' in result and len(result['data']) >= 1:
                output_data = result['data'][0]
                if isinstance(output_data, str) and output_data.startswith('data:image'):
                    # Handle base64 encoded image
                    import base64
                    image_data = output_data.split(',')[1]
                    decoded_image = base64.b64decode(image_data)
                    output_image = Image.open(io.BytesIO(decoded_image))
                    return output_image, "Edit completed successfully."
            
            # If we couldn't parse the image from the response
            return image, f"Error: Could not parse response from local server."
        else:
            return image, f"Error: Local server returned status code {response.status_code}."
    
    except Exception as e:
        return image, f"Error connecting to local server: {str(e)}"

# Combined processing function that chooses between CPU and GPU
def process_image(image, feature_type, modification_type, intensity, 
                 num_inference_steps, guidance_scale, resolution,
                 custom_prompt, use_custom_prompt, processing_mode, server_url):
    if processing_mode == "Local CPU":
        return edit_image_cpu(
            image, feature_type, modification_type, intensity,
            num_inference_steps, guidance_scale, resolution,
            custom_prompt, use_custom_prompt
        )
    else:  # "Remote GPU"
        return process_with_local_server(
            image, feature_type, modification_type, intensity,
            num_inference_steps, guidance_scale, resolution,
            custom_prompt, use_custom_prompt, server_url
        )

# UI Components
def create_ui():
    with gr.Blocks(title="AI-Powered Facial & Body Feature Editor") as app:
        gr.Markdown("# AI-Powered Facial & Body Feature Editor")
        gr.Markdown("Upload an image and use the controls to edit specific facial and body features.")
        
        # Processing mode selection
        with gr.Group():
            gr.Markdown("### Processing Mode")
            processing_mode = gr.Radio(
                choices=["Local CPU", "Remote GPU"], 
                label="Select Processing Mode", 
                value="Local CPU",
                info="Choose 'Local CPU' to use Hugging Face's CPU or 'Remote GPU' to use your local GPU server"
            )
            
            # Server connection (only visible in Remote GPU mode)
            with gr.Group(visible=False) as server_group:
                gr.Markdown("### Local GPU Server Connection")
                server_url = gr.Textbox(
                    label="Local Server URL", 
                    placeholder="Enter the URL of your local GPU server (e.g., https://12345.gradio.app)",
                    value=""
                )
                server_status = gr.Textbox(label="Server Status", value="Not connected", interactive=False)
                
                def check_server(url):
                    if not url:
                        return "Not connected"
                    
                    try:
                        # Ensure URL ends with /
                        if not url.endswith("/"):
                            url += "/"
                        
                        # Try to connect to the server
                        response = requests.get(url)
                        if response.status_code == 200:
                            return "Connected successfully"
                        else:
                            return f"Error: Server returned status code {response.status_code}"
                    except Exception as e:
                        return f"Error connecting to server: {str(e)}"
                
                check_button = gr.Button("Check Connection")
                check_button.click(fn=check_server, inputs=server_url, outputs=server_status)
            
            # Show/hide server connection based on processing mode
            def toggle_server_group(mode):
                return gr.Group(visible=(mode == "Remote GPU"))
            
            processing_mode.change(
                fn=toggle_server_group,
                inputs=processing_mode,
                outputs=server_group
            )
        
        with gr.Row():
            with gr.Column(scale=1):
                # Input controls
                input_image = gr.Image(label="Upload Image", type="pil")
                
                with gr.Group():
                    gr.Markdown("### Feature Selection")
                    feature_type = gr.Dropdown(
                        choices=FEATURE_TYPES, 
                        label="Select Feature", 
                        value="Eyes"
                    )
                    
                    # Initialize with choices for the default feature (Eyes)
                    modification_type = gr.Dropdown(
                        choices=MODIFICATION_PRESETS["Eyes"], 
                        label="Modification Type", 
                        value="Larger"
                    )
                    
                    intensity = gr.Slider(
                        minimum=0.1, 
                        maximum=1.0, 
                        value=0.5, 
                        step=0.1, 
                        label="Intensity"
                    )
                
                with gr.Group():
                    gr.Markdown("### Custom Prompt (Advanced)")
                    use_custom_prompt = gr.Checkbox(
                        label="Use Custom Prompt", 
                        value=False
                    )
                    custom_prompt = gr.Textbox(
                        label="Custom Prompt", 
                        placeholder="e.g., blue eyes with long eyelashes"
                    )
                
                with gr.Group():
                    gr.Markdown("### Performance Settings")
                    num_inference_steps = gr.Slider(
                        minimum=5,
                        maximum=50,
                        value=20,
                        step=1,
                        label="Inference Steps (lower = faster, higher = better quality)"
                    )
                    guidance_scale = gr.Slider(
                        minimum=1.0,
                        maximum=15.0,
                        value=7.5,
                        step=0.5,
                        label="Guidance Scale (lower = more creative, higher = more accurate)"
                    )
                    resolution = gr.Dropdown(
                        choices=["Original", "512x512", "768x768", "1024x1024"],
                        label="Processing Resolution",
                        value="512x512"
                    )
                
                edit_button = gr.Button("Apply Edit", variant="primary")
                reset_button = gr.Button("Reset")
                status_text = gr.Textbox(label="Status", interactive=False)
            
            with gr.Column(scale=1):
                # Output display
                output_image = gr.Image(label="Edited Image", type="pil")
                
                with gr.Accordion("Edit History", open=False):
                    edit_history = gr.State([])
                    history_gallery = gr.Gallery(label="Previous Edits")
        
        # Event handlers
        def update_modification_choices(feature):
            return gr.Dropdown(choices=MODIFICATION_PRESETS[feature])
        
        feature_type.change(
            fn=update_modification_choices,
            inputs=feature_type,
            outputs=modification_type
        )
        
        edit_button.click(
            fn=process_image,
            inputs=[
                input_image, 
                feature_type, 
                modification_type, 
                intensity,
                num_inference_steps,
                guidance_scale,
                resolution,
                custom_prompt, 
                use_custom_prompt,
                processing_mode,
                server_url
            ],
            outputs=[output_image, status_text]
        )
        
        def reset_image():
            return None, "Image reset."
        
        reset_button.click(
            fn=reset_image,
            inputs=[],
            outputs=[output_image, status_text]
        )
        
        # Add ethical usage notice
        gr.Markdown("""
        ## Ethical Usage Notice
        
        This tool is designed for creative and personal use. Please ensure:
        
        - You have appropriate rights to edit the images you upload
        - You use this tool responsibly and respect the dignity of individuals
        - You understand that AI-generated modifications are artificial and may not represent reality
        
        By using this application, you agree to these terms.
        """)
        
        # Add dual-mode instructions
        gr.Markdown("""
        ## Processing Modes
        
        This application supports two processing modes:
        
        ### 1. Local CPU Mode
        - Uses Hugging Face's CPU for processing
        - Works immediately without additional setup
        - Slower processing (may take several minutes per edit)
        - No additional software required
        
        ### 2. Remote GPU Mode
        - Uses your local computer's GPU for processing
        - Requires setting up the local server component
        - Much faster processing (typically seconds per edit)
        - Requires PyTorch and other dependencies installed locally
        
        To use Remote GPU Mode:
        1. Download and run the local_server.py file on your computer
        2. The server will provide a URL (copy this URL)
        3. Select "Remote GPU" mode above
        4. Paste the URL into the "Local Server URL" field
        5. Click "Check Connection" to verify
        """)
        
    return app

# Launch the app
if __name__ == "__main__":
    app = create_ui()
    app.launch(server_name="0.0.0.0", share=False)