"""
Object-centric procedural video QA with temporal grounding, using
BrightGuo/ProcObject-Qwen3-VL-4B (Qwen3-VL-4B-Instruct fine-tuned on ProcObject-10K,
NeurIPS 2026 Evaluations & Datasets Track).

The inference protocol replicates the authors' released benchmark harness
(https://github.com/WenliangGuo/ProcObject-10K) exactly:
  * 48 frames sampled uniformly over the clip (the fine-tuned model's training
    frame budget), resized to 512x512;
  * each frame is sent as an image preceded by "[Timestamp: <t>s]" (seconds from
    the clip start), then the question;
  * the system prompt from benchmark/prompts/system_prompt.txt;
  * greedy decoding (temperature 0);
  * the model returns {"answer": "...", "evidence": [[start, end], ...]} JSON,
    parsed with the same fallback logic as benchmark/utils.parse_model_response.
"""

import json
import os
import re
import time
from typing import Any, Dict, List, Optional, Tuple

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

import spaces  # noqa: E402  (must come before torch)
import torch  # noqa: E402

import cv2  # noqa: E402
import gradio as gr  # noqa: E402
from PIL import Image, ImageDraw  # noqa: E402
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration  # noqa: E402

MODEL_ID = "BrightGuo/ProcObject-Qwen3-VL-4B"
NUM_FRAMES = 48                    # the fine-tuned model's frame budget
RESIZE = (512, 512)                # frame size used by the benchmark protocol
MAX_NEW_TOKENS = 1024              # same as the harness's --max_tokens

SYSTEM_PROMPT = """You are a video question-answering assistant specializing in procedural, object-centric videos.

You will be given a series of uniformly sampled frames from a video clip. Each frame is labeled with its timestamp in seconds relative to the start of the clip (e.g., [Timestamp: 0.0s] is the first moment, [Timestamp: 45.2s] is 45.2 seconds into the clip). Use these timestamps to localize events temporally.

Your task:
1. Carefully analyze all provided frames and their timestamps.
2. Answer the question in detail, describing what you observe across the video frames.
3. Identify the relevant time intervals as [start_second, end_second] pairs using the frame timestamps. The evidence intervals should correspond to the timestamps shown in the frames. For example, if the relevant action occurs between frames labeled [Timestamp: 10.0s] and [Timestamp: 25.0s], report evidence as [10, 25].

Return your response ONLY as a valid JSON object with exactly these two fields:
- "answer": a detailed string describing what you observed to answer the question
- "evidence": a list of [start_second, end_second] pairs (integers) indicating the relevant temporal segments based on the frame timestamps

Example output: (json format)
{"answer": "The apple is first visible on the cutting board and is then sliced into pieces.", "evidence": [[10, 25], [40, 55]]}

Do not include any text outside the JSON object. Do not wrap the JSON in markdown code fences."""

processor = AutoProcessor.from_pretrained(MODEL_ID, min_pixels=256 * 28 * 28, max_pixels=512 * 512)
model = Qwen3VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda").eval()


# ─────────────────────────────────────────────────────────────────────────────
# Frame sampling — identical to benchmark/utils.extract_frames (cv2-based,
# uniform over the whole clip, resized to 512x512, timestamps 0-based).
# ─────────────────────────────────────────────────────────────────────────────

def extract_frames(video_path: str, num_frames: int = NUM_FRAMES,
                   resize: Optional[Tuple[int, int]] = RESIZE) -> Tuple[List[Dict[str, Any]], float]:
    """Uniformly sample num_frames frames from the video, with 0-based timestamps.

    Returns (frames, clip_duration); each frame is {"image": PIL, "timestamp": float}.
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f"Cannot open video: {video_path}")

    fps = cap.get(cv2.CAP_PROP_FPS)
    n_total = cap.get(cv2.CAP_PROP_FRAME_COUNT)
    clip_duration = (n_total / fps) if (fps and fps > 0 and n_total and n_total > 0) else 0.0
    if clip_duration <= 0:  # some muxers report no frame count — seek to the end instead
        cap.set(cv2.CAP_PROP_POS_AVI_RATIO, 1)
        clip_duration = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000.0
        cap.set(cv2.CAP_PROP_POS_MSEC, 0)
    if clip_duration <= 0 or fps <= 0:
        cap.release()
        return [], 0.0

    if num_frames == 1:
        timestamps = [clip_duration / 2]
    else:
        timestamps = [clip_duration * i / (num_frames - 1) for i in range(num_frames)]

    frames: List[Dict[str, Any]] = []
    for ts in timestamps:
        cap.set(cv2.CAP_PROP_POS_MSEC, ts * 1000)
        ret, frame = cap.read()
        back = 1
        while not ret and back <= 5:  # cut clips can end before the last timestamp
            cap.set(cv2.CAP_PROP_POS_MSEC, max(0.0, ts - back / fps) * 1000)
            ret, frame = cap.read()
            back += 1
        if not ret:
            continue
        img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        if resize:
            img = img.resize(resize, Image.LANCZOS)
        frames.append({"image": img, "timestamp": round(ts, 2)})

    cap.release()
    return frames, clip_duration


# ─────────────────────────────────────────────────────────────────────────────
# Response parsing — identical to benchmark/utils.parse_model_response.
# ─────────────────────────────────────────────────────────────────────────────

def parse_model_response(raw_text: str) -> Dict[str, Any]:
    text = raw_text.strip()
    text = re.sub(r"^```(?:json)?\s*", "", text)
    text = re.sub(r"\s*```$", "", text)
    text = text.strip()

    try:
        parsed = json.loads(text)
        return {"answer": str(parsed.get("answer", "")), "evidence": parsed.get("evidence", [])}
    except json.JSONDecodeError:
        pass

    match = re.search(r'\{[^{}]*"answer"[^{}]*\}', text, re.DOTALL)
    if match:
        try:
            parsed = json.loads(match.group())
            return {"answer": str(parsed.get("answer", "")), "evidence": parsed.get("evidence", [])}
        except json.JSONDecodeError:
            pass

    return {"answer": text, "evidence": []}


# ─────────────────────────────────────────────────────────────────────────────
# Evidence timeline visualization
# ─────────────────────────────────────────────────────────────────────────────

def _clean_intervals(evidence: List, duration: float) -> List[Tuple[float, float]]:
    spans = []
    for item in evidence or []:
        try:
            s, e = float(item[0]), float(item[1])
        except (TypeError, ValueError, IndexError):
            continue
        s, e = max(0.0, min(s, duration)), max(0.0, min(e, duration))
        if e > s:
            spans.append((s, e))
    return spans


def render_timeline(duration: float, spans: List[Tuple[float, float]]) -> Image.Image:
    """A thin timeline bar with the grounded evidence spans highlighted."""
    W, H = 780, 92
    pad_l, pad_r, pad_t, bar_h = 14, 14, 38, 28
    img = Image.new("RGB", (W, H), (249, 249, 247))
    d = ImageDraw.Draw(img)

    track = (222, 222, 218)
    d.rectangle([pad_l, pad_t, W - pad_r, pad_t + bar_h], fill=track)
    if duration > 0:
        for (s, e) in spans:
            x0 = pad_l + (s / duration) * (W - pad_l - pad_r)
            x1 = pad_l + (e / duration) * (W - pad_l - pad_r)
            d.rectangle([x0, pad_t, max(x0 + 2.0, x1), pad_t + bar_h], fill=(224, 78, 46))

    # second ticks + labels
    if duration > 0:
        step = 5 if duration <= 60 else (10 if duration <= 180 else 30)
        for t in range(0, int(duration) + 1, step):
            x = pad_l + (t / duration) * (W - pad_l - pad_r)
            d.line([x, pad_t + bar_h + 3, x, pad_t + bar_h + 9], fill=(160, 160, 156))
            label = f"{t}s"
            tw = d.textlength(label)
            d.text((min(max(x - tw / 2, 2), W - tw - 2), pad_t + bar_h + 12), label, fill=(110, 110, 106))

    span_txt = ", ".join(f"[{s:.0f}–{e:.0f}s]" for s, e in spans) if spans else "no evidence returned"
    d.text((pad_l, 12), span_txt, fill=(60, 60, 58))
    return img


def format_result(answer: str, spans: List[Tuple[float, float]]) -> str:
    evidence = ", ".join(f"[{s:.0f}s – {e:.0f}s]" for s, e in spans)
    return f"{answer}\n\n\n⏱ Grounded evidence: {evidence if evidence else 'none returned'}"


# ─────────────────────────────────────────────────────────────────────────────
# Inference
# ─────────────────────────────────────────────────────────────────────────────

def _estimate_duration(video, question, num_frames=NUM_FRAMES, max_new_tokens=MAX_NEW_TOKENS,
                       *args, **kwargs) -> int:
    """Measured ~5s of GPU time for 48 frames + 1024 tokens on ZeroGPU; scale by
    both budgets and add margin for cold weight streaming."""
    try:
        n = int(num_frames)
    except (TypeError, ValueError):
        n = NUM_FRAMES
    try:
        m = int(max_new_tokens)
    except (TypeError, ValueError):
        m = MAX_NEW_TOKENS
    return min(240, 20 + int(n * 0.35) + int(m * 0.004))


@spaces.GPU(duration=_estimate_duration)
def answer_question(video, question, num_frames=NUM_FRAMES, max_new_tokens=MAX_NEW_TOKENS):
    """Answer an object-centric question about a video clip and ground the evidence.

    Args:
        video: an instructional video clip (file path).
        question: a question about an object in the clip.
        num_frames: number of frames sampled uniformly from the clip (48 = the
            model's training budget).
        max_new_tokens: maximum number of tokens to generate.

    Returns:
        (formatted answer + evidence, evidence timeline image, raw model JSON)
    """
    t0 = time.perf_counter()
    if not video:
        raise gr.Error("Please upload a video clip first.")
    question = (question or "").strip()
    if not question:
        raise gr.Error("Please enter a question about the clip.")
    try:
        num_frames = max(4, min(int(num_frames), 64))
    except (TypeError, ValueError):
        num_frames = NUM_FRAMES
    try:
        max_new_tokens = max(128, min(int(max_new_tokens), 2048))
    except (TypeError, ValueError):
        max_new_tokens = MAX_NEW_TOKENS

    frames, clip_duration = extract_frames(video, num_frames=num_frames)
    if not frames:
        raise gr.Error("Could not read any frames from this video. Try an .mp4 file.")

    # Message layout identical to the benchmark's VLLMRunner: one
    # "[Timestamp: <t>s]" text part before every frame image, then the question.
    content: List[Dict[str, Any]] = []
    for f in frames:
        content.append({"type": "text", "text": f"[Timestamp: {f['timestamp']}s]"})
        content.append({"type": "image", "image": f["image"]})
    content.append({"type": "text", "text": f"\nQuestion: {question}"})
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": content},
    ]

    chat_text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    images = [f["image"] for f in frames]
    inputs = processor(
        text=[chat_text], images=images if images else None, padding=True, return_tensors="pt"
    ).to("cuda")

    with torch.inference_mode():
        out_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            temperature=None,
            top_p=None,
            top_k=None,
        )
    gen_ids = out_ids[:, inputs["input_ids"].shape[1]:]
    raw = processor.batch_decode(gen_ids, skip_special_tokens=True)[0].strip()

    parsed = parse_model_response(raw)
    spans = _clean_intervals(parsed["evidence"], clip_duration)
    timeline = render_timeline(clip_duration, spans)
    elapsed = time.perf_counter() - t0

    payload = {"answer": parsed["answer"], "evidence": parsed["evidence"],
               "clip_duration_s": round(clip_duration, 2), "frames": len(frames),
               "inference_s": round(elapsed, 1)}
    return format_result(parsed["answer"], spans), timeline, json.dumps(payload, ensure_ascii=False, indent=2)


# ─────────────────────────────────────────────────────────────────────────────
# UI
# ─────────────────────────────────────────────────────────────────────────────

EXAMPLES = [
    ["examples/3e9350958273.mp4",
     "How does the state of the egg change from the start of the whisking action to the end of this sequence?"],
    ["examples/ba6ce3240abf.mp4",
     "What mistake occurs on the pan with cheese, and why is it a mistake?"],
    ["examples/b039e8c9fafd.mp4",
     "How does the state of the bowl contents change from being stirred initially to the end of this sequence?"],
    ["examples/37aabc657f3b.mp4",
     "How does the state of the coffee grounds change from the start of the blooming phase to the end of this sequence?"],
    ["examples/670e2b9f02bf.mp4",
     "Before cutting the tomato into two pieces begins, is the tomato ready, and what visible cues show this?"],
]

CSS = """
#col-container { max-width: 1150px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🍳 ProcObject-Qwen3-VL-4B — Object-Centric Video QA & Temporal Grounding

            Ask a question about an **object** in an instructional video clip. The model answers and localizes the
            supporting moments as `[start, end]` second intervals — the object-centric reasoning of
            [ProcObject-10K](https://huggingface.co/datasets/BrightGuo/ProcObject-10K)
            ([paper](https://arxiv.org/abs/2512.03479), NeurIPS 2026 Datasets & Benchmarks).
            """
        )

        with gr.Row():
            with gr.Column(scale=1):
                video = gr.Video(label="Video clip")
                question = gr.Textbox(
                    label="Question",
                    placeholder="e.g. How does the state of the egg change from the start of the whisking to the end?",
                    lines=3,
                )
                run = gr.Button("Answer", variant="primary")
                with gr.Accordion("Advanced settings", open=False):
                    num_frames = gr.Slider(minimum=8, maximum=64, value=NUM_FRAMES, step=2,
                                           label="Frames sampled", info="48 is the fine-tuned model's training budget")
                    max_new_tokens = gr.Slider(minimum=128, maximum=2048, value=MAX_NEW_TOKENS, step=64,
                                               label="Max new tokens")

            with gr.Column(scale=1):
                answer_out = gr.Textbox(label="Answer + grounded evidence", lines=12, buttons=["copy"])
                timeline_out = gr.Image(label="Evidence timeline", height=140, buttons=[])
                json_out = gr.Textbox(label="Model output (JSON)", lines=6, buttons=["copy"])

        gr.Markdown(
            """
            #### Tips
            - Ask about an **object** (the egg, the pan, the coffee grounds…), not only about the person: the model is
            trained for object-centric reasoning — preconditions, state evolution, counterfactuals, mistakes, readiness.
            - Evidence intervals are in seconds from the start of the clip, matching the `[Timestamp: …s]` frame labels.
            - The bundled example clips are from [CaptainCook4D](https://captaincook4d.github.io/captain-cook/), one of
            the four source datasets of ProcObject-10K, cut to the exact benchmark clip windows.
            """
        )

        gr.Examples(
            examples=EXAMPLES,
            inputs=[video, question],
            outputs=[answer_out, timeline_out, json_out],
            fn=answer_question,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        answer_question,
        inputs=[video, question, num_frames, max_new_tokens],
        outputs=[answer_out, timeline_out, json_out],
        api_name="answer",
    )
    question.submit(
        answer_question,
        inputs=[video, question, num_frames, max_new_tokens],
        outputs=[answer_out, timeline_out, json_out],
        api_name="answer_submit",
    )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)