import gradio as gr

import os
from PIL import Image
import requests
from io import BytesIO
import io
import base64

# Bria API token, stored as a private Space secret.
hf_token = os.environ.get("HF_TOKEN_API_DEMO")
auth_headers = {"api_token": hf_token, "Content-Type": "application/json", "User-Agent": "BriaHFSpace/generative-fill"}

# Bria v2 Generative Fill endpoint.
GEN_FILL_URL = "https://engine.prod.bria-api.com/v2/image/edit/gen_fill"


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()


def pil_to_base64(image):
    """Encode a PIL image as a raw base64 PNG string (no data-URI prefix)."""
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


def download_image(url):
    response = requests.get(url, timeout=60)
    response.raise_for_status()
    return Image.open(BytesIO(response.content)).convert("RGB")


def gen_fill_api_call(image_base64, mask_base64, prompt):
    """Call the Bria v2 Generative Fill endpoint synchronously and return the result image."""
    payload = {
        "image": image_base64,
        "mask": mask_base64,
        "prompt": prompt,
        "sync": True,
    }

    response = requests.post(GEN_FILL_URL, json=payload, headers=auth_headers, timeout=180)

    if response.status_code != 200:
        raise gr.Error(f"Bria API error ({response.status_code}): {response.text}")

    data = response.json()
    result_url = data.get("result", {}).get("image_url")
    if not result_url:
        raise gr.Error(f"Unexpected API response: {data}")

    return download_image(result_url)


def predict(editor_value, prompt):
    if editor_value is None or editor_value.get("background") is None:
        raise gr.Error("Please upload an image first.")

    if not prompt or not prompt.strip():
        raise gr.Error("Please describe what to generate in the prompt.")

    background = editor_value["background"]
    layers = editor_value.get("layers") or []

    init_image = Image.fromarray(background[:, :, :3], "RGB")

    # The brush strokes live in the alpha channel of the first layer:
    # painted pixels become white (255) in the mask = the region to fill.
    if not layers or layers[0][:, :, 3].max() == 0:
        raise gr.Error("Please brush over the area you want to fill.")

    mask = Image.fromarray(layers[0][:, :, 3], "L")

    image_base64 = pil_to_base64(init_image)
    mask_base64 = pil_to_base64(mask)

    return gen_fill_api_call(image_base64, mask_base64, prompt.strip())


# ------------------------------------------------------------------ #
#  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; }

/* ---- Canvas / output ---- */
#image_upload, #output-img { border-radius: 14px; }
#output-img { min-height: 440px; }
#output-img img { object-fit: contain; width: 100%; height: auto; }
/* Gradio sets an inline `height: 450px; overflow: hidden` on the editor, which
   clips the brush toolbar (it sits below the image) — worse on tall/large
   screens. Override inline with !important so the box grows to fit image + tools. */
#image_upload { min-height: 320px; height: auto !important; overflow: visible !important; }
#image_upload .wrap { overflow: visible !important; }
#image_upload .stage-wrap { max-height: 70vh; }

/* ---- Fill 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;
}
#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 Generative Fill API") as demo:
    gr.HTML(f"""
        <div id="bria-header">
            <img src="{BRIA_LOGO}" alt="BRIA logo" />
            <div class="title-block">
                <h1>Generative Fill</h1>
                <p>Add or replace anything in your image — brush an area, describe it, and generate on licensed data.</p>
                <span class="pill">Powered by Bria API v2</span>
            </div>
        </div>
    """)

    gr.HTML("""
        <div id="steps">
            <b>1.</b> Upload an image &nbsp;·&nbsp;
            <b>2.</b> Brush over the area to fill &nbsp;·&nbsp;
            <b>3.</b> Describe what to generate &nbsp;·&nbsp;
            <b>4.</b> Hit <b>Fill</b>
        </div>
    """)

    with gr.Row(equal_height=False):
        with gr.Column(scale=1):
            image = gr.ImageEditor(
                label="Input",
                elem_id="image_upload",
                sources=["upload"],
                layers=False,
                transforms=[],
                brush=gr.Brush(default_color="#5300C9", colors=["#5300C9"], color_mode="fixed"),
            )
            prompt = gr.Textbox(
                label="Prompt",
                elem_id="prompt",
                placeholder="e.g. a vase of fresh flowers",
                lines=2,
            )
            btn = gr.Button("✨  Fill", elem_id="run_button", variant="primary")

        with gr.Column(scale=1):
            image_out = gr.Image(label="Result", elem_id="output-img")

    btn.click(fn=predict, inputs=[image, prompt], outputs=[image_out], api_name="run")

    gr.HTML("""
        <div id="bria-footer">
            <p>
                Built with the <a href="https://docs.bria.ai/image-editing/v2-endpoints/gen-fill" target="_blank">Bria Generative Fill API</a>,
                trained exclusively on licensed data for full legal coverage.<br/>
                Also available on
                <a href="https://fal.ai/models/fal-ai/bria/genfill" target="_blank">fal.ai</a>,
                <a href="https://replicate.com/bria/genfill" target="_blank">Replicate</a>, and as a
                <a href="https://github.com/Bria-AI/ComfyUI-BRIA-API" target="_blank">ComfyUI node</a>.
            </p>
        </div>
    """)

demo.queue(max_size=25, api_open=False).launch(show_api=False)