# -*- coding: utf-8 -*-
"""
AI Background Remover – Hugging Face Space
"""

import spaces
import gradio as gr
from transparent_background import Remover
from PIL import Image
import numpy as np
import base64
import tempfile
import zipfile

# --------------------------------------------------------------------------- #
# Core processing                                                             #
# --------------------------------------------------------------------------- #
@spaces.GPU  # Remove if CPU‑only is desired
def remove_background(image: Image.Image | np.ndarray):
    remover = Remover()

    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    elif not isinstance(image, Image.Image):
        raise TypeError("Unsupported image type")

    # 1) Remove background
    output = remover.process(image)

    # 2) Save PNG
    png_tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
    output.save(png_tmp, format="PNG")
    png_tmp.close()
    png_path = png_tmp.name

    # 3) Build base64 download button
    with open(png_path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    download_link_html = f"""
    <a href="data:application/octet-stream;base64,{b64}"
       download="background_removed.png"
       style="
         display:inline-block;
         text-decoration:none;
         color:#fff;
         background-color:#007bff;
         padding:0.45em 1.2em;
         border-radius:0.3rem;
         font-weight:600;
       ">
       📥 Save Image
    </a>
    """

    # 4) Zip
    zip_tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
    with zipfile.ZipFile(zip_tmp.name, "w", zipfile.ZIP_DEFLATED) as zf:
        zf.write(png_path, arcname="background_removed.png")
    zip_tmp.close()

    return png_path, download_link_html, zip_tmp.name


# --------------------------------------------------------------------------- #
# UI                                                                          #
# --------------------------------------------------------------------------- #
TITLE = "AI Background Remover – Automatically Remove Image Backgrounds"

DESCRIPTION_MD = """
Upload an image and our AI‑powered background remover will return a PNG with a transparent background.

### How to Use
1. Click **Upload Image** or drag & drop a picture.  
2. Press **Start Processing** and wait a moment.  
3. The processed image appears on the right.  
4. Click **📥 Save Image** to download the PNG, or download the ZIP below.

### Application Scenarios
- **E‑commerce** – Clean product photos with transparent backgrounds.  
- **Graphic Design** – Quick cut‑outs for composites or layouts.  
- **Social Media** – Avatars and icons that blend perfectly anywhere.

Our tool leverages state‑of‑the‑art AI to remove backgrounds quickly and accurately—saving you time and effort.
"""

with gr.Blocks(title=TITLE) as demo:
    gr.Markdown(f"# {TITLE}")

    # 说明折叠
    # with gr.Accordion("Instructions & Scenarios (click to collapse)", open=True):
    gr.Markdown(DESCRIPTION_MD)

    # ---------------- 输入区 ---------------- #
    with gr.Column():                          # 整体左列
        input_img = gr.Image(label="Upload Image", type="numpy")
        process_btn = gr.Button("Start Processing", variant="primary")

    # ---------------- 输出区 ---------------- #
    with gr.Row():  # 两列：左图右按钮
        out_img = gr.Image(label="Output Image (PNG)", type="filepath", height=350, scale=3)

        with gr.Column(min_width=200, scale=1):
            download_html = gr.HTML()
            zip_file = gr.File(label="Download ZIP")

    # 绑定事件
    process_btn.click(
        fn=remove_background,
        inputs=input_img,
        outputs=[out_img, download_html, zip_file]
    )

# --------------------------------------------------------------------------- #
if __name__ == "__main__":
    demo.launch(show_error=True, show_api=True,debug=False, share=False)