#!/usr/bin/env python3
"""Gradio demo for Robust-U1 on Hugging Face Spaces."""

from __future__ import annotations

import os
import random
import subprocess
import sys
import threading
from pathlib import Path
from typing import Any, Optional

import gradio as gr
import numpy as np
import torch
from huggingface_hub import snapshot_download
from PIL import Image

REPO_ROOT = Path(__file__).resolve().parent
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

is_spaces = os.getenv("SPACE_ID") is not None
spaces_available = False
GPU = None

if is_spaces:
    try:
        from spaces import GPU
        spaces_available = True
    except ImportError:
        pass


def gpu_decorator(func):
    if spaces_available and GPU is not None:
        return GPU(duration=240)(func)
    return func


project_dir = os.path.dirname(os.path.abspath(__file__))
if not is_spaces:
    temp_dir = os.path.join(project_dir, ".gradio_temp")
    os.makedirs(temp_dir, exist_ok=True)
    os.environ["GRADIO_TEMP_DIR"] = temp_dir

MODEL_PATH = os.getenv("MODEL_PATH", "Jiaqi-hkust/Robust-U1")
MAX_MEM_PER_GPU = os.environ.get("MAX_MEM_PER_GPU", "80GiB")
OFFLOAD_DIR = Path(os.environ.get("ROBUST_U1_OFFLOAD_DIR", "/tmp/robust_u1_offload"))
MODEL_ALLOW_PATTERNS = (
    "*.json",
    "*.safetensors",
    "*.bin",
    "*.txt",
    "*.model",
    "*.py",
    "*.md",
)

DEFAULT_PROMPT = (
    "Which direction is the vehicle directly in front of us traveling?\n"
    "Options:\n"
    "A. Straight\n"
    "B. Left\n"
    "C. Right"
)

HIDDEN_PROMPT_SUFFIX = (
    "Please restore this corrupted image to its clean version.\n"
    "Based on what you observe in the restored image, please select the correct answer from the options above."
)

APP_CSS = """
:root {
    --page-max-width: 100vw;
    --ink-strong: #102a43;
    --ink-soft: #486581;
    --surface-primary: rgba(255, 255, 255, 0.92);
    --surface-secondary: rgba(244, 247, 251, 0.88);
    --surface-border: rgba(148, 163, 184, 0.24);
    --surface-shadow: 0 24px 60px rgba(15, 23, 42, 0.08);
    --accent-start: #0f766e;
    --accent-end: #0b5ed7;
}

body {
    background:
        radial-gradient(circle at top left, rgba(15, 118, 110, 0.16), transparent 34%),
        radial-gradient(circle at top right, rgba(14, 116, 144, 0.14), transparent 30%),
        linear-gradient(180deg, #eef4f7 0%, #f8fbfd 52%, #eef3f8 100%);
}

body,
.gradio-container,
input,
textarea,
button {
    font-family: "IBM Plex Sans", "Avenir Next", "Segoe UI", sans-serif !important;
}

.gradio-container {
    width: 100% !important;
    max-width: none !important;
    padding: 24px 12px 40px !important;
}

.app-shell {
    gap: 20px;
}

.hero-card {
    position: relative;
    overflow: hidden;
    padding: 26px 34px 24px;
    border-radius: 28px;
    background: linear-gradient(135deg, rgba(255, 255, 255, 0.9) 0%, rgba(232, 243, 248, 0.96) 100%);
    border: 1px solid rgba(148, 163, 184, 0.2);
    box-shadow: var(--surface-shadow);
}

.hero-card::after {
    content: "";
    position: absolute;
    inset: auto -80px -110px auto;
    width: 260px;
    height: 260px;
    border-radius: 999px;
    background: radial-gradient(circle, rgba(11, 94, 215, 0.15), transparent 68%);
}

.hero-title {
    margin: 0;
    font-size: clamp(2.1rem, 4vw, 3.4rem);
    line-height: 1.05;
    color: var(--ink-strong);
}

.hero-subtitle {
    max-width: 760px;
    margin: 12px 0 0;
    font-size: 0.98rem;
    line-height: 1.6;
    color: #52606d;
}

.panel {
    border: 1px solid var(--surface-border);
    border-radius: 24px;
    background: var(--surface-primary);
    box-shadow: var(--surface-shadow);
    padding: 20px;
    backdrop-filter: blur(12px);
}

.panel-header {
    margin-bottom: 14px;
}

.eyebrow {
    font-size: 12px;
    font-weight: 700;
    letter-spacing: 0.16em;
    text-transform: uppercase;
    color: #486581;
}

.panel-title {
    margin-top: 4px;
    font-size: 1.22rem;
    font-weight: 600;
    color: var(--ink-strong);
}

.control-bar {
    margin-top: 16px;
}

.toolbar {
    align-items: center;
    gap: 12px;
}

.accordion-shell {
    border-radius: 24px !important;
    border: 1px solid var(--surface-border) !important;
    background: var(--surface-primary) !important;
    box-shadow: var(--surface-shadow);
}

.accordion-shell > .label-wrap {
    padding-top: 4px;
    padding-bottom: 4px;
}

#run-button {
    min-height: 52px;
    border: none !important;
    background: linear-gradient(135deg, var(--accent-start) 0%, var(--accent-end) 100%) !important;
    box-shadow: 0 18px 32px rgba(11, 94, 215, 0.22);
}

#run-button:hover {
    filter: brightness(1.03);
}

#secondary-button {
    min-height: 52px;
    border: 1px solid rgba(148, 163, 184, 0.24) !important;
    background: rgba(248, 250, 252, 0.9) !important;
    color: var(--ink-strong) !important;
}

.input-image,
.output-image {
    width: 100% !important;
}

.input-image img,
.output-image img {
    border-radius: 18px;
}

.reasoning-box textarea {
    line-height: 1.55;
}

.footer-note {
    margin-top: 6px;
    font-size: 0.95rem;
    line-height: 1.65;
    color: #52606d;
}

@media (max-width: 900px) {
    .hero-card {
        padding: 22px 22px 20px;
    }

    .panel {
        padding: 16px;
    }

    .toolbar {
        flex-direction: column;
        align-items: stretch;
    }
}
"""

_runtime_lock = threading.Lock()
_runtime_cache: dict[str, Any] = {
    "inferencer": None,
    "pil_img2rgb": None,
}


def _build_header_html() -> str:
    return """
    <section class="hero-card">
      <h1 class="hero-title">Robust-U1</h1>
      <p class="hero-subtitle">
        Robust-U1: Can MLLMs Self-Recover Corrupted Visual Content for Robust Understanding?
      </p>
    </section>
    """


def _ensure_flash_attn_available() -> None:
    try:
        import flash_attn  # noqa: F401
    except ModuleNotFoundError:
        env = os.environ.copy()
        env.setdefault("FLASH_ATTENTION_SKIP_CUDA_BUILD", "TRUE")
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "flash-attn", "--no-build-isolation"],
            check=True,
            env=env,
        )


def _resolve_checkpoint(model_path: Path) -> Path:
    for name in ("model_bf16.safetensors", "model.safetensors", "pytorch_model.bin"):
        candidate = model_path / name
        if candidate.is_file():
            return candidate
    raise FileNotFoundError(f"No checkpoint file found under {model_path}")


def _resolve_model_path() -> Path:
    candidate = Path(MODEL_PATH).expanduser()
    if candidate.exists():
        return candidate.resolve()

    token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
    return Path(
        snapshot_download(
            repo_id=MODEL_PATH,
            repo_type="model",
            token=token,
            allow_patterns=list(MODEL_ALLOW_PATTERNS),
        )
    )


def set_seed(seed: int) -> None:
    if seed <= 0:
        return
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def _build_device_map(model: Any) -> tuple[dict[str, Any], str]:
    from accelerate import infer_auto_device_map

    gpu_count = torch.cuda.device_count()
    if gpu_count < 1:
        raise RuntimeError("Robust-U1 inference requires a GPU-enabled Hugging Face Space.")

    device_map = infer_auto_device_map(
        model,
        max_memory={i: MAX_MEM_PER_GPU for i in range(gpu_count)},
        no_split_module_classes=["Bagel", "Qwen2MoTDecoderLayer"],
    )

    same_device_modules = [
        "language_model.model.embed_tokens",
        "time_embedder",
        "latent_pos_embed",
        "vae2llm",
        "llm2vae",
        "connector",
        "vit_pos_embed",
    ]

    if gpu_count == 1:
        first_device = device_map.get(same_device_modules[0], "cuda:0")
        for module_name in same_device_modules:
            device_map[module_name] = first_device if module_name in device_map else "cuda:0"
    else:
        first_device = device_map.get(same_device_modules[0], "cuda:0")
        for module_name in same_device_modules:
            if module_name in device_map:
                device_map[module_name] = first_device

    return device_map, same_device_modules[0]


def _clean_reasoning_trace(text: str) -> str:
    for token in ("<|im_start|>", "<|im_end|>", "<|endoftext|>"):
        text = text.replace(token, "")
    end_tag = "</think>"
    if end_tag in text:
        text = text.split(end_tag, 1)[0] + end_tag
    return text

def ensure_runtime() -> tuple[Any, Any]:
    inferencer = _runtime_cache["inferencer"]
    pil_img2rgb = _runtime_cache["pil_img2rgb"]
    if inferencer is not None and pil_img2rgb is not None:
        return inferencer, pil_img2rgb

    with _runtime_lock:
        inferencer = _runtime_cache["inferencer"]
        pil_img2rgb = _runtime_cache["pil_img2rgb"]
        if inferencer is not None and pil_img2rgb is not None:
            return inferencer, pil_img2rgb

        _ensure_flash_attn_available()

        from accelerate import init_empty_weights, load_checkpoint_and_dispatch
        from data.data_utils import add_special_tokens, pil_img2rgb
        from data.transforms import ImageTransform
        from inferencer import InterleaveInferencer
        from modeling.autoencoder import load_ae
        from modeling.bagel import (
            Bagel,
            BagelConfig,
            Qwen2Config,
            Qwen2ForCausalLM,
            SiglipVisionConfig,
            SiglipVisionModel,
        )
        from modeling.qwen2 import Qwen2Tokenizer

        model_path = _resolve_model_path()
        checkpoint_path = _resolve_checkpoint(model_path)

        llm_config = Qwen2Config.from_json_file(os.path.join(model_path, "llm_config.json"))
        llm_config.qk_norm = True
        llm_config.tie_word_embeddings = False
        llm_config.layer_module = "Qwen2MoTDecoderLayer"

        vit_config = SiglipVisionConfig.from_json_file(os.path.join(model_path, "vit_config.json"))
        vit_config.rope = False
        vit_config.num_hidden_layers -= 1

        vae_model, vae_config = load_ae(local_path=os.path.join(model_path, "ae.safetensors"))

        config = BagelConfig(
            visual_gen=True,
            visual_und=True,
            llm_config=llm_config,
            vit_config=vit_config,
            vae_config=vae_config,
            vit_max_num_patch_per_side=70,
            connector_act="gelu_pytorch_tanh",
            latent_patch_size=2,
            max_latent_size=64,
        )

        with init_empty_weights():
            language_model = Qwen2ForCausalLM(llm_config)
            vit_model = SiglipVisionModel(vit_config)
            model = Bagel(language_model, vit_model, config)
            model.vit_model.vision_model.embeddings.convert_conv2d_to_linear(vit_config, meta=True)

        tokenizer = Qwen2Tokenizer.from_pretrained(str(model_path))
        tokenizer, new_token_ids, _ = add_special_tokens(tokenizer)

        vae_transform = ImageTransform(1024, 512, 16)
        vit_transform = ImageTransform(980, 224, 14)

        device_map, vae_anchor_module = _build_device_map(model)

        OFFLOAD_DIR.mkdir(parents=True, exist_ok=True)
        model = load_checkpoint_and_dispatch(
            model,
            checkpoint=str(checkpoint_path),
            device_map=device_map,
            offload_buffers=True,
            offload_folder=str(OFFLOAD_DIR),
            dtype=torch.bfloat16,
            force_hooks=True,
        ).eval()

        vae_device = device_map.get(vae_anchor_module, "cuda:0")
        vae_model = vae_model.to(device=vae_device, dtype=torch.bfloat16).eval()

        inferencer = InterleaveInferencer(
            model=model,
            vae_model=vae_model,
            tokenizer=tokenizer,
            vae_transform=vae_transform,
            vit_transform=vit_transform,
            new_token_ids=new_token_ids,
        )

        _runtime_cache["inferencer"] = inferencer
        _runtime_cache["pil_img2rgb"] = pil_img2rgb
        return inferencer, pil_img2rgb


@gpu_decorator
def edit_image(
    image: Optional[Image.Image],
    prompt: str,
    show_thinking: bool = False,
    cfg_text_scale: float = 4.0,
    cfg_img_scale: float = 2.0,
    cfg_interval: float = 0.0,
    timestep_shift: float = 3.0,
    num_timesteps: int = 50,
    cfg_renorm_min: float = 1.0,
    cfg_renorm_type: str = "text_channel",
    max_think_token_n: int = 1024,
    do_sample: bool = False,
    text_temperature: float = 0.3,
    seed: int = 0,
):
    if image is None:
        yield None, "Upload an image to begin."
        return

    try:
        inferencer, pil_img2rgb = ensure_runtime()
    except Exception as exc:
        yield None, f"Unable to initialize Robust-U1 on this Space: {exc}"
        return

    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)

    set_seed(seed)
    image = pil_img2rgb(image)

    inference_hyper = dict(
        max_think_token_n=max_think_token_n if show_thinking else 1024,
        do_sample=do_sample if show_thinking else False,
        text_temperature=text_temperature if show_thinking else 0.3,
        cfg_text_scale=cfg_text_scale,
        cfg_img_scale=cfg_img_scale,
        cfg_interval=[cfg_interval, 1.0],
        timestep_shift=timestep_shift,
        num_timesteps=num_timesteps,
        cfg_renorm_min=cfg_renorm_min,
        cfg_renorm_type=cfg_renorm_type,
    )

    result_text = ""
    last_image = None
    model_prompt = f"{prompt.rstrip()}\n\n{HIDDEN_PROMPT_SUFFIX}" if prompt.strip() else HIDDEN_PROMPT_SUFFIX

    for chunk in inferencer(
        image=image,
        text=model_prompt,
        think=show_thinking,
        understanding_output=False,
        **inference_hyper,
    ):
        if isinstance(chunk, str):
            chunk = (
                chunk.replace("<|im_start|>", "")
                     .replace("<|im_end|>", "")
                     .replace("<|endoftext|>", "")
            )
            result_text += chunk
        else:
            last_image = chunk
        yield last_image, _clean_reasoning_trace(result_text)


def update_edit_thinking_visibility(show: bool):
    return gr.update(visible=show), gr.update(visible=show)


def load_example_image(image_path: Path) -> Optional[Image.Image]:
    if not image_path.is_file():
        return None
    return Image.open(image_path).convert("RGB")


def build_demo() -> gr.Blocks:
    default_in = REPO_ROOT / "test_images" / "Dis_image.jpg"
    default_pil = load_example_image(default_in)

    def reset_workspace():
        return (
            default_pil,
            DEFAULT_PROMPT,
            None,
            gr.update(value="", visible=True),
            True,
            4.0,
            2.0,
            0.0,
            3.0,
            50,
            1.0,
            "text_channel",
            1024,
            False,
            0.3,
            0,
            gr.update(visible=False),
        )

    theme = gr.themes.Soft(
        primary_hue=gr.themes.colors.cyan,
        secondary_hue=gr.themes.colors.blue,
        neutral_hue=gr.themes.colors.slate,
    )

    with gr.Blocks(
        theme=theme,
        css=APP_CSS,
        title="Robust-U1",
        elem_classes="app-shell",
        fill_width=True,
    ) as demo:
        gr.HTML(_build_header_html())

        with gr.Row(equal_height=False):
            with gr.Column(scale=1, elem_classes="panel"):
                gr.HTML(
                    """
                    <div class="panel-header">
                      <div class="eyebrow">Input</div>
                      <div class="panel-title">Source Image and Edit Instruction</div>
                    </div>
                    """
                )
                edit_image_input = gr.Image(
                    label="Input image",
                    value=default_pil,
                    type="pil",
                    sources=["upload", "clipboard"],
                    elem_classes="input-image",
                )
                edit_prompt = gr.Textbox(
                    label="Prompt",
                    value=DEFAULT_PROMPT,
                    placeholder="Enter an edit instruction or question.",
                    lines=4,
                )
                with gr.Row(elem_classes="control-bar toolbar"):
                    edit_show_thinking = gr.Checkbox(label="Show reasoning trace", value=True)
                    edit_btn = gr.Button("Run Inference", variant="primary", elem_id="run-button")
                    reset_btn = gr.Button("Reset Workspace", elem_id="secondary-button")

            with gr.Column(scale=1, elem_classes="panel"):
                gr.HTML(
                    """
                    <div class="panel-header">
                      <div class="eyebrow">Output</div>
                      <div class="panel-title">Generated Result and Trace</div>
                    </div>
                    """
                )
                edit_image_output = gr.Image(
                    label="Edited output",
                    type="pil",
                    interactive=False,
                    show_download_button=True,
                    elem_classes="output-image",
                )
                edit_thinking_output = gr.Textbox(
                    label="Reasoning trace",
                    visible=True,
                    lines=8,
                    show_copy_button=True,
                    elem_classes="reasoning-box",
                )

        with gr.Accordion("Advanced Inference Controls", open=False, elem_classes="accordion-shell"):
            with gr.Row():
                edit_seed = gr.Slider(
                    minimum=0,
                    maximum=1_000_000,
                    value=0,
                    step=1,
                    label="Seed",
                    info="Use 0 for a non-deterministic run, or a positive integer for reproducibility.",
                )
                edit_cfg_text_scale = gr.Slider(
                    minimum=1.0,
                    maximum=8.0,
                    value=4.0,
                    step=0.1,
                    label="CFG text scale",
                )
            with gr.Row():
                edit_cfg_img_scale = gr.Slider(
                    minimum=1.0,
                    maximum=4.0,
                    value=2.0,
                    step=0.1,
                    label="CFG image scale",
                )
                edit_cfg_interval = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=0.0,
                    step=0.1,
                    label="CFG activation start",
                )
            with gr.Row():
                edit_cfg_renorm_type = gr.Dropdown(
                    choices=["global", "local", "text_channel"],
                    value="text_channel",
                    label="CFG renormalization type",
                )
                edit_cfg_renorm_min = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=1.0,
                    step=0.1,
                    label="CFG renormalization floor",
                )
            with gr.Row():
                edit_num_timesteps = gr.Slider(
                    minimum=10,
                    maximum=100,
                    value=50,
                    step=5,
                    label="Diffusion timesteps",
                )
                edit_timestep_shift = gr.Slider(
                    minimum=1.0,
                    maximum=10.0,
                    value=3.0,
                    step=0.5,
                    label="Timestep shift",
                )
            edit_thinking_params = gr.Group(visible=True)
            with edit_thinking_params:
                with gr.Row():
                    edit_do_sample = gr.Checkbox(label="Sample reasoning tokens", value=False)
                    edit_max_think_token_n = gr.Slider(
                        minimum=64,
                        maximum=4006,
                        value=1024,
                        step=64,
                        label="Max reasoning tokens",
                    )
                    edit_text_temperature = gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        value=0.3,
                        step=0.1,
                        label="Reasoning temperature",
                    )

        edit_show_thinking.change(
            fn=update_edit_thinking_visibility,
            inputs=[edit_show_thinking],
            outputs=[edit_thinking_output, edit_thinking_params],
        )

        reset_btn.click(
            fn=reset_workspace,
            inputs=None,
            outputs=[
                edit_image_input,
                edit_prompt,
                edit_image_output,
                edit_thinking_output,
                edit_show_thinking,
                edit_cfg_text_scale,
                edit_cfg_img_scale,
                edit_cfg_interval,
                edit_timestep_shift,
                edit_num_timesteps,
                edit_cfg_renorm_min,
                edit_cfg_renorm_type,
                edit_max_think_token_n,
                edit_do_sample,
                edit_text_temperature,
                edit_seed,
                edit_thinking_params,
            ],
        )

        gr.on(
            triggers=[edit_btn.click, edit_prompt.submit],
            fn=edit_image,
            inputs=[
                edit_image_input,
                edit_prompt,
                edit_show_thinking,
                edit_cfg_text_scale,
                edit_cfg_img_scale,
                edit_cfg_interval,
                edit_timestep_shift,
                edit_num_timesteps,
                edit_cfg_renorm_min,
                edit_cfg_renorm_type,
                edit_max_think_token_n,
                edit_do_sample,
                edit_text_temperature,
                edit_seed,
            ],
            outputs=[edit_image_output, edit_thinking_output],
        )

        gr.HTML(
            f"""
            <div class="footer-note">
              Inference is powered by the bundled <code>modeling</code> package in this Space.
              Model weights are fetched from the Hugging Face Hub repo/path <code>{MODEL_PATH}</code> when the runtime is first initialized.
            </div>
            """
        )

    return demo


demo = build_demo()


if __name__ == "__main__":
    demo.queue(max_size=4).launch()