pipeline/frames.py

"""Stage 5 + 6: candidate frames and the weighted screenshot indicator.

Flow:
  * ``extract_candidates`` pulls a dense, downscaled grid of frames from the video in a
    single ffmpeg pass, *before* the video is deleted.
  * ``delete_video`` removes the source video (auto-delete once the transcript + frames
    exist).
  * ``select_screenshots`` runs *after* the LLM. For each step it computes a capture time
    that blends the LLM's suggested timestamp with the Whisper segment timing of the text
    the step quotes, then snaps that target to the nearest pre-extracted candidate frame.
"""
from __future__ import annotations

import difflib
import glob
import os
import subprocess

# A 1 fps grid gives <=0.5 s snapping error. Capped so a long video can't explode disk.
GRID_FPS = 1.0
MAX_CANDIDATES = 1500


def extract_candidates(video_path: str, out_dir: str, duration: float) -> list[tuple[float, str]]:
    """Extract a downscaled candidate-frame grid; return sorted ``[(time_s, path), ...]``."""
    os.makedirs(out_dir, exist_ok=True)
    eff_fps = GRID_FPS
    if duration and duration * GRID_FPS > MAX_CANDIDATES:
        eff_fps = MAX_CANDIDATES / duration  # thin the grid for very long videos

    pattern = os.path.join(out_dir, "cand_%06d.jpg")
    cmd = [
        "ffmpeg", "-y", "-i", video_path,
        "-vf", f"fps={eff_fps:.6f},scale=-2:720",
        "-q:v", "4", pattern,
    ]
    proc = subprocess.run(cmd, capture_output=True, text=True)
    files = sorted(glob.glob(os.path.join(out_dir, "cand_*.jpg")))
    if not files:
        raise RuntimeError(
            "ffmpeg extracted no candidate frames.\n" + proc.stderr[-600:]
        )
    # fps filter emits frame i (1-based) at ~ (i-1)/eff_fps seconds.
    return [((i) / eff_fps, f) for i, f in enumerate(files)]


def delete_video(video_path: str) -> None:
    """Delete the downloaded video (and any sibling temp media) if present."""
    for p in (video_path,):
        try:
            if p and os.path.exists(p):
                os.remove(p)
        except OSError:
            pass


def _clamp(x: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, x))


def best_match_segment(quote: str, t_llm: float, segs: list[dict]) -> dict:
    """Find the Whisper segment a step refers to.

    Prefer the segment whose text best matches ``quote``; if there's no decent textual
    match, fall back to the segment nearest ``t_llm``.
    """
    quote = (quote or "").strip().lower()
    if quote:
        scored = [
            (difflib.SequenceMatcher(None, quote, s["text"].lower()).ratio(), s)
            for s in segs
        ]
        ratio, best = max(scored, key=lambda t: t[0])
        if ratio >= 0.3:
            return best
    return min(segs, key=lambda s: abs((s["start"] + s["end"]) / 2 - t_llm))


def pick_time(step: dict, segs: list[dict], cand_times: list[float], *,
              w_llm: float = 0.4, w_whisper: float = 0.6, lead: float = 1.0,
              max_drift: float = 20.0) -> float:
    """Weighted capture time, snapped to the nearest candidate frame.

    ``t_whisper`` (segment mid + lead) is ground truth for *when* the referenced text is
    spoken; ``t_llm`` is the model's guess of *what* matters. We blend them, but if the
    LLM time drifts far from the matched segment we clamp it back into that segment so a
    hallucinated timestamp can't drag the shot off-topic.
    """
    seg = best_match_segment(step.get("quote", ""), float(step.get("t_llm", 0.0)), segs)
    t_whisper = (seg["start"] + seg["end"]) / 2 + lead
    t_llm = float(step.get("t_llm", t_whisper))
    if abs(t_llm - t_whisper) > max_drift:
        t_llm = _clamp(t_llm, seg["start"], seg["end"])

    total = w_llm + w_whisper or 1.0
    t_target = (w_llm * t_llm + w_whisper * t_whisper) / total
    return min(cand_times, key=lambda c: abs(c - t_target))


def select_screenshots(steps: list[dict], segs: list[dict],
                       candidates: list[tuple[float, str]], *,
                       w_llm: float = 0.4, w_whisper: float = 0.6, lead: float = 1.0,
                       max_shots: int = 8) -> dict[int, dict]:
    """Return ``{step_index: {"time", "path"}}`` for the steps that get a screenshot.

    When there are more steps than ``max_shots``, the highest-``importance`` steps win.
    """
    if not candidates:
        return {}
    cand_times = [t for t, _ in candidates]
    time_to_path = dict(candidates)

    order = sorted(range(len(steps)), key=lambda i: -float(steps[i].get("importance", 0.5)))
    chosen = set(order[:max_shots])

    out: dict[int, dict] = {}
    for i, step in enumerate(steps):
        if i not in chosen:
            continue
        t = pick_time(step, segs, cand_times, w_llm=w_llm, w_whisper=w_whisper, lead=lead)
        out[i] = {"time": t, "path": time_to_path[t]}
    return out