import os
from typing import Tuple, Union

import gradio as gr
import spaces
import torch
from loadimg import load_img
from PIL import Image
from torchvision import transforms
from transformers import AutoModelForImageSegmentation

torch.set_float32_matmul_precision("high")

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

birefnet = AutoModelForImageSegmentation.from_pretrained(
    "ZhengPeng7/BiRefNet", trust_remote_code=True
)
birefnet.to(DEVICE)
birefnet.eval()

# Keep inference inputs aligned with the loaded model precision (fp16/fp32).
MODEL_DTYPE = next(birefnet.parameters()).dtype

transform_image = transforms.Compose(
    [
        transforms.Resize((1024, 1024)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)

EXAMPLE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")
EXAMPLE_IMAGE = os.path.join(EXAMPLE_DIR, "butterfly.jpg")
EXAMPLE_URL = (
    "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
)


@spaces.GPU
def process(image: Image.Image) -> Image.Image:
    """
    Apply BiRefNet-based image segmentation to remove the background.
    Args:
        image (PIL.Image): The input RGB image.
    Returns:
        PIL.Image: The image with the background removed, using the segmentation mask as transparency.
    """
    image_size = image.size
    input_images = (
        transform_image(image).unsqueeze(0).to(device=DEVICE, dtype=MODEL_DTYPE)
    )
    with torch.no_grad():
        preds = birefnet(input_images)[-1].sigmoid().cpu()
    pred = preds[0].squeeze()
    pred_pil = transforms.ToPILImage()(pred)
    mask = pred_pil.resize(image_size)
    image.putalpha(mask)
    return image


def remove_bg(image: Union[Image.Image, str]) -> Tuple[Image.Image, Image.Image]:
    """Remove the background from an uploaded image or an image URL."""
    im = load_img(image, output_type="pil").convert("RGB")
    origin = im.copy()
    return origin, process(im)


def remove_bg_to_file(f: str) -> str:
    """Load an image file from disk, remove the background, and save as a transparent PNG."""
    name_path = os.path.splitext(f)[0] + ".png"
    im = load_img(f, output_type="pil").convert("RGB")
    process(im).save(name_path)
    return name_path


AUTHOR_NAME = "Bijaya Kumar Tiadi"
AUTHOR_TITLE = "Full-Stack Developer"
UPWORK_URL = "https://www.upwork.com/freelancers/~01d86f12209c236752"
LINKEDIN_URL = "https://www.linkedin.com/in/bijayakumartiadi/"

THEME = gr.themes.Soft(
    primary_hue="violet",
    secondary_hue="purple",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
)

CSS = """
.gradio-container { max-width: 980px !important; margin: 0 auto !important; }

#hero {
    text-align: center;
    padding: 34px 24px 30px;
    margin-bottom: 8px;
    border-radius: 20px;
    background: linear-gradient(135deg, #7c3aed 0%, #a855f7 45%, #d946ef 100%);
    box-shadow: 0 10px 30px -12px rgba(124, 58, 237, 0.55);
}
#hero h1 { font-size: 2.3rem; margin-bottom: 6px; color: #ffffff; font-weight: 800; letter-spacing: -0.02em; }
#hero p { color: rgba(255,255,255,0.92); font-size: 1.05rem; margin: 0 auto; max-width: 560px; }
#hero .badges { margin-top: 14px; display: flex; gap: 8px; justify-content: center; flex-wrap: wrap; }
#hero .badge {
    display: inline-block; padding: 4px 12px; border-radius: 999px;
    background: rgba(255,255,255,0.18); color: #fff; font-size: 0.78rem;
    font-weight: 600; backdrop-filter: blur(4px); border: 1px solid rgba(255,255,255,0.3);
}

#run-btn { min-height: 46px; font-weight: 600; }

.hire-card {
    margin-top: 18px;
    padding: 22px 26px;
    border-radius: 18px;
    background: var(--block-background-fill);
    border: 1px solid var(--border-color-primary);
    display: flex;
    align-items: center;
    justify-content: space-between;
    gap: 20px;
    flex-wrap: wrap;
}
.hire-card .who { display: flex; align-items: center; gap: 14px; }
.hire-card .avatar {
    width: 48px; height: 48px; border-radius: 50%; flex-shrink: 0;
    background: linear-gradient(135deg, #7c3aed, #d946ef);
    color: #fff; display: flex; align-items: center; justify-content: center;
    font-weight: 700; font-size: 1.05rem;
}
.hire-card .name { font-weight: 700; font-size: 1.05rem; margin: 0; }
.hire-card .tagline { margin: 0; opacity: 0.7; font-size: 0.9rem; }
.hire-card .cta { margin: 0 0 0 2px; opacity: 0.75; font-size: 0.85rem; }
.btn-row { display: flex; gap: 10px; flex-wrap: wrap; }
.btn-pill {
    display: inline-flex; align-items: center; gap: 8px;
    padding: 10px 18px; border-radius: 999px; font-weight: 600; font-size: 0.9rem;
    text-decoration: none !important; transition: transform 0.15s ease, box-shadow 0.15s ease;
    box-shadow: 0 2px 8px rgba(0,0,0,0.12);
}
.btn-pill:hover { transform: translateY(-2px); box-shadow: 0 6px 16px rgba(0,0,0,0.18); }
.btn-pill svg { width: 16px; height: 16px; fill: currentColor; }
.btn-upwork { background: #14a800; color: #fff !important; }
.btn-linkedin { background: #0a66c2; color: #fff !important; }

.footer-note { text-align: center; opacity: 0.6; font-size: 0.82rem; margin-top: 14px; }
.footer-note a { color: inherit; }
"""

with gr.Blocks(title="Background Remover") as demo:
    gr.Markdown(
        """
        <div id="hero">
        <h1>🪄 Background Remover</h1>
        <p>Drop an image, get a clean transparent cutout. Powered by the open-source
        <a href="https://huggingface.co/ZhengPeng7/BiRefNet" target="_blank" style="color:#fff;text-decoration:underline;">BiRefNet</a> model — free, no sign-up, no watermark.</p>
        <div class="badges">
            <span class="badge">⚡ Free to use</span>
            <span class="badge">🖼️ No watermark</span>
            <span class="badge">🔌 API / MCP ready</span>
        </div>
        </div>
        """
    )

    with gr.Tabs():
        with gr.Tab("📤 Upload"):
            with gr.Row():
                with gr.Column():
                    upload_in = gr.Image(
                        label="Source image", type="pil", sources=["upload", "clipboard"]
                    )
                    upload_btn = gr.Button(
                        "Remove background", variant="primary", elem_id="run-btn"
                    )
                with gr.Column():
                    upload_out = gr.ImageSlider(
                        label="Result — drag to compare", type="pil", format="png"
                    )
            gr.Examples(examples=[EXAMPLE_IMAGE], inputs=upload_in, label="Try an example")
            upload_btn.click(remove_bg, inputs=upload_in, outputs=upload_out, api_name="image")

        with gr.Tab("🔗 Image URL"):
            with gr.Row():
                with gr.Column():
                    url_in = gr.Textbox(
                        label="Image URL", placeholder="https://example.com/photo.jpg"
                    )
                    url_btn = gr.Button(
                        "Remove background", variant="primary", elem_id="run-btn"
                    )
                with gr.Column():
                    url_out = gr.ImageSlider(
                        label="Result — drag to compare", type="pil", format="png"
                    )
            gr.Examples(examples=[EXAMPLE_URL], inputs=url_in, label="Try an example")
            url_btn.click(remove_bg, inputs=url_in, outputs=url_out, api_name="text")

        with gr.Tab("📁 File → PNG"):
            gr.Markdown("Upload a file and get back a downloadable transparent PNG — handy for automation, the API, or MCP.")
            with gr.Row():
                file_in = gr.Image(label="Source image", type="filepath")
                file_out = gr.File(label="Transparent PNG")
            file_btn = gr.Button("Generate PNG", variant="primary", elem_id="run-btn")
            gr.Examples(examples=[EXAMPLE_IMAGE], inputs=file_in, label="Try an example")
            file_btn.click(remove_bg_to_file, inputs=file_in, outputs=file_out, api_name="png")

    gr.HTML(
        f"""
        <div class="hire-card">
            <div class="who">
                <div class="avatar">BT</div>
                <div>
                    <p class="name">{AUTHOR_NAME}</p>
                    <p class="tagline">{AUTHOR_TITLE}</p>
                    <p class="cta">Need something like this built for your product? Let's talk 👇</p>
                </div>
            </div>
            <div class="btn-row">
                <a class="btn-pill btn-upwork" href="{UPWORK_URL}" target="_blank" rel="noopener">
                    <svg viewBox="0 0 24 24"><path d="M18.561 13.158c-1.102 0-2.135-.467-3.074-1.227l.228-1.076.008-.042c.207-1.143.849-3.06 2.839-3.06 1.492 0 2.703 1.212 2.703 2.703-.001 1.489-1.212 2.702-2.704 2.702zm0-8.14c-2.539 0-4.51 1.649-5.31 4.366-1.22-1.834-2.148-4.036-2.687-5.892H7.828v7.112c-.002 1.406-1.141 2.546-2.547 2.546-1.405 0-2.543-1.14-2.545-2.546V3.492H0v7.112c0 2.914 2.37 5.303 5.281 5.303 2.913 0 5.283-2.389 5.283-5.303v-1.19c.529 1.107 1.182 2.229 1.974 3.221l-1.673 7.873h2.797l1.213-5.71c1.063.696 2.285 1.129 3.686 1.129 3.017 0 5.469-2.474 5.469-5.508 0-3.034-2.452-5.484-5.469-5.484z"/></svg>
                    Hire me on Upwork
                </a>
                <a class="btn-pill btn-linkedin" href="{LINKEDIN_URL}" target="_blank" rel="noopener">
                    <svg viewBox="0 0 24 24"><path d="M20.447 20.452h-3.554v-5.569c0-1.328-.027-3.037-1.852-3.037-1.853 0-2.136 1.445-2.136 2.939v5.667H9.351V9h3.414v1.561h.046c.477-.9 1.637-1.85 3.37-1.85 3.601 0 4.267 2.37 4.267 5.455v6.286zM5.337 7.433a2.062 2.062 0 1 1 0-4.125 2.062 2.062 0 0 1 0 4.125zM7.114 20.452H3.558V9h3.556v11.452z"/></svg>
                    Connect on LinkedIn
                </a>
            </div>
        </div>
        """
    )

    gr.Markdown(
        """
        <div class="footer-note">
        Model: <a href="https://huggingface.co/ZhengPeng7/BiRefNet" target="_blank">ZhengPeng7/BiRefNet</a> ·
        UI: <a href="https://gradio.app" target="_blank">Gradio</a> ·
        Runs free on <a href="https://huggingface.co/docs/hub/spaces-zerogpu" target="_blank">ZeroGPU</a><br/>
        Prefer a dedicated API? Use <a href="https://huggingface.co/spaces/hf-applications/background-removal" target="_blank">hf-applications/background-removal</a>,
        or call the <code>png</code> endpoint from the "File → PNG" tab above.
        </div>
        """
    )

if __name__ == "__main__":
    demo.launch(show_error=True, mcp_server=True, theme=THEME, css=CSS)