# Import libraries
import gradio as gr
from rembg import remove
from PIL import Image
import io

# Background remover function with error handling
def remove_background(image: Image.Image) -> Image.Image:
    try:
        # Convert image to bytes
        with io.BytesIO() as buffered:
            image.save(buffered, format="PNG", optimize=True)
            img_bytes = buffered.getvalue()
        
        # Process image with error handling
        output_bytes = remove(
            img_bytes,
            alpha_matting=True,
            alpha_matting_foreground_threshold=240,
            alpha_matting_background_threshold=10,
            alpha_matting_erode_size=10
        )
        
        output_image = Image.open(io.BytesIO(output_bytes)).convert("RGBA")
        return output_image
    
    except Exception as e:
        print(f"Error processing image: {str(e)}")
        raise gr.Error("Failed to process image. Please try another image.")

# Example Images
examples = [
    ["https://upload.wikimedia.org/wikipedia/commons/thumb/1/15/Cat_August_2010-4.jpg/480px-Cat_August_2010-4.jpg"],
    ["https://upload.wikimedia.org/wikipedia/commons/thumb/4/4d/Cat_November_2010-1a.jpg/480px-Cat_November_2010-1a.jpg"],
    ["https://upload.wikimedia.org/wikipedia/commons/thumb/7/75/Internet1.jpg/640px-Internet1.jpg"]
]

# Custom CSS for professional look
custom_css = """
:root {
    --primary: #4f46e5;
    --secondary: #f9fafb;
}
body {
    font-family: 'Inter', sans-serif;
}
.gradio-container {
    max-width: 900px !important;
    margin: auto !important;
}
.dark {
    --primary: #6366f1;
    --secondary: #1e293b;
}
footer {
    display: none !important;
}
.processing-message {
    color: var(--primary);
    font-weight: bold;
}
"""

# Social Media & Credits
socials = """
<div class="social-container">
    <h3>About AI ClearCut</h3>
    <p>Professional background removal tool powered by AI computer vision.</p>
    
    <div class="credits">
        <h4>Technical Details</h4>
        <p>Powered by <a href="https://github.com/danielgatis/rembg" target="_blank">Rembg</a> • UI by Gradio</p>
        <p>Model: u2net • Processing: CPU optimized</p>
    </div>
    
    <div class="developer">
        <h4>Developer</h4>
        <div class="links">
            <a href="https://github.com/prokaiiddo" target="_blank">GitHub</a> •
            <a href="https://linkedin.com/in/kaiiddo" target="_blank">LinkedIn</a> •
            <a href="https://twitter.com/kaiiddo" target="_blank">Twitter</a>
        </div>
    </div>
</div>
"""

# Gradio Interface
with gr.Blocks(title="AI ClearCut - Professional Background Remover", css=custom_css) as interface:
    gr.Markdown("""
    # <center>AI ClearCut</center>
    ### <center>Professional Background Removal Tool</center>
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Upload Image", type="pil", elem_id="input-image")
            with gr.Row():
                submit_btn = gr.Button("Remove Background", variant="primary")
                clear_btn = gr.Button("Clear", variant="secondary")
            gr.Examples(
                examples=examples,
                inputs=[input_image],
                label="Example Images",
                examples_per_page=3
            )
        
        with gr.Column():
            output_image = gr.Image(label="Result with Transparent Background", type="pil", elem_id="output-image")
            with gr.Row():
                download_btn = gr.Button("Download PNG")
                download_jpg = gr.Button("Download JPG (with white bg)")
    
    with gr.Accordion("Advanced Options", open=False):
        with gr.Row():
            alpha_matting = gr.Checkbox(label="Enable Alpha Matting (better for hair/fur)", value=True)
            with gr.Column(visible=False) as adv_options:
                foreground_thresh = gr.Slider(0, 255, value=240, label="Foreground Threshold")
                background_thresh = gr.Slider(0, 255, value=10, label="Background Threshold")
                erode_size = gr.Slider(0, 40, value=10, label="Erode Size")
    
    with gr.Accordion("About & Credits", open=False):
        gr.Markdown(socials)
    
    # Show/hide advanced options
    alpha_matting.change(
        fn=lambda x: gr.Column(visible=x),
        inputs=[alpha_matting],
        outputs=[adv_options]
    )
    
    # Process image
    submit_btn.click(
        fn=remove_background,
        inputs=[input_image],
        outputs=[output_image],
        api_name="remove_bg"
    )
    
    # Download handlers
    download_btn.click(
        fn=lambda x: x,
        inputs=[output_image],
        outputs=[gr.File(label="Downloading...")],
    )
    
    download_jpg.click(
        fn=lambda x: x.convert("RGB") if x else None,
        inputs=[output_image],
        outputs=[gr.File(label="Downloading JPG...")],
    )
    
    # Clear button
    clear_btn.click(
        fn=lambda: [None, None],
        outputs=[input_image, output_image]
    )

# Simplified launch for Hugging Face Spaces
interface.launch()