import gradio as gr
from transformers import pipeline
import torch
import numpy as np
from PIL import Image
import io

def remove_background(input_image):
    try:
        # Initialize the pipeline with trust_remote_code=True
        segmentor = pipeline(
            "image-segmentation", 
            model="briaai/RMBG-1.4",
            trust_remote_code=True,
            device="cpu"
        )
        
        # Process the image and get result
        result = segmentor(input_image, return_mask=True)
        
        # Convert result to RGBA
        if isinstance(result, Image.Image):
            # Create transparent background
            output = Image.new('RGBA', result.size, (0, 0, 0, 0))
            output.paste(input_image, mask=result)
        else:
            output = result['output_image']
        
        return output

    except Exception as e:
        raise gr.Error(f"Error processing image: {str(e)}")

# Custom theme and styling
theme = gr.themes.Soft(
    primary_hue="gold",
    secondary_hue="orange",
).set(
    body_background_fill="linear-gradient(135deg, #1a1a1a 0%, #2d2d2d 100%)",
    body_text_color="#ffffff",
    button_primary_background_fill="linear-gradient(45deg, #FFD700, #FFA500)",
    button_primary_text_color="#000000",
    border_color_primary="#FFD700"
)

css = """
.gradio-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
    padding: 20px !important;
}
.image-container {
    border-radius: 15px !important;
    border: 2px solid rgba(255, 215, 0, 0.3) !important;
    padding: 10px !important;
    background: rgba(255, 255, 255, 0.1) !important;
    transition: transform 0.3s ease !important;
}
.image-container:hover {
    transform: translateY(-5px) !important;
}
.gr-button {
    min-width: 200px !important;
    height: 45px !important;
    font-size: 16px !important;
    margin: 10px !important;
    transition: all 0.3s ease !important;
}
.gr-button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 5px 15px rgba(255, 215, 0, 0.3) !important;
}
.footer {
    text-align: center;
    margin-top: 20px;
    color: #666;
}
"""

# Create Gradio interface
with gr.Blocks(theme=theme, css=css) as demo:
    gr.HTML(
        """
        <div style="text-align: center; margin-bottom: 2rem;">
            <h1 style="font-size: 3rem; margin-bottom: 1rem; background: linear-gradient(45deg, #FFD700, #FFA500); -webkit-background-clip: text; -webkit-text-fill-color: transparent;">
                AI Background Remover Pro
            </h1>
            <p style="color: #cccccc; font-size: 1.2rem;">
                Remove backgrounds instantly using advanced AI technology
            </p>
        </div>
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="Upload Your Image",
                type="pil",
                elem_classes="image-container"
            )
            
            with gr.Row():
                clear_btn = gr.Button("Clear", variant="secondary")
                process_btn = gr.Button("Remove Background", variant="primary")
                download_btn = gr.Button("Download", variant="primary", visible=False)

        with gr.Column():
            output_image = gr.Image(
                label="Result",
                type="pil",
                elem_classes="image-container"
            )
            
    # Status message
    status_msg = gr.Textbox(
        label="Status",
        placeholder="Ready to process your image...",
        interactive=False
    )

    # Event handlers
    def process_and_update(image):
        if image is None:
            return None, "Please upload an image first", gr.Button.update(visible=False)
        try:
            result = remove_background(image)
            return (
                result, 
                "✨ Background removed successfully!", 
                gr.Button.update(visible=True)
            )
        except Exception as e:
            return None, f"❌ Error: {str(e)}", gr.Button.update(visible=False)

    process_btn.click(
        fn=process_and_update,
        inputs=[input_image],
        outputs=[output_image, status_msg, download_btn],
    )
    
    clear_btn.click(
        fn=lambda: (None, None, "Ready to process your image...", gr.Button.update(visible=False)),
        outputs=[input_image, output_image, status_msg, download_btn],
    )

    gr.HTML(
        """
        <div class="footer">
            <p>Powered by BRIA AI's RMBG V1.4 Model</p>
        </div>
        """
    )

# Launch the app
demo.launch()