"""TimeLens2-2B video temporal grounding demo."""
import os

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

import spaces  # noqa: E402  MUST come before any CUDA-touching import
import torch
from pathlib import Path
import gradio as gr
from transformers import AutoModelForImageTextToText, AutoProcessor
from qwen_vl_utils import process_vision_info

MODEL_ID = "MCG-NJU/TimeLens2-2B"

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda").eval()


@spaces.GPU(duration=30)
def ground(
    video_path: str,
    query: str,
    max_new_tokens: int = 512,
) -> str:
    """Given a video and a natural-language query, locate the time spans where the query is relevant.

    Args:
        video_path: Path or URL of the input video file.
        query: Natural-language description of the event to locate.
        max_new_tokens: Maximum tokens the model may generate.

    Returns:
        JSON array of [start, end] time pairs in seconds, e.g. [[1.5, 3.2], [10.0, 12.5]].
    """
    if not video_path or not query.strip():
        return "Please provide both a video and a query."

    prompt = (
        f'Given the query: "{query}", return ALL time spans (in seconds) '
        "where the query is relevant.\n"
        "Output format MUST be a JSON array of [start, end] pairs.\n"
    )

    if not video_path.startswith(("http://", "https://", "file://")):
        video_uri = Path(video_path).resolve().as_uri()
    else:
        video_uri = video_path

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "video",
                    "video": video_uri,
                    "fps": 2.0,
                    "min_pixels": 32 * 32,
                    "max_pixels": 480 * 480,
                    "total_pixels": 128000 * 32 * 32,
                },
                {"type": "text", "text": prompt},
            ],
        }
    ]

    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    images, videos, video_kwargs = process_vision_info(
        messages,
        image_patch_size=16,
        return_video_kwargs=True,
        return_video_metadata=True,
    )

    if videos is not None:
        videos, video_metadatas = zip(*videos)
        videos, video_metadatas = list(videos), list(video_metadatas)
    else:
        video_metadatas = None

    inputs = processor(
        text=text,
        images=images,
        videos=videos,
        video_metadata=video_metadatas,
        do_resize=False,
        return_tensors="pt",
        **video_kwargs,
    ).to(model.device)

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.01,
            top_p=0.001,
            top_k=1,
            repetition_penalty=1.0,
        )

    trimmed = [out[len(inp):] for inp, out in zip(inputs.input_ids, output_ids)]
    response = processor.batch_decode(
        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    return response[0].strip()


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

with gr.Blocks() as demo:
    gr.Markdown(
        "# 🎬 TimeLens2-2B — Video Temporal Grounding\n"
        "Upload a video and describe an event. The model returns the time spans "
        "(in seconds) where that event occurs.\n\n"
        "Model: [MCG-NJU/TimeLens2-2B](https://huggingface.co/MCG-NJU/TimeLens2-2B) — "
        "a compact 2B Qwen3-VL fine-tune that achieves SOTA at the 2B scale on seven "
        "temporal-grounding benchmarks (average mIoU 44.5), even surpassing TimeLens-8B."
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            video_in = gr.Video(label="Input video", sources=["upload"])
            query_in = gr.Textbox(
                label="Query",
                placeholder="e.g. A person opens the refrigerator",
                lines=2,
                scale=2,
            )

        run_btn = gr.Button("Ground", variant="primary")

        output = gr.Textbox(
            label="Predicted time spans (JSON)",
            lines=4,
            interactive=False,
        )

        with gr.Accordion("Advanced settings", open=False):
            max_tokens = gr.Slider(
                label="Max new tokens",
                minimum=64,
                maximum=4096,
                value=512,
                step=64,
            )

    gr.Examples(
        examples=[
            ["examples/barista_frothing.mp4", "milk is being frothed"],
            ["examples/cat_laser.mp4", "a cat chases a laser pointer"],
            ["examples/slicing_veggie.mp4", "vegetables are being sliced"],
        ],
        inputs=[video_in, query_in],
        outputs=output,
        fn=ground,
        cache_examples=True,
        cache_mode="lazy",
    )

    run_btn.click(
        fn=ground,
        inputs=[video_in, query_in, max_tokens],
        outputs=output,
        api_name="ground",
    )

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