"""StreamOPD-4B-ST-CueGate — streaming video QA demo.

Reproduces the paper's memory-free *recent-window* inference protocol:
the model answers a question about a live video stream while seeing only the
N most recent frames sampled at 1 fps — no memory bank, no retrieval, no
KV-cache compression, no reasoning trace.

The prompt construction here mirrors the reference implementation in
https://github.com/UniX-AI-Lab/StreamOPD (Apache-2.0),
`streamopd/streaming/qwen3_5.py` :: RecentWindowQAModel:
  * each recent frame is fed as a separate image inside its own
    <|vision_start|> ... <|vision_end|> block (image tokens, not video tokens),
  * the assistant turn is opened with thinking already closed
    ("<think>\\n\\n</think>\\n\\n") so the model answers in instruct mode,
  * greedy decoding.
Frame selection mirrors `decode_video_to_chunks_qwen` with
chunk_duration = 1.0 s, fps = 1.0 and `recent_frames_only = N`.
"""

from __future__ import annotations

import spaces  # must come before torch / transformers

import math
import os
import re
import time

import av
import gradio as gr
import torch
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor

MODEL_ID = "UniX-Lab/StreamOPD-4B-ST-CueGate"

# qwen_vl_utils sampling constants, mirrored from the reference decoder.
FRAME_FACTOR = 2
FPS_MIN_FRAMES = 4
FPS_MAX_FRAMES = 768

# Vision budget. The upper bound keeps a 4K upload from exploding the prefill;
# the lower bound is the checkpoint's own preprocessor default.
MIN_PIXELS = 64 * 32 * 32          # 65_536
MAX_PIXELS = 900 * 32 * 32         # 921_600  (~900 vision tokens per frame)

CHUNK_DURATION = 1.0               # seconds per streaming chunk
SAMPLE_FPS = 1.0                   # frames sampled per second of stream

IMAGE_TOKEN = "<|image_pad|>"
VISION_START = "<|vision_start|>"
VISION_END = "<|vision_end|>"

processor = AutoProcessor.from_pretrained(
    MODEL_ID, min_pixels=MIN_PIXELS, max_pixels=MAX_PIXELS
)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID, dtype=torch.bfloat16, attn_implementation="sdpa"
)
model.eval()
model.to("cuda")


# --------------------------------------------------------------------------- #
# Video decoding — recent-window frame selection
# --------------------------------------------------------------------------- #
def _floor_by(value: float, factor: int) -> int:
    return int(math.floor(value / factor) * factor)


def _ceil_by(value: float, factor: int) -> int:
    return int(math.ceil(value / factor) * factor)


def probe_video(path: str) -> tuple[float, float, int]:
    """Return (fps, duration_seconds, total_frames) for a video file."""
    with av.open(path) as container:
        stream = container.streams.video[0]
        fps = float(stream.average_rate) if stream.average_rate else 25.0
        duration = 0.0
        if stream.duration is not None and stream.time_base:
            duration = float(stream.duration * stream.time_base)
        elif container.duration is not None:
            duration = float(container.duration) / av.time_base
        total = int(stream.frames or 0)
        if total <= 0:
            total = max(1, int(round(duration * fps)))
        if duration <= 0:
            duration = total / fps
    return fps, duration, total


def plan_window(
    path: str,
    stream_position_pct: float,
    window_chunks: int,
    sample_fps: float = SAMPLE_FPS,
    chunk_duration: float = CHUNK_DURATION,
):
    """Pick the frame indices of the `window_chunks` most recent 1-second chunks.

    Mirrors `decode_video_to_chunks_qwen` (uniform 1 fps sampling over the
    already-streamed prefix, bucketed into 1 s chunks, keep the last N buckets).
    """
    fps, duration, total_frames = probe_video(path)
    pct = min(max(float(stream_position_pct), 1.0), 100.0)
    cut_seconds = duration * pct / 100.0

    n_prefix = max(1, min(total_frames, int(round(cut_seconds * fps))))

    nframes = n_prefix / fps * sample_fps
    min_frames = _ceil_by(FPS_MIN_FRAMES, FRAME_FACTOR)
    max_frames = _floor_by(min(FPS_MAX_FRAMES, n_prefix), FRAME_FACTOR)
    nframes = min(min(max(nframes, min_frames), max_frames), n_prefix)
    nframes = max(1, _floor_by(nframes, FRAME_FACTOR))
    nframes = int(min(nframes, n_prefix))

    if nframes <= 1:
        indices = [n_prefix - 1]
    else:
        indices = [
            int(round(i * (n_prefix - 1) / (nframes - 1))) for i in range(nframes)
        ]

    buckets: dict[int, list[int]] = {}
    for idx in indices:
        buckets.setdefault(int((idx / fps) // chunk_duration), []).append(idx)

    keys = sorted(buckets)[-max(1, int(window_chunks)):]
    selected = [idx for key in keys for idx in buckets[key]]
    timestamps = [idx / fps for idx in selected]
    return selected, timestamps, fps, duration, cut_seconds


def grab_frames(path: str, indices: list[int]) -> list[Image.Image]:
    """Decode the given absolute frame indices as PIL images."""
    wanted = sorted(set(indices))
    found: dict[int, Image.Image] = {}
    with av.open(path) as container:
        stream = container.streams.video[0]
        stream.thread_type = "AUTO"
        fps = float(stream.average_rate) if stream.average_rate else 25.0
        seek_seconds = max(0.0, (wanted[0] / fps) - 2.0)
        try:
            container.seek(
                int(seek_seconds / float(stream.time_base)),
                stream=stream,
                backward=True,
            )
        except Exception:
            container.seek(0)
        remaining = list(wanted)
        last = None
        for frame in container.decode(stream):
            if frame.pts is None:
                continue
            position = int(round(float(frame.pts * stream.time_base) * fps))
            last = frame
            while remaining and position >= remaining[0]:
                found[remaining.pop(0)] = frame.to_image()
            if not remaining:
                break
        if remaining and last is not None:
            image = last.to_image()
            for idx in remaining:
                found[idx] = image
    if not found:
        raise gr.Error("Could not decode any frame from this video.")
    fallback = next(iter(found.values()))
    return [found.get(i, fallback).convert("RGB") for i in indices]


# --------------------------------------------------------------------------- #
# Prompt construction + generation (mirrors RecentWindowQAModel)
# --------------------------------------------------------------------------- #
def build_prompt(num_frames: int, question: str, thinking: bool) -> str:
    blocks = "".join(f"{VISION_START}{IMAGE_TOKEN}{VISION_END}" for _ in range(num_frames))
    tail = "<think>\n" if thinking else "<think>\n\n</think>\n\n"
    return (
        "<|im_start|>user\n"
        + blocks
        + "\n"
        + question
        + "<|im_end|>\n<|im_start|>assistant\n"
        + tail
    )


def strip_thinking(text: str) -> str:
    return re.sub(r"<think>.*?</think>\s*", "", text, flags=re.DOTALL).strip()


def _gpu_duration(
    video=None,
    question="",
    stream_position=100.0,
    window_chunks=4,
    max_new_tokens=192,
    thinking=False,
) -> int:
    """Size the ZeroGPU reservation to the requested work (measured ~2-8 s typical)."""
    return int(min(150, 12 + 1.5 * float(window_chunks) + 0.05 * float(max_new_tokens)))


@spaces.GPU(duration=_gpu_duration)
def answer_stream(
    video: str | None,
    question: str,
    stream_position: float = 100.0,
    window_chunks: int = 4,
    max_new_tokens: int = 192,
    thinking: bool = False,
) -> tuple[str, list, str]:
    """Answer a question about a video stream using only its most recent frames.

    Args:
        video: Path to the video file that plays the stream.
        question: Free-form question about what is happening in the stream right now,
            e.g. "What is the person doing with the object in their hands?".
        stream_position: Where the stream is "now", as a percentage of the video
            length (100 = the very end). Everything after this point is unseen.
        window_chunks: How many of the most recent 1-second chunks the model may see.
        max_new_tokens: Maximum number of tokens to generate.
        thinking: Emit a reasoning trace before answering (off in the paper's protocol).

    Returns:
        The model's answer, the frames it actually saw, and run details.
    """
    if not video:
        raise gr.Error("Please upload or pick a video first.")
    if not question or not question.strip():
        raise gr.Error("Please enter a question about the video.")

    indices, timestamps, fps, duration, cut_seconds = plan_window(
        video, stream_position, int(window_chunks)
    )
    frames = grab_frames(video, indices)

    prompt = build_prompt(len(frames), question.strip(), bool(thinking))
    inputs = processor(text=[prompt], images=frames, return_tensors="pt")
    inputs = {k: (v.to("cuda") if hasattr(v, "to") else v) for k, v in inputs.items()}

    prompt_len = int(inputs["input_ids"].shape[1])
    num_vision_tokens = int((inputs["input_ids"] == model.config.image_token_id).sum())

    started = time.perf_counter()
    with torch.inference_mode():
        generated = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=False,
        )
    elapsed = time.perf_counter() - started

    text = processor.tokenizer.decode(
        generated[0][prompt_len:], skip_special_tokens=True
    ).strip()
    text = strip_thinking(text) or text

    gallery = [
        (frame, f"t = {ts:.2f}s")
        for frame, ts in zip(frames, timestamps)
    ]
    details = (
        f"**Recent window:** {len(frames)} frame(s) at "
        f"{', '.join(f'{t:.2f}s' for t in timestamps)}  \n"
        f"**Stream cut-off:** {cut_seconds:.2f}s of {duration:.2f}s "
        f"(source {fps:.1f} fps)  \n"
        f"**Prompt:** {prompt_len} tokens ({num_vision_tokens} vision)  \n"
        f"**Generation:** {elapsed:.2f}s"
    )
    return text, gallery, details


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #
EXAMPLES = [
    [
        "examples/barista_frothing.mp4",
        "What is the person doing with the metal pitcher right now?",
        100,
    ],
    [
        "examples/slicing_veggie.mp4",
        "Describe what has appeared on the cutting board by now.",
        100,
    ],
    [
        "examples/basketball_dribble.mp4",
        "What is the man on the court doing at this moment?",
        60,
    ],
    [
        "examples/train_arriving.mp4",
        "Briefly describe what is happening in the scene right now.",
        100,
    ],
]

CSS = """
#answer_box textarea { font-size: 1.05rem; }
"""

with gr.Blocks(title="StreamOPD Streaming Video QA") as demo:
    gr.Markdown(
        """
# StreamOPD-4B — Streaming Video QA

[**StreamOPD: A Post-Training Recipe with Spatio-Temporal Cue Gating for Streaming Video
Understanding**](https://huggingface.co/papers/2608.16320) ·
[model](https://huggingface.co/UniX-Lab/StreamOPD-4B-ST-CueGate) ·
[code](https://github.com/UniX-AI-Lab/StreamOPD)

Ask a question about a video **as if it were a live stream**. Following the paper's
memory-free *recent-window* protocol, the model only ever sees the **4 most recent frames
at 1 fps** — no memory bank, no retrieval, no KV-cache compression and no reasoning trace.
Move the *stream position* slider to pretend the stream is still running, and everything
after that point stays unseen.
"""
    )

    with gr.Row():
        with gr.Column(scale=1):
            video = gr.Video(label="Video stream", height=320)
            stream_position = gr.Slider(
                1,
                100,
                value=100,
                step=1,
                label="Stream position (% of the video)",
                info="Where the stream is 'now'. 100% = the very end of the clip.",
            )
            question = gr.Textbox(
                label="Question",
                lines=4,
                placeholder=(
                    "Ask anything about what is happening right now, e.g.\n"
                    "What is the person doing with the object in their hands?"
                ),
            )
            run = gr.Button("Answer", variant="primary")

            with gr.Accordion("Advanced", open=False):
                window_chunks = gr.Slider(
                    1,
                    16,
                    value=4,
                    step=1,
                    label="Recent window (1-second chunks)",
                    info="4 is the paper's protocol.",
                )
                max_new_tokens = gr.Slider(
                    8, 512, value=192, step=8, label="Max new tokens"
                )
                thinking = gr.Checkbox(
                    value=False,
                    label="Enable thinking",
                    info="Off matches the reported instruct-mode results.",
                )

        with gr.Column(scale=1):
            answer = gr.Textbox(label="Answer", lines=4, elem_id="answer_box")
            frames_view = gr.Gallery(
                label="Frames the model actually saw",
                columns=4,
                height=200,
                object_fit="contain",
            )
            details = gr.Markdown()

    gr.Examples(
        examples=EXAMPLES,
        inputs=[video, question, stream_position],
        outputs=[answer, frames_view, details],
        fn=answer_stream,
        cache_examples=True,
        cache_mode="lazy",
        label="Examples",
    )

    gr.Markdown(
        """
Example clips come from
[`linoyts/repo-to-space-example-videos`](https://huggingface.co/datasets/linoyts/repo-to-space-example-videos)
(CC0-1.0). Prompt construction and frame selection follow
[`streamopd/streaming/qwen3_5.py`](https://github.com/UniX-AI-Lab/StreamOPD/blob/main/streamopd/streaming/qwen3_5.py)
(Apache-2.0). Answers are limited by construction to evidence inside the recent window.
"""
    )

    gr.on(
        triggers=[run.click, question.submit],
        fn=answer_stream,
        inputs=[video, question, stream_position, window_chunks, max_new_tokens, thinking],
        outputs=[answer, frames_view, details],
        api_name="answer_stream",
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)