"""
Dummy algorithm for the ORena SAVE FOCUS challenge — SEGMENT track.

Replace the marked blocks with your model. Every spot to change carries a ``TODO``.

What one run receives
---------------------
A single container run answers a **batch of questions**.
Inputs are mounted read-only at ``/input``:

  request.json            a LIST of focus.Request objects — one per question
  plain/<qID>.mp4         the video clip for question <qID>
  overlayed/<qID>.mp4     the same clip with an HH:MM:SS timestamp burned in
  FO_definitions.json     the foreign-object class definitions, as plain text
  batch.json              a convenience index of the batch (qIDs and layout)

and one file is written to ``/output``:

  answer.json             a LIST of focus.Response objects — one per question

``request.json`` is authoritative: iterate it, and answer every ``qID`` it lists.

Each question carries its own clip
----------------------------------
Every question has a **separate video file, already trimmed to that question's
window**, and no two questions in a batch come from the same source video. Your
answer to one question can therefore never depend on another question's footage.

Two consequences for your pipeline:

* ``request.start_time`` / ``end_time`` refer to the **original, untrimmed
  procedure video** — the same timeline shown in the overlay variant. The clip you
  are given has already been cut to that window, so decode the clip from its own
  beginning; do **not** seek to ``start_time`` inside it. ``request.duration``
  still tells you how long the window is.
* There is nothing to share between questions — no common video to cache or encode
  once. Each clip is decoded on its own.

Both variants show the same footage. Use the plain clip, the overlay, or both.

Clip encoding
-------------
Clips are H.264 MP4 at exactly **5 fps** (one frame every 0.2 s), height-normalised
to at most 576 px. Widths follow the source's aspect ratio (e.g., 1024x576, 720x576 or
640x360) — do not assume a fixed resolution. A
keyframe every 5 s keeps seeking cheap. The overlay clock shows the source
timeline's second, floored.

Load your model once
--------------------
Model loading, imports and warm-up happen **once per container run**, not once per
question. The latency budget is built around that, so keep model loading outside
the per-question loop — as this file does. This is the single most important thing
to get right for speed.

The latency budget
------------------
The platform measures the execution time of the whole run itself; the ``latency``
field you write into each Response is informational and is not used for scoring.

The budget is pooled across the batch. For a batch of B questions:

    allowed = 120 s  +  B x 15 s
              ^setup    ^per-question processing (the SEGMENT budget)

Going over does not fail the run outright. Instead, a share of your questions is
**forfeited** (scored incorrect), growing with the overrun: a small overrun costs a
few questions, and exceeding the budget by 20% forfeits the whole batch.

The forfeited questions are selected for you — not necessarily the slow ones — so an
overrun costs you questions you would have answered correctly. If you are running out
of budget, answer the rest quickly rather than overrun.

Keep going when a question fails
--------------------------------
One bad question should not cost you the batch. Catch errors per question, still
emit a Response for that ``qID``, and carry on — as the loop below does. A question
with no answer is simply scored incorrect.

Frame sampling
--------------
The clip is decoded with `decord <https://pypi.org/project/decord/>`_. Full-frame
solid-blue signal-loss screens are detected from pixels, without relying on a
video identifier. Exactly 80 bin-centre frames are then sampled over the remaining
time and letterboxed to the Tr006 training resolution. Their timestamps are mapped
back to the original timeline from ``Request.start_time``.

Trying it out
-------------
A sample batch (``test/input/interface_1/``) ships with the template. This package
stores the merged Tr006 checkpoint under ``resources/model``; runtime inference is
fully local and does not download model files.

  ./do_test_run.sh   # build the image and run it on the sample batch
  ./do_save.sh       # build and save a .tar.gz for upload
"""

import json
import logging
import re
import sys
import time
from pathlib import Path

import numpy as np
import torch

# decord must be imported AFTER torch: importing it first breaks CUDA
# initialisation (a known decord issue — torch.cuda init then fails with
# "random_device could not be read" or a segfault).
import decord
from focus import Request, Response, load_requests, save_items
from PIL import Image, ImageOps
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration

logging.basicConfig(
    stream=sys.stdout,
    level=logging.DEBUG,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger(__name__)

RESOURCES_PATH = Path(__file__).parent / "resources"
MODEL_PATH = RESOURCES_PATH / "model"

INPUT_PATH = Path("/input")
OUTPUT_PATH = Path("/output")
N_FRAMES = 80
FRAME_WIDTH = 420
FRAME_HEIGHT = 224
BLUE_DETECT_CHUNK = 64
BLUE_DETECT_STRIDE = 16
BLUE_REFINE_ATTEMPTS = 2
MAX_NEW_TOKENS = 64
USE_OVERLAY = False  # The trained Tr006 V2 model used plain frames + metadata.
VIDEO_DIR = INPUT_PATH / ("overlayed" if USE_OVERLAY else "plain")

# Keep this prompt byte-for-byte aligned with the Tr006 V2 fixed80 training
# manifest.  The official package supplies FO_definitions separately; it is
# deliberately not concatenated to the model prompt because the training rows
# used this system message and the question only.
SYSTEM_PROMPT = (
    "You are a surgical video analysis assistant. You are shown a short clip from a laparoscopic procedure and asked a single question about it.\n\n"
    "The clip is given as a sequence of frames in chronological order. Each frame is preceded by its timestamp in the source video, written as "
    "<SECONDS seconds> -- a frame marked <1234.5 seconds> was taken 1234.5 seconds (00:20:34) after the video began. Consecutive frames may be several seconds "
    "apart, so an event can happen between two of them.\n\n"
    "A foreign object (FO) is any object fully introduced into the patient's body\n"
    "cavity during surgery that must be retrieved or accounted for. Importantly,\n"
    "standard surgical instruments that remain connected to the external environment\n"
    "(e.g., graspers, scissors, trocars, staplers, cameras) are not considered foreign\n"
    "objects. Furthermore, we exclude detachable parts of surgical instruments,\n"
    "particularly anvil components of staplers.\n\n"
    "The foreign object classes are exactly: Sponge, Clip, Specimen Bag, Silicone Loop, External Drain, Needle, Gallstone, Specimen, Mesh, Absorbable Hemostatic Agent.\n\n"
    "Reply with the answer and nothing else -- no reasoning, no preamble, no\n"
    "explanation, no restating the question. A single short line.\n\n"
    "Rules for the answer:\n"
    "- Write the value only. No sentence, no explanation, no units, no trailing\n"
    "  period, and never repeat the question.\n"
    "- Asks yes or no -> write exactly: yes   or   no\n"
    "- Asks how many / for a count -> write digits only, e.g. 0 or 1 or 2.\n"
    "- Asks which foreign object class(es) -> write class names exactly as spelled\n"
    "  in the list above, comma-separated (e.g. Clip, Sponge), or exactly: none\n"
    "  Never answer with a generic description such as \"surgical instrument\".\n"
    "- Asks when something happens -> write hh:mm:ss, read off the frame timestamps.\n"
    "  Those timestamps count from the START OF THE VIDEO, not from the start of this\n"
    "  clip, so <1234.5 seconds> is 00:20:34. Convert carefully: 3600 seconds is one\n"
    "  hour, 60 seconds is one minute. If the event falls between two frames, give\n"
    "  your best estimate between their timestamps rather than snapping to one.\n"
    "  For several time points, separate them with commas: 00:09:19, 00:12:44\n"
    "- Asks for a percentage or a share -> write the number only, e.g. 40\n"
    "- Lists options to choose from -> copy exactly one of those options, verbatim.\n"
    "- Anything else -> a short phrase, at most a few words.\n\n"
    "If you are unsure, still commit to your single best answer in the required\n"
    "form. An empty, hedged, or explanatory answer is scored as wrong.\n"
)


def _flatten_token_ids(value) -> list[int]:
    """Flatten tokenizer/config token-id values while dropping invalid entries."""
    if value is None or isinstance(value, bool):
        return []
    if isinstance(value, (int, np.integer)):
        return [int(value)] if int(value) >= 0 else []
    if hasattr(value, "detach"):
        value = value.detach().cpu().tolist()
    elif hasattr(value, "tolist") and not isinstance(value, (str, bytes)):
        value = value.tolist()
    if isinstance(value, (list, tuple, set)):
        out = []
        for item in value:
            out.extend(_flatten_token_ids(item))
        return out
    return []


def _ordered_unique(values: list[int]) -> list[int]:
    return list(dict.fromkeys(int(value) for value in values if int(value) >= 0))


def resolve_eos_token_ids(model, tokenizer) -> list[int]:
    """Merge model and tokenizer EOS values, preserving order and removing duplicates."""
    values = []
    for owner in (
        getattr(model, "generation_config", None),
        getattr(model, "config", None),
        tokenizer,
    ):
        values.extend(_flatten_token_ids(getattr(owner, "eos_token_id", None)))
    return _ordered_unique(values)


def resolve_im_start_token_ids(tokenizer) -> list[int]:
    """Return the tokenizer's (possibly multi-token) ``<|im_start|>`` sequence."""
    try:
        values = _flatten_token_ids(tokenizer.encode("<|im_start|>", add_special_tokens=False))
    except Exception:
        values = []
    if values:
        return values
    try:
        value = tokenizer.convert_tokens_to_ids("<|im_start|>")
    except Exception:
        value = None
    return _flatten_token_ids(value)


def _token_rows(value) -> list[list[int]]:
    if hasattr(value, "detach"):
        value = value.detach().cpu().tolist()
    elif hasattr(value, "tolist") and not isinstance(value, (str, bytes)):
        value = value.tolist()
    if not isinstance(value, list):
        value = [value]
    if value and not isinstance(value[0], list):
        value = [value]
    return [[int(token) for token in row] for row in value]


def extract_new_token_ids(generated, input_ids, attention_mask=None) -> list[list[int]]:
    """Extract only generated IDs, accounting for padded prompt layouts."""
    outputs = _token_rows(generated)
    prompts = _token_rows(input_ids)
    masks = _token_rows(attention_mask) if attention_mask is not None else []
    if not prompts:
        return outputs
    extracted = []
    for index, output in enumerate(outputs):
        prompt = prompts[min(index, len(prompts) - 1)]
        mask = masks[min(index, len(masks) - 1)] if masks else [1] * len(prompt)
        active_prompt = [token for token, flag in zip(prompt, mask) if flag]
        if len(output) >= len(prompt) and output[: len(prompt)] == prompt:
            offset = len(prompt)
        elif active_prompt and len(output) >= len(active_prompt) and output[: len(active_prompt)] == active_prompt:
            offset = len(active_prompt)
        else:
            offset = 0
        extracted.append(output[offset:])
    return extracted


def _first_subsequence(values: list[int], needle: list[int]) -> int | None:
    if not needle or len(needle) > len(values):
        return None
    width = len(needle)
    for index in range(len(values) - width + 1):
        if values[index : index + width] == needle:
            return index
    return None


def trim_generated_token_ids(
    raw_token_ids,
    *,
    eos_token_ids,
    im_start_token_ids,
    max_new_tokens: int,
) -> dict:
    """Stop at the earliest EOS or im_start sequence and retain a raw-token audit."""
    raw = _flatten_token_ids(raw_token_ids)
    eos_ids = set(_flatten_token_ids(eos_token_ids))
    im_start = _flatten_token_ids(im_start_token_ids)
    eos_positions = [index for index, token in enumerate(raw) if token in eos_ids]
    eos_index = eos_positions[0] if eos_positions else None
    im_start_index = _first_subsequence(raw, im_start)
    candidates = [(index, "eos") for index in (eos_index,) if index is not None]
    candidates.extend((index, "im_start") for index in (im_start_index,) if index is not None)
    stop_index, stop_reason = min(candidates, default=(None, None), key=lambda item: (item[0], 0 if item[1] == "eos" else 1))
    hit_max = stop_index is None and len(raw) >= int(max_new_tokens)
    if stop_index is None:
        stop_reason = "max_new_tokens" if hit_max else "completed"
        token_ids = raw
    else:
        token_ids = raw[:stop_index]
    return {
        "raw_token_ids": raw,
        "token_ids": token_ids,
        "stop_index": stop_index,
        "eos_position": eos_index,
        "eos_positions": eos_positions,
        "stop_reason": stop_reason,
        "hit_max_new_tokens": hit_max,
    }


_LEADING_THINK = re.compile(r"\A\s*<think>.*?</think>\s*", re.IGNORECASE | re.DOTALL)
_TRAILING_SPECIAL = re.compile(r"(?:<\|[^>\n]+\|>|</think>|<think>)", re.IGNORECASE)
_ROLE_MARKER = re.compile(r"\A\s*(?:assistant|user|system)\s*(?::|\Z)", re.IGNORECASE)


def clean_decoded_answer(text: str) -> str:
    """Remove only a complete leading think block and trailing template pollution."""
    value = str(text or "").strip()
    leading = _LEADING_THINK.match(value)
    if leading:
        value = value[leading.end() :].strip()
    elif value.lower().startswith("<think>"):
        return value
    value = value.splitlines()[0].strip() if value else ""
    marker = _TRAILING_SPECIAL.search(value)
    if marker:
        value = value[: marker.start()].rstrip()
    role = _ROLE_MARKER.match(value)
    if role:
        value = ""
    return value


def log_environment(device: torch.device) -> None:
    """Report the CUDA stack this container ended up with — once, at startup.

    A torch build carrying no kernels for the GPU it was handed is the classic
    reason a run dies on the platform, and it stays invisible until the first
    CUDA call. Printing this makes a failed job diagnosable from the logs alone.
    Keep it: it costs milliseconds, once per run.
    """
    log.info("--- Environment ---")
    log.info("  torch          : %s (built against CUDA %s)", torch.__version__, torch.version.cuda)

    # get_arch_list() is empty when no GPU is visible; the architecture flags
    # compiled into the wheel can be read either way.
    arch = torch.cuda.get_arch_list() or (torch._C._cuda_getArchFlags() or "").split()
    log.info("  torch kernels  : %s", " ".join(arch) or "unknown")

    driver_file = Path("/proc/driver/nvidia/version")
    log.info(
        "  host driver    : %s",
        driver_file.read_text().strip().splitlines()[0]
        if driver_file.exists()
        else "no NVIDIA driver visible",
    )

    if device.type != "cuda":
        log.warning("  No GPU visible — running on CPU. The platform always provides one.")
        return

    capability = torch.cuda.get_device_capability(0)
    free, total = torch.cuda.mem_get_info(0)
    log.info("  GPU            : %s", torch.cuda.get_device_name(0))
    log.info("  capability     : %s (sm_%d%d)", capability, *capability)
    log.info("  VRAM           : %.1f GiB free of %.1f GiB", free / 1024**3, total / 1024**3)


def clip_path_for(req: Request) -> Path:
    """The video clip belonging to this question.

    Every question has exactly one clip, named after its ``qID`` and already cut to
    that question's ``[start_time, end_time]`` window.
    """
    return VIDEO_DIR / f"{req.qID}.mp4"


def blue_screen_flags(frames: np.ndarray) -> np.ndarray:
    """Classify full-frame solid-blue signal loss on a sparse pixel grid."""
    sampled = frames[:, ::BLUE_DETECT_STRIDE, ::BLUE_DETECT_STRIDE]
    pixels = sampled.reshape(len(sampled), -1, 3).astype(np.float32)
    red, green, blue = pixels[:, :, 0], pixels[:, :, 1], pixels[:, :, 2]
    blue_fraction = ((blue > 180) & (blue > 4 * (red + 1)) & (blue > 4 * (green + 1))).mean(1)
    return (blue_fraction >= 0.98) & (pixels.std(1).mean(1) <= 15)


def sample_non_blue_indices(blue: np.ndarray) -> np.ndarray:
    valid = np.flatnonzero(~blue)
    if not len(valid):
        log.warning("Every decoded frame looks like a blue screen; using the full clip")
        valid = np.arange(len(blue))
    positions = (np.arange(N_FRAMES) + 0.5) * len(valid) / N_FRAMES
    slots = np.floor(positions + 0.5 - 1e-9).astype(np.int64).clip(0, len(valid) - 1)
    return valid[slots]


def detect_blue_intervals(vr: decord.VideoReader) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Probe 80 time bins and expand blue probes to their represented intervals."""
    probes = sample_non_blue_indices(np.zeros(len(vr), dtype=bool))
    probe_frames = vr.get_batch(probes).asnumpy()
    flags = blue_screen_flags(probe_frames)
    if flags.all():
        blue = np.zeros(len(vr), dtype=bool)
        for begin in range(0, len(vr), BLUE_DETECT_CHUNK):
            indices = np.arange(begin, min(begin + BLUE_DETECT_CHUNK, len(vr)))
            blue[indices] = blue_screen_flags(vr.get_batch(indices).asnumpy())
        return blue, probes, probe_frames
    boundaries = np.concatenate((
        [0], ((probes[:-1] + probes[1:]) // 2 + 1), [len(vr)]
    ))
    blue = np.zeros(len(vr), dtype=bool)
    for index in np.flatnonzero(flags):
        blue[boundaries[index]:boundaries[index + 1]] = True
    return blue, probes, probe_frames


def letterbox_frames(frames: np.ndarray) -> np.ndarray:
    return np.stack([
        np.asarray(ImageOps.pad(
            Image.fromarray(frame),
            (FRAME_WIDTH, FRAME_HEIGHT),
            method=Image.Resampling.BICUBIC,
            color=(0, 0, 0),
            centering=(0.5, 0.5),
        ))
        for frame in frames
    ])


def your_preprocessing(video_path: Path, req: Request) -> tuple[np.ndarray, dict]:
    """Decode exactly ``N_FRAMES`` and attach original-timeline millisecond PTS.

    Official clips are already trimmed and encoded at 5fps. Blue signal-loss time
    is detected online, then the same non-blue bin-centre and 420x224 letterbox
    contracts used for Tr006 V2 are applied. ``frames_indices`` remain absolute.
    """
    log.debug("Decoding %s (size on disk: %d bytes)", video_path, video_path.stat().st_size)

    vr = decord.VideoReader(str(video_path), ctx=decord.cpu(0), num_threads=1)
    base_fps = float(vr.get_avg_fps())
    n_frames = len(vr)
    if n_frames <= 0 or base_fps <= 0:
        raise RuntimeError(f"invalid video stream: n_frames={n_frames}, fps={base_fps}")
    blue, indices, frames = detect_blue_intervals(vr)
    if blue.any():
        for _ in range(BLUE_REFINE_ATTEMPTS + 1):
            indices = sample_non_blue_indices(blue)
            frames = vr.get_batch(indices).asnumpy()
            selected_blue = blue_screen_flags(frames)
            if not selected_blue.any() or blue.all():
                break
            radius = max(1, int(np.ceil(n_frames / (2 * N_FRAMES))))
            for index in indices[selected_blue]:
                blue[max(0, index - radius):min(n_frames, index + radius + 1)] = True
    local_seconds = vr.get_frame_timestamp(indices)[:, 0]
    frames = letterbox_frames(frames)
    log.debug(
        "Video: %d frames @ %.3f fps; excluded %d blue frames; selected %d non-blue bin centres",
        n_frames, base_fps, int(blue.sum()), len(indices),
    )
    del vr

    # Qwen3's processor averages timestamps per temporal patch (2 frames), so
    # metadata must have one entry per decoded frame and use the absolute source
    # timeline.  The request endpoint is in seconds; the training manifest used
    # the same millisecond grid and ``precomputed_absolute_pts`` backend.
    absolute_ms = np.rint((float(req.start_time) + local_seconds) * 1000.0).astype(np.int64)
    metadata = {
        "fps": 1000,
        "frames_indices": absolute_ms.tolist(),
        "total_num_frames": int(absolute_ms.max()) + 1,
        "width": FRAME_WIDTH,
        "height": FRAME_HEIGHT,
        "duration": (int(absolute_ms.max()) + 1) / 1000.0,
        "video_backend": "precomputed_absolute_pts",
    }
    return frames, metadata


def answer_one(model, processor, req: Request) -> tuple[str, dict]:
    """Run the merged Qwen3.5 checkpoint on one official request."""
    frames, metadata = your_preprocessing(clip_path_for(req), req)
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": [{"type": "video"}, {"type": "text", "text": req.question}]},
    ]
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
    )
    inputs = processor(
        text=[text],
        videos=[frames],
        video_metadata=[metadata],
        do_sample_frames=False,
        return_tensors="pt",
    )
    model_device = next(model.parameters()).device
    inputs = inputs.to(model_device)
    tokenizer = getattr(processor, "tokenizer", processor)
    eos_token_ids = resolve_eos_token_ids(model, tokenizer)
    im_start_token_ids = resolve_im_start_token_ids(tokenizer)
    generation_kwargs = {"max_new_tokens": MAX_NEW_TOKENS, "do_sample": False}
    if eos_token_ids:
        generation_kwargs["eos_token_id"] = eos_token_ids if len(eos_token_ids) > 1 else eos_token_ids[0]
    with torch.inference_mode():
        generated = model.generate(**inputs, **generation_kwargs)
    if hasattr(generated, "sequences"):
        generated = generated.sequences
    new_token_rows = extract_new_token_ids(
        generated,
        inputs["input_ids"],
        inputs.get("attention_mask"),
    )
    trimmed = trim_generated_token_ids(
        new_token_rows[0] if new_token_rows else [],
        eos_token_ids=eos_token_ids,
        im_start_token_ids=im_start_token_ids,
        max_new_tokens=MAX_NEW_TOKENS,
    )
    decoded = processor.batch_decode(
        [trimmed["token_ids"]], skip_special_tokens=False, clean_up_tokenization_spaces=False
    )[0]
    answer = clean_decoded_answer(decoded)
    audit = {
        "raw_token_ids": trimmed["raw_token_ids"],
        "stop_index": trimmed["stop_index"],
        "eos_position": trimmed["eos_position"],
        "eos_positions": trimmed["eos_positions"],
        "stop_reason": trimmed["stop_reason"],
        "hit_max_new_tokens": trimmed["hit_max_new_tokens"],
        "new_token_count": len(trimmed["raw_token_ids"]),
    }
    return answer, audit


def run() -> int:
    t_start = time.monotonic()
    log.info("=== ORena SAVE FOCUS — inference start ===")

    # ── Device ────────────────────────────────────────────────────────────────
    # On the platform you get one GPU, chosen by you when you create the algorithm:
    # an NVIDIA L40S (48 GB VRAM, 4 vCPUs, 32 GiB RAM) or an NVIDIA RTX PRO 6000
    # Blackwell (96 GB VRAM, 8 vCPUs, 64 GiB RAM). This image runs on both — see
    # the hardware table in the README.
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    log.info("Device: %s", device)
    log_environment(device)

    # ── Inputs ─────────────────────────────────────────────────────────────────
    log.info("--- Loading inputs ---")
    for name in ("batch.json", "request.json", "FO_definitions.json"):
        p = INPUT_PATH / name
        status = f"{p.stat().st_size} bytes" if p.exists() else "MISSING"
        log.info("  %-24s %s", name, status)
    log.info("  %-24s %s", f"{VIDEO_DIR.name}/", f"{len(list(VIDEO_DIR.glob('*.mp4')))} clip(s)" if VIDEO_DIR.is_dir() else "MISSING")

    # request.json holds the whole batch: answer every question it lists.
    requests = load_requests(INPUT_PATH / "request.json")
    if not requests:
        log.error("request.json contains no requests")
        return 1
    log.info("Batch of %d question(s)", len(requests))

    # FO_definitions.json holds plain text, JSON-encoded; json.loads() gives the
    # text back. The definitions are the same for every question.
    fo_definitions = json.loads((INPUT_PATH / "FO_definitions.json").read_text())
    log.info("FO definitions: %d chars", len(fo_definitions))

    # ── Model — loaded ONCE for the whole batch ────────────────────────────────
    # Keep this outside the per-question loop. Loading the model is paid once per
    # run, and the latency budget grants a setup allowance to cover it.
    log.info("--- Loading model (once for the batch) ---")
    if not MODEL_PATH.exists():
        log.error("Merged model directory not found: %s", MODEL_PATH)
        return 1
    log.info("Merged model path: %s", MODEL_PATH)
    processor = AutoProcessor.from_pretrained(MODEL_PATH, local_files_only=True)
    model = Qwen3_5ForConditionalGeneration.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        low_cpu_mem_usage=True,
        local_files_only=True,
    )
    model.eval()
    n_params = sum(p.numel() for p in model.parameters())
    log.info("Model loaded and set to eval mode: %d parameters (setup %.2f s)", n_params, time.monotonic() - t_start)

    if device.type == "cuda":
        allocated = torch.cuda.memory_allocated(0) / 1024**2
        reserved = torch.cuda.memory_reserved(0) / 1024**2
        log.info("  VRAM after model load: %.1f MiB allocated / %.1f MiB reserved", allocated, reserved)

    # ── Inference — one pass per question ──────────────────────────────────────
    log.info("--- Running inference over %d question(s) ---", len(requests))
    responses = []
    generation_audits = []
    n_failed = 0
    t_batch = time.monotonic()

    for i, req in enumerate(requests, start=1):
        log.info(
            "[%d/%d] qID=%s  videoID=%s  window=[%.2f, %.2f] s (%.2f s)",
            i, len(requests), req.qID, req.videoID, req.start_time, req.end_time, req.duration,
        )
        t0 = time.monotonic()
        try:
            answer, audit = answer_one(model, processor, req)
        except Exception:
            # Keep going: one failed question should not cost you the rest of the
            # batch. An empty answer is simply scored incorrect.
            n_failed += 1
            log.exception("[%d/%d] qID=%s failed; emitting empty answer", i, len(requests), req.qID)
            answer = ""
            audit = {
                "raw_token_ids": [],
                "stop_index": None,
                "eos_position": None,
                "eos_positions": [],
                "stop_reason": "error",
                "hit_max_new_tokens": False,
                "new_token_count": 0,
            }
        latency = time.monotonic() - t0

        # Informational only. The platform times the run itself and scores against
        # that measurement — see "The latency budget" in the module docstring.
        responses.append(Response(qID=req.qID, content=answer, latency=latency))
        generation_audits.append({"qID": req.qID, **audit})
        log.info(
            "[%d/%d] qID=%s answered in %.4f s answer_len=%d stop_reason=%s",
            i,
            len(requests),
            req.qID,
            latency,
            len(answer),
            audit["stop_reason"],
        )

    batch_seconds = time.monotonic() - t_batch
    log.info(
        "Inference complete: %d answered (%d failed) in %.2f s (%.3f s/question)",
        len(responses), n_failed, batch_seconds, batch_seconds / max(len(responses), 1),
    )

    if device.type == "cuda":
        allocated = torch.cuda.memory_allocated(0) / 1024**2
        log.debug("  VRAM after inference: %.1f MiB allocated", allocated)

    # ── Output ─────────────────────────────────────────────────────────────────
    log.info("--- Writing output ---")
    OUTPUT_PATH.mkdir(parents=True, exist_ok=True)
    output_path = OUTPUT_PATH / "answer.json"
    save_items(responses, output_path)
    log.info("Wrote %d response(s) to %s", len(responses), output_path)
    audit_path = OUTPUT_PATH / "generation_audit.json"
    audit_path.write_text(json.dumps(generation_audits, ensure_ascii=False, indent=2) + "\n")
    log.info("Wrote generation audit for %d response(s) to %s", len(generation_audits), audit_path)
    log.info("=== inference done in %.2f s total ===", time.monotonic() - t_start)
    return 0


if __name__ == "__main__":
    raise SystemExit(run())