import gradio as gr
from src.utils.change_format import change_format
from src.utils.remove_background import remove_background
from src.utils.generate_image import generate_image
from src.utils.add_text import add_text_to_image_base64
from src.utils.compress import compress_image_memory
from src.utils.generate_image import generate_image
from src.utils.apply_filter import apply_filter_direct
from src.utils.watermark import add_watermark, remove_watermark
from src.utils.describe import describe_image
import base64
from PIL import Image
import io
import requests
from io import BytesIO
from typing import Union

def change_format(image: Union[str, BytesIO], target_format: str) -> str:
    """
    Change the format of an image from a URL to the specified target format.
    """
    
    if not isinstance(image, BytesIO):
        response = requests.get(image, timeout=30)
        response.raise_for_status()
        img = Image.open(BytesIO(response.content))
    else: 
        img = Image.open(image)
    
    output = BytesIO()
    img.save(output, format=target_format)
    output.seek(0)
    
    encoded_image = base64.b64encode(output.getvalue()).decode('utf-8')
    
    return encoded_image

def image_to_base64(image):
    if image is None:
        return None
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return base64.b64encode(buffer.getvalue()).decode()

def base64_to_image(base64_str):
    if not base64_str:
        return None
    
    # Remove data URI prefix if present (e.g., "data:image/png;base64,")
    if isinstance(base64_str, str) and "base64," in base64_str:
        base64_str = base64_str.split("base64,", 1)[1]
    
    try:
        # Strip any whitespace that might be in the base64 string
        if isinstance(base64_str, str):
            base64_str = base64_str.strip()
        
        # Decode the base64 data
        image_data = base64.b64decode(base64_str)
        
        # Check if we have data
        if not image_data:
            print("Decoded base64 data is empty")
            return None
        
        # Attempt to open the image
        image = Image.open(io.BytesIO(image_data))
        
        # Convert the image to ensure it's valid
        return image.copy()
        
    except base64.binascii.Error as e:
        print(f"Base64 decoding error: {str(e)}")
        if isinstance(base64_str, str):
            preview = base64_str[:30] + "..." if len(base64_str) > 30 else base64_str
            print(f"Base64 preview: {preview}")
        return None
        
    except Exception as e:
        print(f"Error converting base64 to image: {str(e)}")
        
        # Print preview of the base64 string for debugging
        if isinstance(base64_str, str):
            preview = base64_str[:30] + "..." if len(base64_str) > 30 else base64_str
            print(f"Base64 preview: {preview}")
        
        # Additional debug information
        if 'image_data' in locals() and image_data:
            try:
                magic_bytes = image_data[:12].hex()
                print(f"First 12 bytes: {magic_bytes}")
            except:
                pass
        
        return None
    
def url_to_base64(url):
    response = requests.get(url)
    return base64.b64encode(response.content).decode()

def gradio_remove_background(image):
    if image is None:
        return None
    base64_img = image_to_base64(image)
    result = remove_background(f"data:image/png;base64,{base64_img}")
    
    # Check if the result is directly a base64 string or has an image_data key
    if isinstance(result, str):
        return base64_to_image(result)
    elif isinstance(result, dict) and "image_data" in result:
        # If image_data contains a data URI prefix
        if isinstance(result["image_data"], str) and result["image_data"].startswith("data:"):
            # The response already contains the full data URI
            return base64_to_image(result["image_data"])
        else:
            # Try to process it as a regular base64 string
            try:
                return base64_to_image(result["image_data"])
            except Exception as e:
                print(f"Error processing image data: {e}")
                return None
    else:
        print(f"Unexpected response format from remove_background: {type(result)}")
        return None

def gradio_describe_image(image):
    if image is None:
        return "No image provided"
    try:
        base64_img = image_to_base64(image)
        return describe_image(base64_img)
    except Exception as e:
        print(f"Error describing image: {e}")
        return f"Error: {str(e)}"

def gradio_change_format(image, format_type):
    if image is None:
        return None
    try:
        base64_img = image_to_base64(image)
        result = change_format(base64_img, format_type)
        return base64_to_image(result)
    except Exception as e:
        print(f"Error changing format: {e}")
        return image

def gradio_generate_image(prompt, width=512, height=512):
    result = generate_image(prompt, width, height)
    return base64_to_image(result["b64"])

def gradio_apply_filter(image, filter_type, intensity=1.0):
    if image is None:
        print("No image provided")
        return None
    
    return apply_filter_direct(image, filter_type, intensity)

def update_text_image(image, text, centered, x, y, font_size, color):
    if image is None:
        return None
    if not text or text.strip() == "":
        return image
    
    result = add_text_to_image_base64(image, text, int(x), int(y), int(font_size), color, centered)
    return result

def toggle_position_fields(centered):
    return (
        gr.Number(interactive=not centered), 
        gr.Number(interactive=not centered)
    )

def toggle_intensity_slider(filter_type):
    intensity_filters = ['blur', 'brightness', 'contrast', 'saturation']
    return gr.Slider(interactive=filter_type in intensity_filters)

def gradio_add_watermark(image, watermark_text, opacity=0.5):
    if image is None:
        return None
    try:
        base64_img = image_to_base64(image)
        result = add_watermark(base64_img, watermark_text, opacity)
        return base64_to_image(result)
    except Exception as e:
        print(f"Error adding watermark: {e}")
        return image

def gradio_remove_watermark(image):
    if image is None:
        return None
    try:
        base64_img = image_to_base64(image)
        result = remove_watermark(base64_img)
        return base64_to_image(result)
    except Exception as e:
        print(f"Error removing watermark: {e}")
        return image

def gradio_compress_image(image, quality=80):
    """
    Compress image for Gradio interface
    """
    if image is None:
        return None
    try:
        compressed_image = compress_image_memory(image, quality, "JPEG")
        return compressed_image
    except Exception as e:
        print(f"Error compressing image: {e}")
        return image

def create_gradio_interface():
    with gr.Blocks(title="ImageUtilitiesMCP", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🖼️ ImageUtilitiesMCP")
        gr.Markdown("Complete processing image tools")
        
        with gr.Tabs():
            with gr.Tab("🎨 Generate Image"):
                with gr.Row():
                    prompt_input = gr.Textbox(label="Prompt", placeholder="Describe the image you want to generate")
                    with gr.Column():
                        width_input = gr.Slider(256, 1024, 512, label="Width")
                        height_input = gr.Slider(256, 1024, 512, label="Height")
                generate_btn = gr.Button("Generate", variant="primary")
                generated_output = gr.Image(label="Generated Image")
                
                generate_btn.click(
                    gradio_generate_image,
                    [prompt_input, width_input, height_input],
                    generated_output
                )
            
            with gr.Tab("🔍 Describe Image"):
                with gr.Row():
                    describe_input = gr.Image(label="Upload Image", type="pil")
                    description_output = gr.Textbox(label="Description", lines=4)
                
                describe_input.change(gradio_describe_image, describe_input, description_output)
            
            with gr.Tab("✂️ Remove Background"):
                with gr.Row():
                    bg_input = gr.Image(label="Upload Image", type="pil")
                    bg_output = gr.Image(label="Background Removed")
                
                bg_input.change(gradio_remove_background, bg_input, bg_output)
            
            with gr.Tab("🎭 Apply Filters"):
                with gr.Row():
                    filter_input = gr.Image(label="Upload Image", type="pil")
                    with gr.Column():
                        filter_type = gr.Dropdown(
                            ["blur", "sharpen", "vintage", "black_white", "sepia", "emboss", "edge", "smooth", "brightness", "contrast", "saturation", "grayscale"],
                            label="Filter Type",
                            value="blur"
                        )
                        
                        intensity_slider = gr.Slider(
                            minimum=0.1,
                            maximum=300.0,
                            value=1.0,
                            step=0.1,
                            label="Intensity",
                            interactive=True
                        )
                        
                        filter_output = gr.Image(label="Filtered Image")
                
                filter_type.change(
                    toggle_intensity_slider,
                    filter_type,
                    intensity_slider
                )
                
                filter_inputs = [filter_input, filter_type, intensity_slider]
                
                for inp in filter_inputs:
                    inp.change(gradio_apply_filter, filter_inputs, filter_output)
            
            with gr.Tab("📝 Add Text"):
                with gr.Row():
                    text_input = gr.Image(label="Upload Image", type="pil")
                    with gr.Column():
                        text_content = gr.Textbox(
                            label="Text", 
                            placeholder="Enter text to add",
                            value=""
                        )
                        text_centered = gr.Checkbox(label="Center Text", value=False)
                        
                        with gr.Row():
                            text_x = gr.Number(
                                label="X Position", 
                                value=50, 
                                interactive=True,
                                minimum=0
                            )
                            text_y = gr.Number(
                                label="Y Position", 
                                value=50, 
                                interactive=True,
                                minimum=0
                            )
                        
                        with gr.Row():
                            font_size = gr.Slider(
                                minimum=10, 
                                maximum=100, 
                                value=20, 
                                label="Font Size"
                            )
                            text_color = gr.ColorPicker(
                                label="Color", 
                                value="#FFFFFF"
                            )
                        
                        add_text_btn = gr.Button("Add Text", variant="primary")
                        text_output = gr.Image(label="Image with Text")
                
                text_centered.change(
                    toggle_position_fields,
                    text_centered,
                    [text_x, text_y]
                )
                
                inputs = [text_input, text_content, text_centered, text_x, text_y, font_size, text_color]
                
                add_text_btn.click(
                    update_text_image,
                    inputs,
                    text_output
                )
                
                for inp in inputs:
                    inp.change(update_text_image, inputs, text_output)
            
            with gr.Tab("💧 Watermark"):
                with gr.Tabs():
                    with gr.Tab("Add Watermark"):
                        with gr.Row():
                            watermark_input = gr.Image(label="Upload Image", type="pil")
                            with gr.Column():
                                watermark_text = gr.Textbox(label="Watermark Text")
                                watermark_opacity = gr.Slider(0.1, 1.0, 0.5, label="Opacity")
                                watermark_output = gr.Image(label="Watermarked Image")
                        
                        inputs = [watermark_input, watermark_text, watermark_opacity]
                        for inp in inputs:
                            inp.change(gradio_add_watermark, inputs, watermark_output)
                    
                    with gr.Tab("Remove Watermark"):
                        with gr.Row():
                            unwatermark_input = gr.Image(label="Upload Image", type="pil")
                            unwatermark_output = gr.Image(label="Watermark Removed")
                        
                        unwatermark_input.change(gradio_remove_watermark, unwatermark_input, unwatermark_output)
            
            with gr.Tab("🗜️ Compress"):
                with gr.Row():
                    compress_input = gr.Image(label="Upload Image", type="pil")
                    with gr.Column():
                        quality_slider = gr.Slider(0, 100, 80, label="Quality %")
                        compress_output = gr.Image(label="Compressed Image")
                
                compress_input.change(gradio_compress_image, [compress_input, quality_slider], compress_output)
                quality_slider.change(gradio_compress_image, [compress_input, quality_slider], compress_output)
            
            with gr.Tab("🔄 Change Format"):
                with gr.Row():
                    format_input = gr.Image(label="Upload Image", type="pil")
                    with gr.Column():
                        format_type = gr.Dropdown(
                            ["PNG", "JPEG", "WEBP", "BMP"],
                            label="Output Format",
                            value="PNG"
                        )
                        format_output = gr.Image(label="Converted Image")
                
                format_input.change(gradio_change_format, [format_input, format_type], format_output)
                format_type.change(gradio_change_format, [format_input, format_type], format_output)
        
        gr.Markdown("---")
        gr.Markdown("💡 **Status**: Active | Procesamiento de imágenes en tiempo real")
    
    return demo

if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch(
        mcp_server=True,
        server_name="0.0.0.0", 
        server_port=7860,
        show_error=True
    )