"""Rosetta — Composable Native Multimodal (Tencent Hunyuan / HKUST).

A ZeroGPU Gradio demo for tencent/Rosetta-inference (Rosetta-3.8B-A1B):
  * Text -> Image generation
  * Image + Text understanding (VLM)

The upstream repo (github.com/Lxiangyue/Rosetta) is a research/eval codebase
that places tensors on CUDA eagerly and uses accelerate device_map dispatch.
To stay compatible with ZeroGPU's fork model, we build + load the model the
first time a @spaces.GPU function runs (inside the GPU worker where real CUDA
exists), then keep it resident on that warm worker.
"""

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  (must precede torch/CUDA imports)

import struct
import threading
import urllib.request
import zlib
from pathlib import Path

import gradio as gr
import torch
from huggingface_hub import snapshot_download

# --------------------------------------------------------------------------- #
# Paths / config
# --------------------------------------------------------------------------- #
REPO_ID = "tencent/Rosetta-inference"
DATA_DIR = Path(os.environ.get("ROSETTA_DATA", "/tmp/rosetta_data"))
CKPT_DIR = DATA_DIR / "checkpoints" / "Rosetta-3.8B-A1B" / "hf_weights"
ASSETS_DIR = DATA_DIR / "public_assets"
CONFIG_PATH = Path(__file__).parent / "rosetta.yaml"
GEN_CFG_DIR = ASSETS_DIR / "generation_configs"

# The rosetta package reads asset locations from these env vars at import time.
os.environ["ASSETS_BASE"] = str(ASSETS_DIR)

ZIP_URL = f"https://huggingface.co/{REPO_ID}/resolve/main/public_assets.zip"

_sampler = None
_load_lock = threading.Lock()


# --------------------------------------------------------------------------- #
# Asset download (runs once at startup, on the CPU disk)
# --------------------------------------------------------------------------- #
def _get_range(a, b):
    for attempt in range(6):
        try:
            req = urllib.request.Request(ZIP_URL, headers={"Range": f"bytes={a}-{b}"})
            return urllib.request.urlopen(req, timeout=180).read()
        except Exception as e:  # noqa: BLE001
            print(f"  range retry {attempt}: {e}", flush=True)
    raise RuntimeError("range request failed")


def _parse_zip64_extra(extra, usize, csize, lho):
    off = 0
    while off + 4 <= len(extra):
        hid, dsize = struct.unpack("<HH", extra[off:off + 4])
        data = extra[off + 4:off + 4 + dsize]
        if hid == 0x0001:
            p = 0
            if usize == 0xFFFFFFFF:
                usize = struct.unpack("<Q", data[p:p + 8])[0]; p += 8
            if csize == 0xFFFFFFFF:
                csize = struct.unpack("<Q", data[p:p + 8])[0]; p += 8
            if lho == 0xFFFFFFFF:
                lho = struct.unpack("<Q", data[p:p + 8])[0]; p += 8
            break
        off += 4 + dsize
    return usize, csize, lho


def _zip_entries():
    size = int(urllib.request.urlopen(
        urllib.request.Request(ZIP_URL, method="HEAD")).headers["Content-Length"])
    tail = _get_range(size - 131072, size - 1)
    i = tail.rfind(b"PK\x06\x07")
    cd_off = struct.unpack("<Q", tail[i + 8:i + 16])[0]
    z = _get_range(cd_off, cd_off + 55)
    cdsize = struct.unpack("<Q", z[40:48])[0]
    cdoff = struct.unpack("<Q", z[48:56])[0]
    cd = _get_range(cdoff, cdoff + cdsize - 1)
    off, out = 0, []
    while off < len(cd) and cd[off:off + 4] == b"PK\x01\x02":
        comp = struct.unpack("<H", cd[off + 10:off + 12])[0]
        csize = struct.unpack("<I", cd[off + 20:off + 24])[0]
        usize = struct.unpack("<I", cd[off + 24:off + 28])[0]
        nlen = struct.unpack("<H", cd[off + 28:off + 30])[0]
        elen = struct.unpack("<H", cd[off + 30:off + 32])[0]
        clen = struct.unpack("<H", cd[off + 32:off + 34])[0]
        lho = struct.unpack("<I", cd[off + 42:off + 46])[0]
        name = cd[off + 46:off + 46 + nlen].decode("utf-8", "replace")
        extra = cd[off + 46 + nlen:off + 46 + nlen + elen]
        usize, csize, lho = _parse_zip64_extra(extra, usize, csize, lho)
        out.append((name, usize, csize, comp, lho))
        off += 46 + nlen + elen + clen
    return out


def _extract(name, csize, comp, lho, dest):
    lh = _get_range(lho, lho + 29)
    assert lh[:4] == b"PK\x03\x04", f"bad local header for {name}"
    nlen = struct.unpack("<H", lh[26:28])[0]
    elen = struct.unpack("<H", lh[28:30])[0]
    ds = lho + 30 + nlen + elen
    raw = _get_range(ds, ds + csize - 1)
    data = raw if comp == 0 else zlib.decompress(raw, -15)
    dest.parent.mkdir(parents=True, exist_ok=True)
    dest.write_bytes(data)


def ensure_data():
    """Download main checkpoint + selectively extract model assets from the
    18 GB public_assets.zip (skipping the ~56k evaluation-dataset files)."""
    index = CKPT_DIR / "model.safetensors.index.json"
    if not index.exists():
        print("=== downloading Rosetta-3.8B-A1B checkpoint ===", flush=True)
        snapshot_download(
            REPO_ID, repo_type="model",
            allow_patterns=["checkpoints/Rosetta-3.8B-A1B/hf_weights/*"],
            local_dir=str(DATA_DIR), max_workers=8,
        )
    vae = ASSETS_DIR / "image_encoder" / "flux2-vae" / "model.pt"
    if not vae.exists():
        print("=== extracting shared assets (VAE / ViT / tokenizer) ===", flush=True)
        for name, usize, csize, comp, lho in _zip_entries():
            if name.startswith("public_assets/evaluation") or name.endswith("/"):
                continue
            dest = DATA_DIR / name
            if dest.exists() and dest.stat().st_size == usize:
                continue
            print(f"  extract {name}", flush=True)
            _extract(name, csize, comp, lho, dest)
    print("=== data ready ===", flush=True)


# Kick off the download at import time (CPU-only work, safe under the hijack).
ensure_data()


# --------------------------------------------------------------------------- #
# Model loading (inside the GPU worker)
# --------------------------------------------------------------------------- #
def get_sampler():
    global _sampler
    if _sampler is not None:
        return _sampler
    with _load_lock:
        if _sampler is not None:
            return _sampler
        from evaluation.multimodal_sampler import MultimodalSampler
        print("=== building + loading Rosetta model on GPU ===", flush=True)
        _sampler = MultimodalSampler.from_pretrained(
            ckpt_path=str(CKPT_DIR),
            config_path=str(CONFIG_PATH),
            device=0,
            extra_args=[
                "--sequence-template", "instruct",
                # A valid generation config is required at load time; the ckpt
                # dir ships none, so point at the shipped T2I config.
                "--generation-config", str(GEN_CFG_DIR / "qwen3_0.6b_t2i_eval.json"),
            ],
        )
        print("=== model ready ===", flush=True)
    return _sampler


def _apply_gen_config(model, cfg_name):
    """Load one of the shipped generation configs onto the model."""
    model.load_generation_config(str(GEN_CFG_DIR / cfg_name))


# --------------------------------------------------------------------------- #
# Inference
# --------------------------------------------------------------------------- #
@spaces.GPU(duration=180)
def generate_image(prompt, image_ratio, steps, guidance, seed, progress=gr.Progress(track_tqdm=True)):
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")
    sampler = get_sampler()
    model = sampler.model
    _apply_gen_config(model, "qwen3_0.6b_t2i_eval.json")
    model.generation_config.diff_infer_steps = int(steps)
    model.generation_config.diff_guidance_scale = float(guidance)

    seed = int(seed)
    if seed < 0:
        seed = torch.randint(0, 2**31 - 1, (1,)).item()

    with torch.no_grad():
        outputs = model.generate_image(
            prompt=prompt.strip(),
            seed=seed,
            image_size=image_ratio,
            bot_task="image",
        )
    images = outputs.images
    img = images[0] if isinstance(images, list) else images
    return img, seed


@spaces.GPU(duration=120)
def understand_image(image, question, max_new_tokens, progress=gr.Progress(track_tqdm=True)):
    if image is None:
        raise gr.Error("Please upload an image.")
    if not question or not question.strip():
        question = "Describe this image in detail."
    sampler = get_sampler()
    model = sampler.model
    _apply_gen_config(model, "qwen3_0.6b_mmu_eval_instruct.json")
    model.generation_config.max_new_tokens = int(max_new_tokens)

    with torch.no_grad():
        # Mirror MultimodalSampler.run(): pass prompt + image directly.
        inputs = model.prepare_model_inputs(
            prompt=question.strip(),
            image=[image.convert("RGB")],
            mode="gen_text",
            bot_task="auto",
        )
        outputs = model.generate(
            **inputs, decode_text=True, skip_special_tokens=True, verbose=0,
        )
    texts = outputs.texts
    text = texts[0] if isinstance(texts, (list, tuple)) else texts
    while isinstance(text, (list, tuple)):
        text = text[0]
    return str(text).strip()


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #
RATIOS = ["1:1", "3:4", "4:3", "9:16", "16:9"]

with gr.Blocks(theme=gr.themes.Citrus(), title="Rosetta Multimodal") as demo:
    gr.Markdown(
        """
        # 🪨 Rosetta — Composable Native Multimodal
        Unified text↔image model **[tencent/Rosetta-inference](https://huggingface.co/tencent/Rosetta-inference)**
        (Rosetta-3.8B-A1B, Tencent Hunyuan × HKUST). One backbone does both
        **text-to-image generation** and **image understanding**.
        """
    )

    with gr.Tab("🎨 Text → Image"):
        with gr.Row():
            with gr.Column():
                t2i_prompt = gr.Textbox(
                    label="Prompt", lines=3,
                    placeholder="A serene Japanese garden with a red maple tree beside a koi pond, soft morning light",
                )
                t2i_btn = gr.Button("Generate", variant="primary")
                with gr.Accordion("Advanced options", open=False):
                    t2i_ratio = gr.Radio(RATIOS, value="1:1", label="Aspect ratio")
                    t2i_steps = gr.Slider(10, 100, value=50, step=1, label="Diffusion steps")
                    t2i_guidance = gr.Slider(1.0, 10.0, value=3.5, step=0.1, label="Guidance scale")
                    t2i_seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
            with gr.Column():
                t2i_out = gr.Image(label="Generated image", type="pil")
                t2i_used_seed = gr.Number(label="Seed used", interactive=False)

        t2i_btn.click(
            fn=generate_image,
            inputs=[t2i_prompt, t2i_ratio, t2i_steps, t2i_guidance, t2i_seed],
            outputs=[t2i_out, t2i_used_seed],
        )
        gr.Examples(
            examples=[
                ["A serene Japanese garden with a red maple tree beside a koi pond, soft morning light"],
                ["A cute corgi puppy wearing tiny sunglasses, studio photo, sharp focus"],
                ["An astronaut riding a horse on Mars, cinematic, highly detailed"],
                ["A cozy bookstore cafe interior, warm lighting, watercolor illustration"],
            ],
            inputs=[t2i_prompt],
            cache_examples=False,
        )

    with gr.Tab("🖼️ Image Understanding"):
        with gr.Row():
            with gr.Column():
                vlm_image = gr.Image(label="Image", type="pil")
                vlm_question = gr.Textbox(
                    label="Question", lines=2,
                    placeholder="Describe this image in detail.",
                )
                vlm_btn = gr.Button("Ask", variant="primary")
                with gr.Accordion("Advanced options", open=False):
                    vlm_max_tokens = gr.Slider(32, 512, value=256, step=8, label="Max new tokens")
            with gr.Column():
                vlm_out = gr.Textbox(label="Answer", lines=12)

        vlm_btn.click(
            fn=understand_image,
            inputs=[vlm_image, vlm_question, vlm_max_tokens],
            outputs=[vlm_out],
        )
        gr.Examples(
            examples=[
                ["examples/cat_tabby.jpg", "Describe this image in detail."],
                ["examples/bird_kingfisher.jpg", "What kind of bird is this and what is it doing?"],
                ["examples/cake.jpg", "What food is shown and how is it decorated?"],
            ],
            inputs=[vlm_image, vlm_question],
            cache_examples=False,
        )

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