"""VideoMaMa++ demo: SAM3 mask track -> VideoMaMa++ alpha matte.

The VideoMaMa++ weights and model code live in a private Hugging Face repo.
They are downloaded at startup with the HF_TOKEN Space secret into the Hub
cache (outside the app folder) and are never served to visitors.
"""

import os
import sys
import time
import uuid
from pathlib import Path

try:
    import spaces  # ZeroGPU; must be imported before torch touches CUDA
    GPU = spaces.GPU
except ImportError:  # local run
    def GPU(fn=None, **_):
        return fn if fn is not None else (lambda f: f)

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

from painter import paint_mask, paint_points

PRIVATE_REPO = os.environ.get("VIDEOMAMAPP_REPO", "SammyLim/VideoMaMa-PlusPlus")
HF_TOKEN = os.environ.get("HF_TOKEN")
APP_DIR = Path(__file__).resolve().parent
OUT_DIR = Path(os.environ.get("VIDEOMAMAPP_OUT", "/tmp/videomamapp_outputs"))
# Wan2.1 base files (public) are cached here, outside the app folder.
os.environ.setdefault("DIFFSYNTH_MODEL_BASE_PATH", str(Path.home() / ".cache" / "videomamapp_models"))
os.environ.setdefault("DIFFSYNTH_DOWNLOAD_SOURCE", "huggingface")

MAX_FRAMES = 48          # hard cap per run
DEFAULT_FRAMES = 24
MAX_LONG_SIDE = 1536     # the paper's processing size; larger videos are capped here
MIN_LONG_SIDE = 256
SCALES = {"100% (original)": 1.0, "75%": 0.75, "50%": 0.5, "25%": 0.25}
DEFAULT_SCALE = "50%"
GREEN = np.array([120, 255, 155], np.float32)

# ---------------------------------------------------------------- models
print("[startup] downloading VideoMaMa++ (private repo)...", flush=True)
_local = os.environ.get("VIDEOMAMAPP_LOCAL_DIR")  # local testing only
if not _local and not HF_TOKEN:
    raise RuntimeError(
        "HF_TOKEN is not set. Add a read token with access to the private model repo "
        f"{PRIVATE_REPO} and to facebook/sam3 as a Space secret named HF_TOKEN.")
_private = Path(_local) if _local else Path(snapshot_download(PRIVATE_REPO, token=HF_TOKEN))
sys.path.insert(0, str(_private / "code"))
import videomamapp_inference as vmpp  # noqa: E402
from sam3_tracker import Sam3Tracker  # noqa: E402

print("[startup] loading VideoMaMa++ ...", flush=True)
PIPE = vmpp.load_pipeline(_private / "videomamapp_dit.safetensors", _private / "prompt_context.pt", device="cpu")
PIPE.to("cuda")
print("[startup] loading SAM3 ...", flush=True)
SAM3 = Sam3Tracker(device="cuda", token=HF_TOKEN)
print("[startup] ready", flush=True)


# ---------------------------------------------------------------- video io
def read_video(path, max_frames, stride):
    cap = cv2.VideoCapture(str(path))
    fps = cap.get(cv2.CAP_PROP_FPS) or 24.0
    orig = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    frames, i = [], 0
    while len(frames) < max_frames:
        ok, f = cap.read()
        if not ok:
            break
        if i % stride == 0:
            f = cv2.cvtColor(f, cv2.COLOR_BGR2RGB)
            h, w = f.shape[:2]
            s = MAX_LONG_SIDE / max(h, w)
            if s < 1:
                f = cv2.resize(f, (round(w * s / 2) * 2, round(h * s / 2) * 2), interpolation=cv2.INTER_AREA)
            frames.append(f)
        i += 1
    cap.release()
    return frames, fps / stride, orig


def processing_size(orig, scale):
    """(W, H) the model runs at: `scale` of the original size, where the
    original's long side is first capped at MAX_LONG_SIDE; floored at
    MIN_LONG_SIDE, both sides multiples of 16."""
    w, h = orig
    long_side = max(MIN_LONG_SIDE, min(max(w, h), MAX_LONG_SIDE) * SCALES[scale])
    f = long_side / max(w, h)
    return max(16, round(w * f / 16) * 16), max(16, round(h * f / 16) * 16)


def gpu_seconds(n, size):
    """Estimated GPU time. VideoMaMa++ per frame (6-frame chunks, overlap 2) fitted
    to A100 runs at 512-1536 px (0.15-2.78 s/frame), plus SAM3 at ~0.1 s/frame."""
    mp = size[0] * size[1] / 1e6
    return n * (0.1 + 1.135 * mp + 0.718 * mp * mp)


def write_video(frames, path, fps):
    """H.264 mp4 that browsers can play."""
    h, w = frames[0].shape[:2]
    W, H = w - w % 2, h - h % 2
    cmd = [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error",
           "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{W}x{H}", "-r", f"{fps:.4f}", "-i", "-",
           "-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "18", "-movflags", "+faststart", str(path)]
    p = subprocess.Popen(cmd, stdin=subprocess.PIPE)
    for f in frames:
        if f.ndim == 2:
            f = np.repeat(f[..., None], 3, 2)
        p.stdin.write(np.ascontiguousarray(f[:H, :W]).tobytes())
    p.stdin.close()
    p.wait()
    return str(path)


# ---------------------------------------------------------------- GPU work
@GPU(duration=15)
def gpu_first_frame_mask(frame, points, labels):
    return SAM3.first_frame_mask(frame, points, labels)


def _run_duration(frames, points, labels, text, max_size):
    # GPU seconds to reserve. Visitors are refused when this exceeds their
    # remaining ZeroGPU quota, so keep it close to the real cost (model
    # loading + the estimate, with 50% headroom).
    h, w = frames[0].shape[:2]
    return int(min(240, 20 + 1.5 * gpu_seconds(len(frames), (w, h))))


@GPU(duration=_run_duration)
def gpu_track_and_matte(frames, points, labels, text, max_size):
    t0 = time.time()
    if text:
        masks = SAM3.track_text(frames, text)
    else:
        masks = SAM3.track_points(frames, points, labels)
    torch.cuda.synchronize()
    t1 = time.time()
    mattes = vmpp.matte(PIPE, frames, masks, max_size=int(max_size))
    t2 = time.time()
    timing = {"gpu": torch.cuda.get_device_name(), "sam3": t1 - t0, "matte": t2 - t1}
    return masks, mattes, timing


# ---------------------------------------------------------------- handlers
def on_load(video, n_frames, stride, scale, state):
    if video is None:
        raise gr.Error("Upload a video or pick an example first.")
    frames, fps, orig = read_video(video, int(n_frames), int(stride))
    if not frames:
        raise gr.Error("Could not read any frames from this video.")
    state = {"frames": frames, "fps": fps, "orig": orig, "points": [], "labels": [], "mask": None}
    msg = f"Loaded {len(frames)} frames ({orig[0]}×{orig[1]}). Click on the object, or use a text prompt."
    return state, Image.fromarray(frames[0]), msg, on_scale(state, scale)


def on_scale(state, scale):
    if not state or "frames" not in state:
        return ""
    size = processing_size(state["orig"], scale)
    sec = gpu_seconds(len(state["frames"]), size)
    return (f"Runs at **{size[0]}×{size[1]}** on {len(state['frames'])} frames: "
            f"about **{sec:.0f} s** of GPU time (A100 estimate).")


def on_click(state, kind, evt: gr.SelectData):
    if not state or "frames" not in state:
        raise gr.Error("Load a video first.")
    x, y = evt.index[0], evt.index[1]
    state["points"].append([x, y])
    state["labels"].append(1 if kind == "Positive" else 0)
    frame = state["frames"][0]
    state["mask"] = gpu_first_frame_mask(frame, state["points"], state["labels"])
    view = paint_points(paint_mask(frame, state["mask"]), state["points"], state["labels"])
    return state, Image.fromarray(view)


def on_clear(state):
    if not state or "frames" not in state:
        return state, None
    state.update(points=[], labels=[], mask=None)
    return state, Image.fromarray(state["frames"][0])


def on_run(state, text, scale, progress=gr.Progress()):
    if not state or "frames" not in state:
        raise gr.Error("Load a video first.")
    text = (text or "").strip()
    if not text and not state["points"]:
        raise gr.Error("Click on the object in the first frame, or type a text prompt.")
    frames, points = state["frames"], state["points"]
    # Work at the processing resolution: SAM3, VideoMaMa++ and the output
    # videos all use the downscaled frames.
    h, w = frames[0].shape[:2]
    size = processing_size(state["orig"], scale)
    max_size = max(size)
    if size != (w, h):
        interp = cv2.INTER_AREA if size[0] < w else cv2.INTER_CUBIC
        frames = [cv2.resize(f, size, interpolation=interp) for f in frames]
        points = [[x * size[0] / w, y * size[1] / h] for x, y in points]
    progress(0.1, desc=f"SAM3 + VideoMaMa++ on {len(frames)} frames ...")
    t0 = time.time()
    masks, mattes, timing = gpu_track_and_matte(frames, points, state["labels"], text, max_size)
    took = time.time() - t0

    run = OUT_DIR / uuid.uuid4().hex[:10]
    run.mkdir(parents=True, exist_ok=True)
    fps = state["fps"]
    comps = []
    for f, a in zip(frames, mattes):
        al = (a.astype(np.float32) / 255.0)[..., None]
        comps.append((f.astype(np.float32) * al + GREEN * (1 - al)).astype(np.uint8))
    alpha_v = write_video(mattes, run / "alpha.mp4", fps)
    green_v = write_video(comps, run / "greenscreen.mp4", fps)
    mask_v = write_video(masks, run / "sam3_mask.mp4", fps)
    write = time.time() - t0 - took
    wait = took - timing["sam3"] - timing["matte"]
    print(f"[run] {len(frames)} frames {frames[0].shape[1]}x{frames[0].shape[0]} ({scale}) on {timing['gpu']}: "
          f"SAM3 {timing['sam3']:.1f}s, VideoMaMa++ {timing['matte']:.1f}s, GPU wait/transfer {wait:.1f}s, "
          f"encode {write:.1f}s", flush=True)
    return alpha_v, green_v, mask_v, (
        f"Done: {len(frames)} frames at {size[0]}×{size[1]} in {took + write:.0f} s "
        f"(SAM3 {timing['sam3']:.0f} s, VideoMaMa++ {timing['matte']:.0f} s).")


def on_new_video():
    return None, None, "", None, None, None, ""


# ---------------------------------------------------------------- UI
CSS = """
#title {text-align: center; margin-bottom: 0;}
#links {text-align: center;}
#run_btn {font-size: 1.1rem;}
#note {background: var(--color-accent-soft); border-left: 4px solid var(--color-accent);
       border-radius: 6px; padding: 10px 14px;}
"""

with gr.Blocks(title="VideoMaMa++ demo") as demo:
    gr.Markdown("# VideoMaMa++: Temporally Consistent Video Matting via Preserve-and-Refine Tokens", elem_id="title")
    gr.Markdown(
        "**NeurIPS 2026** · [Project page](https://cvlab-kaist.github.io/VideoMaMa-PlusPlus) · "
        "[Paper](https://cvlab-kaist.github.io/VideoMaMa-PlusPlus/source/_NeurIPS_2026__LongMaMa.pdf) · "
        "[GitHub](https://github.com/cvlab-kaist/VideoMaMa-PlusPlus)",
        elem_id="links",
    )
    gr.Markdown(
        "1. Upload a video (or pick an example) and press **Load video**. "
        "2. Click the object on the first frame (or type what it is, e.g. `person`). "
        "3. Press **Run matting**: SAM3 tracks the object and VideoMaMa++ turns the mask track into an alpha matte."
    )
    gr.Markdown(
        f"**Speed vs. quality:** to keep runs short, the demo starts with a reduced setting of "
        f"**{DEFAULT_FRAMES} frames at {DEFAULT_SCALE} resolution**, so these results are lower in quality than "
        "VideoMaMa++ can produce. For the best mattes, set **Processing resolution** to **100% (original)**. "
        "That takes longer and uses more of your GPU quota.",
        elem_id="note",
    )

    state = gr.State(None)
    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Input video", sources=["upload"])
            with gr.Row():
                n_frames = gr.Slider(6, MAX_FRAMES, value=DEFAULT_FRAMES, step=1, label="Frames to process",
                                     info="Taken from the start of the video. Press Load video again after changing.")
                stride = gr.Slider(1, 4, value=1, step=1, label="Frame stride",
                                   info="Use every k-th frame to cover a longer span at the same cost.")
            load_btn = gr.Button("Load video", variant="secondary")
            status = gr.Markdown("")
        with gr.Column():
            first_frame = gr.Image(label="First frame: click the object", type="pil", interactive=False)
            with gr.Row():
                click_kind = gr.Radio(["Positive", "Negative"], value="Positive", label="Click type")
                clear_btn = gr.Button("Clear clicks")
            text_prompt = gr.Textbox(label="Or a text prompt (overrides clicks)", placeholder="e.g. person, dog, hair")
            with gr.Accordion("Settings", open=True):
                scale = gr.Radio(list(SCALES), value=DEFAULT_SCALE, label="Processing resolution",
                                 info=f"Share of the original size; 100% is capped at {MAX_LONG_SIDE} px on the long side, "
                                      "the paper's setting. Lower is much faster; fewer frames also helps.")
                estimate = gr.Markdown("")
            run_btn = gr.Button("Run matting", variant="primary", elem_id="run_btn")

    with gr.Row():
        alpha_out = gr.Video(label="Alpha matte (VideoMaMa++)", autoplay=True, loop=True)
        green_out = gr.Video(label="Greenscreen composite", autoplay=True, loop=True)
        mask_out = gr.Video(label="SAM3 mask track (input guide)", autoplay=True, loop=True)

    examples = sorted(str(p) for p in (APP_DIR / "samples").glob("*.mp4"))
    if examples:
        gr.Examples(examples=examples, inputs=[video_in], label="Examples")

    load_btn.click(on_load, [video_in, n_frames, stride, scale, state], [state, first_frame, status, estimate])
    scale.change(on_scale, [state, scale], [estimate])
    first_frame.select(on_click, [state, click_kind], [state, first_frame])
    clear_btn.click(on_clear, [state], [state, first_frame])
    run_btn.click(on_run, [state, text_prompt, scale], [alpha_out, green_out, mask_out, status])
    video_in.change(on_new_video, None, [state, first_frame, status, alpha_out, green_out, mask_out, estimate])


if __name__ == "__main__":
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    demo.queue(max_size=20).launch(
        css=CSS,
        allowed_paths=[str(OUT_DIR), str(APP_DIR / "samples")],
        # Never serve the private weights or the model cache.
        blocked_paths=[str(_private), os.environ["DIFFSYNTH_MODEL_BASE_PATH"], str(Path.home() / ".cache")],
    )