import base64
import os
from pathlib import Path

os.system("pip freeze")
import spaces
import tempfile
import shutil

import gradio as gr
import torch as torch
from gradio_dualvision import DualVisionApp
from huggingface_hub import login
from PIL import Image
from windowseat_inference import load_network, run_inference

uri_base = "Qwen/Qwen-Image-Edit-2509"
uri_lora = "huawei-bayerlab/windowseat-reflection-removal-v1-0"

if "HF_TOKEN_LOGIN" in os.environ:
    login(token=os.environ["HF_TOKEN_LOGIN"])

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

vae, transformer, embeds_dict, processing_resolution = load_network(uri_base, uri_lora, device)

# # As of transformers==4.57.1 , xformers is not supported in QwenImageTransformer2DModel
# try:
#     transformer.enable_xformers_memory_efficient_attention()
#     print("xformers enabled")
# except:
#     print("xformers not enabled")


SHIELDS_DIR = Path(__file__).parent / "assets" / "shields"
LINKS = [
    ("website", "Website", "https://hf.co/spaces/huawei-bayerlab/windowseat-reflection-removal-web"),
    ("paper", "Paper", "https://arxiv.org/abs/2512.05000"),
    ("code", "Code", "https://github.com/huawei-bayerlab/windowseat-reflection-removal"),
    ("weights", "Weights", "https://hf.co/huawei-bayerlab/windowseat-reflection-removal-v1-0"),
    ("follow", "Follow", "https://twitter.com/antonobukhov1"),
]
# Badges are the committed SVGs, shared with the README, model card, and website. Inlined
# as data URIs so the HTML header needs no static file route.
BADGES = {
    name: "data:image/svg+xml;base64,"
    + base64.b64encode((SHIELDS_DIR / f"{name}.svg").read_bytes()).decode()
    for name, _, _ in LINKS
}
HEADER_CSS = """

    .ws-header { max-width: 70vw; margin: 0 auto; display: grid; grid-template-columns: minmax(0, 6fr) minmax(0, 7fr); gap: 28px; align-items: center; color: var(--body-text-color); }
    .ws-identity { display: flex; flex-direction: column; gap: 6px; border-right: 1px solid var(--border-color-primary); padding-right: 28px; }
    .ws-titles { display: flex; flex-wrap: wrap; align-items: baseline; gap: 4px 14px; }
    .ws-title { font-family: 'Instrument Serif', Georgia, serif; font-weight: 400; font-size: 36px; line-height: 40px; color: var(--body-text-color); }
    .ws-subtitle { font-size: 12px; line-height: 16px; letter-spacing: 0.16em; text-transform: uppercase; color: var(--body-text-color-subdued); }
    .ws-links { display: flex; flex-wrap: wrap; align-items: center; gap: 6px 8px; margin-top: 8px; }
    .ws-link-group { display: inline-flex; gap: inherit; white-space: nowrap; }
    .ws-links a { display: inline-flex; line-height: 0; }
    .ws-links img { height: 22px; width: auto; transition: transform 0.15s ease; }
    .ws-links a:hover img { transform: translateY(-1px); }
    .ws-intro { margin: 0; font-size: 14px; line-height: 20px; }
    @media (max-width: 900px) {
        .ws-header { max-width: none; grid-template-columns: minmax(0, 1fr); gap: 10px; }
        .ws-identity { display: grid; grid-template-columns: max-content minmax(0, 1fr); column-gap: 16px; row-gap: 4px; align-items: center; border-right: 0; border-bottom: 1px solid var(--border-color-primary); padding: 0 0 10px 0; }
        .ws-titles { display: contents; }
        .ws-title { grid-column: 1; grid-row: 1; font-size: 30px; line-height: 34px; }
        .ws-links { grid-column: 2; grid-row: 1; justify-content: flex-end; gap: 6px 8px; margin: 0; }
        .ws-subtitle { grid-column: 1 / -1; grid-row: 2; }
        .ws-intro { font-size: 13px; line-height: 18px; }
    }
    @media (max-width: 529px) {
        .ws-identity { grid-template-columns: minmax(0, 1fr) max-content; }
        .ws-subtitle { grid-column: 1; font-size: 10px; letter-spacing: 0.12em; }
        .ws-links { grid-row: 1 / span 2; align-self: center; flex-direction: column; align-items: flex-end; gap: 6px; }
        .ws-link-group { gap: 6px; }
        .ws-links img { height: 20px; }
    }
    @media (max-width: 370px) {
        .ws-title { font-size: 26px; line-height: 30px; }
    }
    /* Footer links wrap as whole items; outside .contain, so plain rules only. */
    footer.svelte-1byz9vf { flex-wrap: wrap; justify-content: center; gap: 2px 24px; }
    footer.svelte-1byz9vf a { white-space: nowrap; }
    footer.svelte-1byz9vf .divider { display: none; }
    /* Examples block as wide as the drop zone, tiles in a full rectangle (column counts that divide 12). */
    .block:has(> .gallery.svelte-p5q82i) { max-width: 70vw; margin-left: auto; margin-right: auto; }
    .gallery.svelte-p5q82i.svelte-p5q82i.svelte-p5q82i { display: grid; grid-template-columns: repeat(6, minmax(0, 1fr)); gap: var(--spacing-lg); }
    .gallery-item .gallery.svelte-a9zvka { width: 100%; min-width: 0; max-width: none; height: auto; min-height: 0; aspect-ratio: 1; }
    .gallery.svelte-p5q82i .gallery-item img { width: 100%; min-width: 0; max-width: none; height: 100%; min-height: 0; object-fit: cover; }
    @media (max-width: 900px) {
        .block:has(> .gallery.svelte-p5q82i) { max-width: none; }
        .sliderrow .slider { max-width: none; width: 100%; }
        .slider .wrap.half-wrap { width: 100%; }
    }
    @media (max-width: 829px) { .gallery.svelte-p5q82i.svelte-p5q82i.svelte-p5q82i { grid-template-columns: repeat(4, minmax(0, 1fr)); } }
    @media (max-width: 529px) { .gallery.svelte-p5q82i.svelte-p5q82i.svelte-p5q82i { grid-template-columns: repeat(3, minmax(0, 1fr)); } }
"""
HEADER_HEAD = '<link rel="stylesheet" href="https://fonts.googleapis.com/css2?family=Instrument+Serif&display=swap">'


class WindowSeatApp(DualVisionApp):
    DEFAULT_SEED = 2025

    def make_header(self):
        links = "".join(
            '<span class="ws-link-group">'
            + "".join(
                f'<a href="{href}" target="_blank" rel="noopener noreferrer">'
                f'<img src="{BADGES[name]}" alt="{label}"></a>'
                for name, label, href in group
            )
            + "</span>"
            for group in (LINKS[:3], LINKS[3:])
        )
        gr.HTML(
            f"""
            <div class="ws-header">
                <div class="ws-identity">
                    <div class="ws-titles">
                        <div class="ws-title">WindowSeat</div>
                        <div class="ws-subtitle">Image Reflection Removal</div>
                    </div>
                    <div class="ws-links remove-elements">{links}</div>
                </div>
                <p class="ws-intro remove-elements">
                    The authors' demo of <i>Reflection Removal through Efficient Adaptation of Diffusion Transformers</i>.
                    Upload a photo through glass or pick an example below, wait for the result,
                    then drag the slider to compare it with the input and zoom in for detail.
                    If a quota limit appears, duplicate the space to continue.
                </p>
            </div>
            """,
            padding=False,
        )

    def build_user_components(self):
        return {}

    def process(self, image_in: Image.Image, **kwargs):
        input_temp_dir = tempfile.mkdtemp()
        output_temp_dir = tempfile.mkdtemp()

        try:
            input_image_path = os.path.join(input_temp_dir, "image.png")
            image_in.save(input_image_path)
            run_inference(
                vae, 
                transformer, 
                embeds_dict, 
                processing_resolution, 
                input_temp_dir,
                output_temp_dir, 
                use_short_edge_tile=True, 
                save_comparison=False, 
                save_alternating=False,
            )
            output_image_path = os.path.join(output_temp_dir, "image_windowseat_output.png")
            result_image = Image.open(output_image_path)
            result_image.load()

            out_modalities = {
                "Result": result_image,
            }

            out_settings = {}

            return out_modalities, out_settings

        finally:
            if os.path.exists(input_temp_dir):
                shutil.rmtree(input_temp_dir)
            if os.path.exists(output_temp_dir):
                shutil.rmtree(output_temp_dir)


with WindowSeatApp(
    title="WindowSeat Reflection Removal",
    examples_path="example_images",
    examples_per_page=12,
    right_selector_visible=False,
    advanced_settings_visible=False,
    squeeze_canvas=True,
    spaces_zero_gpu_enabled=True,
    css=HEADER_CSS,
    head=HEADER_HEAD,
) as demo:
    demo.queue(
        api_open=False,
    ).launch(
        server_name="0.0.0.0",
        server_port=7860,
        ssr_mode=False,
    )