"""ToolMerge — Decomposing Queries into Tool Calls for Long-Video Keyframe Retrieval.

Live Gradio demo of the SigLIP path of ToolMerge (paper 2605.23826).

Pipeline (faithful port of https://github.com/michalsr/ToolMerge):
  1. A text-only Qwen3-VL planner (michalsr/toolmerge-planner-grpo) decomposes
     the question + answer choices into independent SigLIP search queries combined
     with AND/OR boolean operators.
  2. Frames are sampled from the uploaded video (2 fps, matching the paper's cache
     phase) and encoded with SigLIP-2 (google/siglip2-giant-opt-patch16-384).
  3. Each planner query is TEXT-encoded and dot-producted against the per-frame
     image features -> per-frame percentile ranks.
  4. Ranks are merged with the planner's AND(min)/OR(max) expression.
  5. Greedy NMS with the paper's auto temporal gap tau = min(D/(2K), 10)s selects
     the final top-K keyframes.

Only the SigLIP tool is run live (the paper's T-REN region-text tool and the
OCR-judge stage require offline-built per-video caches / separate weights). The
merge naturally handles a single-tool plan, exactly as the reference code does
when a tool is disabled via `enabled_tools`.
"""

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

import json
import logging
import re
import tempfile
from typing import Dict, List, Tuple

import spaces
import torch
import torch.nn.functional as F
import numpy as np
import cv2
from PIL import Image
import gradio as gr
from transformers import (
    AutoModel,
    AutoProcessor,
    Qwen3VLForConditionalGeneration,
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("toolmerge")

PLANNER_ID = "michalsr/toolmerge-planner-grpo"
SIGLIP_ID = "google/siglip2-giant-opt-patch16-384"

TARGET_FPS = 2.0          # paper cache phase samples at 2 fps
MAX_FRAMES = 900          # safety cap on frames encoded (900 / 2fps = 7.5 min)
GAP_CAP_SECONDS = 10.0    # paper default cap on auto-tau

# ---------------------------------------------------------------------------
# Planner prompt (verbatim from toolmerge/prompts/planner/v7_no_temporal.py,
# with the T-REN tool removed since only SigLIP runs live).
# ---------------------------------------------------------------------------
PLANNER_TEMPLATE = """\
You are a search planner for a video question-answering system. \
Given a question and answer choices, write queries for specific search tools that LOCATE the relevant frames. \
A separate answerer model will look at those frames and determine the correct answer — \
you do NOT answer the question yourself.

## Tools

**siglip** — Visual similarity search.
  - Describe what the scene LOOKS LIKE: settings, actions, spatial layout, object attributes, visual states.
  - Cannot read text. Never include on-screen text in siglip queries.
  - Bad: "sign reading Exit Here" → Good: "hallway with illuminated signs"
  - Bad: "someone is happy" → Good: "person smiling and clapping"

For siglip, focus on the most visually distinctive feature — rare details beat generic descriptions.

## Query design

- Break complex scenes into separate queries.
- Keep siglip queries concrete and visual. Avoid abstract or narrative language.

## Combining queries (1-5 queries per plan)

- **AND** = intersection. Scene has multiple distinctive elements — one query each. \
Never AND queries that describe the same thing differently.
- **OR** = union. Different scenes, or different queries that might each find what you need.

## Rules

1. **Locate, don't answer.** Find the scene; the answerer decides what's happening.
2. **Always output at least one query.** Every question has a visual scene to find.
3. **Use all information.** Extract every visually searchable detail from the question AND \
the answer choices. Entities, objects, settings, actions — if it can help locate the right frames, query for it.
4. **Use answer choices wisely.** Visually different choices → search for each. \
Same scene described differently → one query, let the answerer decide.

## Question

{question}

Options:
{options}

Video duration: {duration}s encoded at {fps} fps.

You MUST first write 1-3 sentences of reasoning before the JSON block. Think about: \
what must be visually true about the frames that contain the answer? What is the most \
distinctive element to search for? Do the answer choices point to different scenes or \
the same scene? Never output the JSON block without reasoning first.

Then output a JSON block:
```json
{"queries": [{"tool": "siglip", "query": "...", "id": "Q1"}], "combine": "Q1"}
```

Fields per query: "tool", "query", "id" (Q1, Q2, ...).

Examples:

---

Question: What does the woman in the red dress do after picking up the book from the table?
Options: A) places it on the shelf B) hands it to the man in glasses C) sits down on the couch and reads D) puts it in her bag E) walks out of the room

The question mentions a woman in a red dress, a book, and a table. The choices describe different actions after picking up the book — each would look different visually. I'll find the woman with the book and search for the distinct scenes from each choice.
```json
{"queries": [{"tool": "siglip", "query": "woman in red dress holding a book", "id": "Q1"}, {"tool": "siglip", "query": "person placing book on shelf", "id": "Q2"}, {"tool": "siglip", "query": "person handing book to someone", "id": "Q3"}, {"tool": "siglip", "query": "person sitting on couch reading", "id": "Q4"}], "combine": "Q1 AND (Q2 OR Q3 OR Q4)"}
```

---

Question: In which room does the child first play with the wooden blocks?
Options: A) the kitchen B) the living room with the blue rug C) the bedroom D) the hallway E) the backyard

The question mentions a child and wooden blocks. The choices are different rooms, each visually distinct. I'll find the child with blocks and search for each room.
```json
{"queries": [{"tool": "siglip", "query": "child playing with wooden blocks", "id": "Q1"}, {"tool": "siglip", "query": "child playing in kitchen", "id": "Q2"}, {"tool": "siglip", "query": "living room with blue rug", "id": "Q3"}, {"tool": "siglip", "query": "child playing in bedroom", "id": "Q4"}], "combine": "Q1 AND (Q2 OR Q3 OR Q4)"}
```

"""


# ---------------------------------------------------------------------------
# Model loading (module scope, eager .to("cuda") per ZeroGPU rules).
# ---------------------------------------------------------------------------
logger.info("Loading planner %s ...", PLANNER_ID)
planner_processor = AutoProcessor.from_pretrained(PLANNER_ID)
planner_model = Qwen3VLForConditionalGeneration.from_pretrained(
    PLANNER_ID,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).eval().to("cuda")

logger.info("Loading SigLIP-2 %s ...", SIGLIP_ID)
siglip_processor = AutoProcessor.from_pretrained(SIGLIP_ID)
siglip_model = AutoModel.from_pretrained(
    SIGLIP_ID,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).eval().to("cuda")
logger.info("Models loaded.")


# ---------------------------------------------------------------------------
# Planner JSON parsing (ported from toolmerge/planner.py).
# ---------------------------------------------------------------------------
def parse_planner_response(response: str) -> Tuple[List[dict], str]:
    json_match = re.search(r"```(?:json)?\s*(.*?)\s*```", response, re.DOTALL)
    if json_match:
        json_str = json_match.group(1)
    else:
        brace_match = re.search(r"\{.*\}", response, re.DOTALL)
        if not brace_match:
            return [], ""
        json_str = brace_match.group(0)
    try:
        data = json.loads(json_str)
    except json.JSONDecodeError:
        return [], ""
    queries = data.get("queries", [])
    combine_expr = data.get("combine", "")
    valid: List[dict] = []
    for q in queries:
        if isinstance(q, dict) and "tool" in q and "query" in q and "id" in q:
            valid.append({"tool": q["tool"], "query": q["query"], "id": q["id"]})
    return valid, combine_expr


# ---------------------------------------------------------------------------
# Boolean merge AST (ported verbatim from toolmerge/merging.py).
# ---------------------------------------------------------------------------
class Node:
    pass


class Leaf(Node):
    __slots__ = ("query_id",)

    def __init__(self, query_id: str):
        self.query_id = query_id

    def __repr__(self):
        return self.query_id


class BinOp(Node):
    __slots__ = ("op", "left", "right")

    def __init__(self, op, left, right):
        self.op = op
        self.left = left
        self.right = right

    def __repr__(self):
        return f"({self.left} {self.op} {self.right})"


def tokenize(expr: str) -> List[str]:
    tokens: List[str] = []
    s = expr.strip()
    i = 0
    while i < len(s):
        c = s[i]
        if c in ("(", ")"):
            tokens.append(c)
            i += 1
        elif c.isspace():
            i += 1
        else:
            j = i
            while j < len(s) and not s[j].isspace() and s[j] not in ("(", ")"):
                j += 1
            tokens.append(s[i:j])
            i = j
    return tokens


def parse_atom(tokens, pos):
    if pos[0] >= len(tokens):
        raise ValueError("Unexpected end of combine expression")
    tok = tokens[pos[0]]
    if tok == "(":
        pos[0] += 1
        node = parse_or(tokens, pos)
        if pos[0] < len(tokens) and tokens[pos[0]] == ")":
            pos[0] += 1
        return node
    pos[0] += 1
    return Leaf(tok)


def parse_and(tokens, pos):
    left = parse_atom(tokens, pos)
    while pos[0] < len(tokens) and tokens[pos[0]].upper() == "AND":
        pos[0] += 1
        right = parse_atom(tokens, pos)
        left = BinOp("AND", left, right)
    return left


def parse_or(tokens, pos):
    left = parse_and(tokens, pos)
    while pos[0] < len(tokens) and tokens[pos[0]].upper() == "OR":
        pos[0] += 1
        right = parse_and(tokens, pos)
        left = BinOp("OR", left, right)
    return left


def parse_combine_expr(expr: str) -> Node:
    tokens = tokenize(expr)
    pos = [0]
    return parse_or(tokens, pos)


def evaluate_combine_scores(node, query_score_maps):
    if isinstance(node, Leaf):
        return dict(query_score_maps.get(node.query_id, {}))
    if isinstance(node, BinOp):
        left = evaluate_combine_scores(node.left, query_score_maps)
        right = evaluate_combine_scores(node.right, query_score_maps)
        all_indices = set(left) | set(right)
        out: Dict[int, float] = {}
        for idx in all_indices:
            ls = left.get(idx, 0.0)
            rs = right.get(idx, 0.0)
            if node.op == "AND":
                out[idx] = min(ls, rs)
            elif node.op == "OR":
                out[idx] = max(ls, rs)
        return out
    return {}


def combine_or_all(query_score_maps):
    out: Dict[int, float] = {}
    for sm in query_score_maps.values():
        for idx, s in sm.items():
            out[idx] = s if idx not in out else max(out[idx], s)
    return out


# ---------------------------------------------------------------------------
# Scoring (ported from toolmerge/tools/scoring.py).
# ---------------------------------------------------------------------------
def normalize_to_percentiles(results):
    if not results:
        return results
    n = len(results)
    if n == 1:
        return [(results[0][0], 1.0)]
    by_score = sorted(results, key=lambda x: x[1])
    percentiles = {idx: rank / (n - 1) for rank, (idx, _) in enumerate(by_score)}
    out = [(idx, percentiles[idx]) for idx in percentiles]
    out.sort(key=lambda x: x[1], reverse=True)
    return out


# ---------------------------------------------------------------------------
# Selection: greedy NMS (ported from toolmerge/selection.py).
# ---------------------------------------------------------------------------
def auto_tau_seconds(num_frames, fps, max_k, cap=GAP_CAP_SECONDS):
    if fps <= 0 or max_k <= 0:
        return 0.0
    duration = num_frames / fps
    return min(duration / (2 * max_k), cap)


def greedy_gap_select(scored, max_k, min_gap_frames):
    ranked = sorted(scored.items(), key=lambda x: x[1], reverse=True)
    selected: Dict[int, float] = {}
    for idx, score in ranked:
        if len(selected) >= max_k:
            break
        if min_gap_frames <= 0 or all(abs(idx - s) >= min_gap_frames for s in selected):
            selected[idx] = score
    return selected


# ---------------------------------------------------------------------------
# Video decoding.
# ---------------------------------------------------------------------------
def decode_frames(video_path, target_fps=TARGET_FPS, max_frames=MAX_FRAMES):
    """Sample frames at target_fps. Returns (list[PIL.Image], list[float timestamps], src_fps)."""
    cap = cv2.VideoCapture(video_path)
    src_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
    if src_fps <= 0:
        src_fps = 30.0
    step = max(1, int(round(src_fps / target_fps)))
    frames: List[Image.Image] = []
    timestamps: List[float] = []
    idx = 0
    grabbed = 0
    while True:
        ret = cap.grab()
        if not ret:
            break
        if idx % step == 0:
            ok, frame = cap.retrieve()
            if not ok:
                break
            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(Image.fromarray(rgb))
            timestamps.append(idx / src_fps)
            grabbed += 1
            if grabbed >= max_frames:
                break
        idx += 1
    cap.release()
    return frames, timestamps, src_fps


# ---------------------------------------------------------------------------
# GPU inference.
# ---------------------------------------------------------------------------
def _as_tensor(out):
    """Coerce a get_*_features() return value to a plain feature tensor.

    Depending on the transformers version, get_image_features / get_text_features
    may return a tensor directly or a ModelOutput wrapping it.
    """
    if isinstance(out, torch.Tensor):
        return out
    for attr in ("image_embeds", "text_embeds", "pooler_output", "last_hidden_state"):
        val = getattr(out, attr, None)
        if isinstance(val, torch.Tensor):
            return val
    # Fall back to first tensor element (ModelOutput is tuple-like).
    try:
        first = out[0]
        if isinstance(first, torch.Tensor):
            return first
    except Exception:
        pass
    raise TypeError(f"Cannot extract feature tensor from {type(out)}")


def _encode_siglip_images(frames, batch_size=64):
    feats = []
    for i in range(0, len(frames), batch_size):
        batch = frames[i:i + batch_size]
        inputs = siglip_processor(images=batch, return_tensors="pt")
        pixel_values = inputs["pixel_values"].to("cuda", dtype=torch.bfloat16)
        with torch.no_grad():
            f = _as_tensor(siglip_model.get_image_features(pixel_values=pixel_values))
        f = f.float()
        f = f / f.norm(dim=-1, keepdim=True)
        feats.append(f.cpu())
    return torch.cat(feats, dim=0)  # (T, D)


def _encode_siglip_text(query):
    inputs = siglip_processor(
        text=[query],
        return_tensors="pt",
        padding="max_length",
        truncation=True,
        max_length=64,
    )
    input_ids = inputs["input_ids"].to("cuda")
    with torch.no_grad():
        f = _as_tensor(siglip_model.get_text_features(input_ids=input_ids))
    f = f.float()
    f = f / f.norm(dim=-1, keepdim=True)
    return f.cpu().squeeze(0)  # (D,)


def _score_siglip(query, image_feats):
    text_feat = _encode_siglip_text(query)
    text_feat = F.normalize(text_feat.unsqueeze(0), p=2, dim=1).squeeze(0)
    emb = F.normalize(image_feats.float(), p=2, dim=1)
    scores = torch.matmul(emb, text_feat)  # (T,)
    results = [(i, float(s)) for i, s in enumerate(scores.tolist())]
    results.sort(key=lambda x: x[1], reverse=True)
    return normalize_to_percentiles(results)


def _run_planner(question, options_text, duration_s):
    prompt_text = (
        PLANNER_TEMPLATE.replace("{question}", question)
        .replace("{options}", options_text)
        .replace("{duration}", str(int(round(duration_s))))
        .replace("{fps}", str(int(TARGET_FPS)))
    )
    messages = [{"role": "user", "content": [{"type": "text", "text": prompt_text}]}]
    text = planner_processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = planner_processor(text=[text], padding=True, return_tensors="pt").to("cuda")
    with torch.no_grad():
        output_ids = planner_model.generate(
            **inputs, max_new_tokens=512, do_sample=False
        )
    trimmed = output_ids[0][inputs.input_ids.shape[1]:]
    return planner_processor.decode(trimmed, skip_special_tokens=True)


def _estimate_duration(video_path, question, options, k, *args, **kwargs):
    return 150


@spaces.GPU(duration=_estimate_duration)
def retrieve_keyframes(video_path, question, options, k=8,
                       progress=gr.Progress(track_tqdm=True)):
    """Retrieve the top-K keyframes for a question using the ToolMerge planner + SigLIP.

    Args:
        video_path: path to the input video file.
        question: the natural-language question about the video.
        options: newline-separated multiple-choice answer options (optional).
        k: number of keyframes to return.

    Returns:
        A tuple of (gallery of keyframe images with timestamp captions,
        markdown report of the planner's decomposition and merge).
    """
    if not video_path:
        raise gr.Error("Please provide a video.")
    if not question or not question.strip():
        raise gr.Error("Please enter a question.")

    k = int(k)
    options_text = (options or "").strip()

    # 1. Decode frames at 2 fps.
    progress(0.05, desc="Decoding video frames (2 fps)...")
    frames, timestamps, src_fps = decode_frames(video_path)
    if len(frames) == 0:
        raise gr.Error("Could not decode any frames from the video.")
    num_frames = len(frames)
    duration_s = timestamps[-1] if timestamps else num_frames / TARGET_FPS
    logger.info("Decoded %d frames @2fps (src %.1f fps), duration ~%.1fs",
                num_frames, src_fps, duration_s)

    # 2. Planner: decompose the query into tool calls.
    progress(0.15, desc="Running planner (Qwen3-VL)...")
    raw_response = _run_planner(question, options_text, duration_s)
    queries, combine_expr = parse_planner_response(raw_response)
    # Keep only SigLIP queries (only tool available live).
    siglip_queries = [q for q in queries if q["tool"] == "siglip"]
    if not siglip_queries:
        # Fallback: use the raw question as one SigLIP query.
        siglip_queries = [{"tool": "siglip", "query": question.strip(), "id": "Q1"}]
        combine_expr = "Q1"

    # 3. Encode frames once with SigLIP.
    progress(0.35, desc=f"Encoding {num_frames} frames with SigLIP-2...")
    image_feats = _encode_siglip_images(frames)

    # 4. Score each query -> percentile maps.
    progress(0.75, desc="Scoring queries...")
    query_score_maps: Dict[str, Dict[int, float]] = {}
    for q in siglip_queries:
        pct = _score_siglip(q["query"], image_feats)
        query_score_maps[q["id"]] = {idx: s for idx, s in pct}

    # 5. Merge via the planner's boolean expression.
    valid_ids = {q["id"] for q in siglip_queries}
    combine_ids = set(re.findall(r"Q\d+", combine_expr or ""))
    use_combine = bool(combine_expr) and combine_ids.issubset(valid_ids) and combine_ids
    if use_combine:
        try:
            ast = parse_combine_expr(combine_expr)
            combined = evaluate_combine_scores(ast, query_score_maps)
        except Exception as e:
            logger.warning("combine parse failed (%r); OR-fallback", e)
            combined = combine_or_all(query_score_maps)
            use_combine = False
    else:
        combined = combine_or_all(query_score_maps)

    # 6. Greedy NMS with auto temporal gap.
    progress(0.9, desc="Selecting keyframes (greedy NMS)...")
    tau_s = auto_tau_seconds(num_frames, TARGET_FPS, k)
    gap_frames = int(tau_s * TARGET_FPS)
    selected = greedy_gap_select(combined, k, gap_frames)
    ordered = sorted(selected.keys())

    # 7. Build gallery + report.
    gallery = []
    for idx in ordered:
        ts = timestamps[idx]
        cap = f"t={ts:5.1f}s  (score {selected[idx]:.2f})"
        gallery.append((frames[idx], cap))

    report_lines = ["## Query decomposition\n"]
    report_lines.append(f"**Question:** {question.strip()}\n")
    report_lines.append(
        f"**Video:** {num_frames} frames sampled @ {int(TARGET_FPS)} fps "
        f"(~{duration_s:.0f}s)\n"
    )
    report_lines.append("### Tool calls (planner output)\n")
    for q in siglip_queries:
        report_lines.append(f"- `{q['id']}` — **{q['tool']}**: \"{q['query']}\"")
    report_lines.append("")
    report_lines.append(
        f"### Combine expression\n`{combine_expr if use_combine else ' OR '.join(sorted(valid_ids))}`\n"
    )
    report_lines.append(
        f"AND = `min` (intersection), OR = `max` (union) over per-frame percentile ranks.\n"
    )
    report_lines.append(
        f"### Selection\nGreedy NMS, K={k}, auto temporal gap "
        f"τ = min(D/(2K), 10) = **{tau_s:.1f}s**. "
        f"Selected {len(ordered)} keyframes.\n"
    )
    report_lines.append("<details><summary>Raw planner response</summary>\n\n```\n"
                        + raw_response.strip()[:2000] + "\n```\n</details>")

    return gallery, "\n".join(report_lines)


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

INTRO = """\
# 🎬 ToolMerge — Long-Video Keyframe Retrieval

Decomposes a question into **SigLIP visual-search tool calls** combined with **AND/OR**
boolean logic, scores every frame, merges the ranks, and picks the top-K keyframes with
temporal non-maximum suppression.

Faithful port of the SigLIP path of **[ToolMerge](https://github.com/michalsr/ToolMerge)**
([paper](https://huggingface.co/papers/2605.23826)) using the released text-only planner
[`michalsr/toolmerge-planner-grpo`](https://huggingface.co/michalsr/toolmerge-planner-grpo)
(Qwen3-VL-8B, GRPO-tuned) + [`google/siglip2-giant-opt-patch16-384`](https://huggingface.co/google/siglip2-giant-opt-patch16-384).

*The paper's T-REN region-text tool and OCR-judge stage need offline-built per-video caches and are omitted here; the boolean merge runs SigLIP-only, exactly as the reference code does when other tools are disabled.*
"""


with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(INTRO)
        with gr.Row():
            with gr.Column(scale=1):
                video_in = gr.Video(label="Video", height=280)
                question_in = gr.Textbox(
                    label="Question",
                    placeholder="e.g. What is the chef doing over the flames?",
                    lines=2,
                )
                options_in = gr.Textbox(
                    label="Answer choices (optional, one per line)",
                    placeholder="A) ...\nB) ...\nC) ...",
                    lines=3,
                )
                with gr.Accordion("Advanced settings", open=False):
                    k_in = gr.Slider(
                        label="K (number of keyframes)",
                        minimum=1, maximum=32, step=1, value=8,
                    )
                run_btn = gr.Button("Retrieve keyframes", variant="primary")
            with gr.Column(scale=1):
                gallery_out = gr.Gallery(
                    label="Retrieved keyframes (in temporal order)",
                    columns=4, height=340, object_fit="contain",
                )
                report_out = gr.Markdown()

        gr.Examples(
            examples=[
                ["examples/chef_wok_flames.mp4",
                 "What is the cook doing with the food in the pan?",
                 "A) stirring vegetables\nB) flipping food over flames\nC) pouring sauce\nD) plating a dish",
                 8],
                ["examples/city_traffic_night.mp4",
                 "What is visible on the busy street at night?",
                 "A) empty road\nB) cars with bright light trails\nC) pedestrians only\nD) a parking lot",
                 8],
                ["examples/couple_dancing.mp4",
                 "What are the two people doing together?",
                 "A) dancing\nB) arguing\nC) eating dinner\nD) walking a dog",
                 8],
            ],
            inputs=[video_in, question_in, options_in, k_in],
            outputs=[gallery_out, report_out],
            fn=retrieve_keyframes,
            cache_examples=True,
            cache_mode="lazy",
        )

        run_btn.click(
            retrieve_keyframes,
            inputs=[video_in, question_in, options_in, k_in],
            outputs=[gallery_out, report_out],
            api_name="retrieve",
        )

if __name__ == "__main__":
    demo.launch(mcp_server=True)