import os
import base64
import tempfile

import spaces
import gradio as gr
import torch
from diffusers import BriaFiboEditPipeline
from PIL import Image

# Load the model on CPU at module level (ZeroGPU moves to CUDA inside @spaces.GPU)
print("Loading Fibo-Edit-RMBG model...")
pipe = BriaFiboEditPipeline.from_pretrained(
    "briaai/Fibo-Edit-RMBG",
    torch_dtype=torch.bfloat16,
)
print("Model loaded on CPU")

INSTRUCTION = {'edit_instruction': 'Generate a detailed grayscale alpha matte. Map the opaque foreground to white and the background to black. Produce soft, anti-aliased grayscale gradients at the edges of the subject to represent fine details and transparency.'}


@spaces.GPU(duration=300)
def process(image, num_steps=10, guidance_scale=1.0):
    if image is None:
        return None, None

    pipe.to("cuda")
    mask = pipe(
        image=image,
        prompt=INSTRUCTION,
        num_inference_steps=int(num_steps),
        guidance_scale=guidance_scale,
    ).images[0]

    original_rgba = image.convert("RGBA")
    mask_gray = mask.convert("L")
    if mask_gray.size != original_rgba.size:
        mask_gray = mask_gray.resize(original_rgba.size, Image.Resampling.LANCZOS)
    original_rgba.putalpha(mask_gray)
    result = original_rgba

    # Save full-res result for download
    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    result.save(tmp.name, format="PNG")

    return (image.convert("RGBA"), result), tmp.name


def _load_logo():
    """Embed the bundled logo as a data URI so the header never depends on an external CDN."""
    try:
        with open(os.path.join(os.path.dirname(__file__), "bria_logo.png"), "rb") as f:
            return "data:image/png;base64," + base64.b64encode(f.read()).decode("utf-8")
    except Exception:
        return "https://cdn.sanity.io/images/zppcnj0l/production/d0f4aef10abdcc4bacaebeade789ea8c3e453882-210x132.png"


BRIA_LOGO = _load_logo()


# ------------------------------------------------------------------ #
#  Theme + styling — Bria brand (purple #5300C9, magenta accent)     #
# ------------------------------------------------------------------ #
bria_theme = gr.themes.Soft(
    primary_hue=gr.themes.Color(
        c50="#f3ebff", c100="#e4d3ff", c200="#c9a6ff", c300="#a86bff",
        c400="#8a3dff", c500="#5300C9", c600="#4700ab", c700="#37008a",
        c800="#28006b", c900="#1c0050", c950="#120038",
    ),
    secondary_hue=gr.themes.Color(
        c50="#ffe9f4", c100="#ffd0e6", c200="#ffa1cd", c300="#ff6fb2",
        c400="#f5399a", c500="#e2138a", c600="#c00b74", c700="#99095c",
        c800="#730745", c900="#520531", c950="#340120",
    ),
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
    radius_size=gr.themes.sizes.radius_lg,
).set(
    body_background_fill="#f6f5fb",
    body_background_fill_dark="#0f0b1a",
    block_background_fill="white",
    block_background_fill_dark="#1a1330",
    block_border_width="1px",
    block_shadow="0 1px 3px rgba(28,0,80,0.06), 0 8px 24px rgba(28,0,80,0.05)",
    button_primary_background_fill="linear-gradient(90deg, #5300C9 0%, #8a3dff 100%)",
    button_primary_background_fill_hover="linear-gradient(90deg, #4700ab 0%, #7a2dff 100%)",
    button_primary_text_color="white",
    button_large_radius="12px",
)

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

/* ---- Header ---- */
#bria-header {
    display: flex;
    align-items: center;
    gap: 18px;
    padding: 22px 26px;
    margin-bottom: 8px;
    background: linear-gradient(120deg, #ffffff 0%, #f3ecff 100%);
    border: 1px solid rgba(83, 0, 201, 0.12);
    border-radius: 18px;
}
.dark #bria-header {
    background: linear-gradient(120deg, #1a1330 0%, #241640 100%);
    border-color: rgba(138, 61, 255, 0.25);
}
#bria-header img { height: 52px; width: auto; }
#bria-header .title-block h1 {
    margin: 0;
    font-size: 1.7rem;
    font-weight: 800;
    letter-spacing: -0.02em;
    color: #1c0050;
}
.dark #bria-header .title-block h1 { color: #f0eafc; }
#bria-header .title-block p {
    margin: 4px 0 0 0;
    font-size: 0.95rem;
    color: #6b6580;
}
.dark #bria-header .title-block p { color: #b3abc9; }
#bria-header .pill {
    display: inline-block;
    margin-top: 8px;
    padding: 3px 12px;
    font-size: 0.72rem;
    font-weight: 700;
    letter-spacing: 0.04em;
    text-transform: uppercase;
    color: #5300C9;
    background: rgba(83, 0, 201, 0.09);
    border-radius: 999px;
}
.dark #bria-header .pill { color: #c9a6ff; background: rgba(138, 61, 255, 0.15); }

/* ---- Steps helper ---- */
#steps { font-size: 0.92rem; color: #4b4661; margin: 2px 4px 10px; }
.dark #steps { color: #b3abc9; }
#steps b { color: #5300C9; }
.dark #steps b { color: #c9a6ff; }

/* ---- Image panels ---- */
#image_upload, #output-img { border-radius: 14px; min-height: 360px; }
#image_upload img, #output-img img { object-fit: contain; }

/* ---- Run button ---- */
#run_button {
    width: 100%;
    height: 54px;
    font-size: 1.05rem;
    font-weight: 700;
    box-shadow: 0 6px 18px rgba(83, 0, 201, 0.28);
}
#run_button:hover { box-shadow: 0 8px 22px rgba(83, 0, 201, 0.38); }

/* ---- Footer ---- */
#bria-footer {
    margin-top: 28px;
    padding-top: 18px;
    border-top: 1px solid rgba(83, 0, 201, 0.12);
    text-align: center;
    font-size: 0.88rem;
    color: #6b6580;
    line-height: 1.7;
}
#bria-footer a { color: #5300C9; text-decoration: none; font-weight: 600; }
#bria-footer a:hover { text-decoration: underline; }
.dark #bria-footer { color: #b3abc9; border-color: rgba(138, 61, 255, 0.22); }
.dark #bria-footer a { color: #c9a6ff; }
"""

with gr.Blocks(theme=bria_theme, css=css, title="BRIA Fibo-Edit-RMBG") as demo:
    gr.HTML(f"""
        <div id="bria-header">
            <img src="{BRIA_LOGO}" alt="BRIA logo" />
            <div class="title-block">
                <h1>Remove Background</h1>
                <p>Fibo-Edit-RMBG — background removal via Bria's 8B FIBO-Edit model, on licensed data.</p>
                <span class="pill">Open model · Fibo-Edit-RMBG</span>
            </div>
        </div>
    """)

    gr.HTML("""
        <div id="steps">
            <b>1.</b> Upload an image &nbsp;·&nbsp;
            <b>2.</b> (optional) tune the settings &nbsp;·&nbsp;
            <b>3.</b> Hit <b>Remove Background</b> &nbsp;·&nbsp;
            <b>4.</b> Drag the slider to compare, then download the PNG
        </div>
    """)

    with gr.Row(equal_height=False):
        with gr.Column(scale=1):
            input_image = gr.Image(label="Input image", type="pil", elem_id="image_upload")

            with gr.Accordion("Advanced settings", open=False):
                num_steps = gr.Slider(
                    minimum=1,
                    maximum=50,
                    value=10,
                    step=1,
                    label="Number of steps",
                )
                guidance_scale = gr.Slider(
                    minimum=1.0,
                    maximum=15.0,
                    value=1.0,
                    step=0.5,
                    label="Guidance scale",
                )

            process_btn = gr.Button("✨  Remove Background", variant="primary", elem_id="run_button")

        with gr.Column(scale=1):
            slider = gr.ImageSlider(
                label="Original vs. background removed",
                type="pil",
                image_mode="RGBA",
                format="png",
                slider_position=50,
                elem_id="output-img",
                container=True,
            )
            download_file = gr.File(label="PNG with transparent background")

    gr.HTML("""
        <div id="bria-footer">
            <p>
                Fibo-Edit-RMBG is built on the 8B <a href="https://huggingface.co/briaai/Fibo-Edit" target="_blank">FIBO-Edit</a> model and is open for non-commercial use (<a href="https://creativecommons.org/licenses/by-nc/4.0/" target="_blank">CC BY-NC 4.0</a>).<br/>
                <a href="https://huggingface.co/briaai/Fibo-Edit-RMBG" target="_blank">Model card</a> ·
                <a href="https://bria.ai/contact-us" target="_blank">Commercial licensing</a> ·
                <a href="https://bria.ai" target="_blank">Bria AI</a>
            </p>
        </div>
    """)

    process_btn.click(
        fn=process,
        inputs=[input_image, num_steps, guidance_scale],
        outputs=[slider, download_file],
    )

if __name__ == "__main__":
    demo.launch(show_error=True)