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

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

# 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"]
}

# Main editing function
def edit_image(image, feature_type, modification_type, intensity, 
               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
            
        # 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
        edited_image = ledits_model.edit_image(
            processed_image, 
            mask, 
            prompt, 
            intensity=intensity
        )
        
        # Postprocess
        final_image = postprocess_image(edited_image, processed_image, mask)
        
        return final_image, "Edit completed successfully."
    
    except Exception as e:
        return image, f"Error during editing: {str(e)}"

# 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.")
        
        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"
                    )
                
                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=edit_image,
            inputs=[
                input_image, 
                feature_type, 
                modification_type, 
                intensity, 
                custom_prompt, 
                use_custom_prompt
            ],
            outputs=[output_image, status_text]
        )
        
        def reset_image():
            return None, "Image reset."
        
        reset_button.click(
            fn=reset_image,
            inputs=[],
            outputs=[output_image, status_text]
        )
        
        # Removed examples section to avoid file path errors
        
        # 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.
        """)
        
    return app

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