import gradio as gr
from PIL import Image
import io
from rembg import remove

# =========================
# Core Function
# =========================
def remove_background(input_image):
    if input_image is None:
        return None, "⚠️ Please upload an image."

    try:
        # Convert PIL image to bytes
        img_byte_arr = io.BytesIO()
        input_image.save(img_byte_arr, format="PNG")
        img_bytes = img_byte_arr.getvalue()

        # Remove background
        output_bytes = remove(img_bytes)

        # Convert back to PIL image
        result_img = Image.open(io.BytesIO(output_bytes)).convert("RGBA")

        return result_img, "✅ Background removed successfully"

    except Exception as e:
        return None, f"❌ Error: {str(e)}"

# =========================
# CSS
# =========================
custom_css = """
.gradio-container { max-width: 1000px !important; margin: auto; }
.main-header {
    text-align: center;
    margin-bottom: 2rem;
    padding: 2rem;
    background: #fcfcfc;
    border-radius: 24px;
    border: 1px solid #eee;
}
.footer-text {
    text-align: center;
    margin-top: 3rem;
    font-size: 0.85rem;
    color: #94a3b8;
}
"""

# =========================
# UI
# =========================
with gr.Blocks(title="Remove BG by Nadish") as demo:

    gr.HTML("""
    <div class="main-header">
        <h1>🎨 Remove Background by Nadish</h1>
        <h3>Fast & Accurate AI Background Removal</h3>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            input_img = gr.Image(
                type="pil",
                label="Upload Image"
            )
            run_btn = gr.Button(
                "✨ Remove Background",
                variant="primary",
                size="lg"
            )

        with gr.Column(scale=1):
            output_img = gr.Image(
                type="pil",
                label="Result (Transparent PNG)",
                interactive=False
            )
            status_text = gr.Textbox(
                label="Status",
                interactive=False
            )

    run_btn.click(
        fn=remove_background,
        inputs=input_img,
        outputs=[output_img, status_text]
    )

    gr.HTML("""
    <div class="footer-text">
        Built by Nadish • Powered by rembg
    </div>
    """)

# =========================
# Launch
# =========================
if __name__ == "__main__":
    demo.launch(css=custom_css)