pipeline/frames.py

"""Stage 5 + 6: real screenshots via yt-dlp stream URL + ffmpeg, at weighted timestamps.

No full download. We resolve a direct video stream URL with yt-dlp (metadata only), then
``ffmpeg -ss T -frames:v 1`` grabs a single frame per timestamp over HTTP range requests
(only a few seconds of data each). For each shot we grab 3 candidates around T and keep
the sharpest (Laplacian variance) to avoid landing on a blurry/transition frame.

The capture timestamp itself comes from the weighted indicator: a blend of the LLM's
suggested timestamp and the transcript segment timing of the text the step quotes.
"""
from __future__ import annotations

import difflib
import os
import subprocess

from yt_dlp import YoutubeDL

# Format preference: muxed (single A/V url) <=720p first, then video-only, then anything.
_FORMAT = ("best[height<=720][vcodec!=none][acodec!=none]/"
           "best[height<=720][vcodec!=none]/bestvideo[height<=720]/best[vcodec!=none]/best")


# --------------------------------------------------------------------------- weighting
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:
    """The transcript segment a step refers to: best text match, else nearest to t_llm."""
    quote = (quote or "").strip().lower()
    if quote:
        ratio, best = max(
            ((difflib.SequenceMatcher(None, quote, s["text"].lower()).ratio(), s) for s in segs),
            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], *, w_llm: float = 0.4, w_whisper: float = 0.6,
              lead: float = 1.0, max_drift: float = 20.0) -> float:
    """Weighted capture time (seconds). No snapping — we grab on demand at any T.

    Transcript timing is ground truth for *when* the text is spoken; the LLM time is its
    guess of *what* matters. Blend them, clamping a far-off LLM time back into the matched
    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
    return max(0.0, (w_llm * t_llm + w_whisper * t_whisper) / total)


def compute_shot_times(steps: list[dict], segs: list[dict], *, w_llm: float = 0.4,
                       w_whisper: float = 0.6, lead: float = 1.0,
                       max_shots: int = 8) -> dict[int, float]:
    """Return ``{step_index: capture_time}`` for the top-``max_shots`` steps by importance."""
    if not segs:
        return {}
    order = sorted(range(len(steps)), key=lambda i: -float(steps[i].get("importance", 0.5)))
    chosen = set(order[:max_shots])
    return {
        i: pick_time(steps[i], segs, w_llm=w_llm, w_whisper=w_whisper, lead=lead)
        for i in range(len(steps)) if i in chosen
    }


# ----------------------------------------------------------------------- stream + grab
def get_stream_url(video_id: str, cookiefile: str | None = None,
                   proxy: str | None = None) -> tuple[str, float]:
    """Resolve a direct video stream URL (no download); return ``(url, duration_seconds)``."""
    opts = {"quiet": True, "no_warnings": True, "skip_download": True, "format": _FORMAT}
    if cookiefile:
        opts["cookiefile"] = cookiefile
    if proxy:
        opts["proxy"] = proxy
    with YoutubeDL(opts) as ydl:
        info = ydl.extract_info(f"https://www.youtube.com/watch?v={video_id}", download=False)

    stream = info.get("url")
    if not stream:
        for f in info.get("requested_formats") or []:
            if f.get("vcodec") not in (None, "none") and f.get("url"):
                stream = f["url"]
                break
    if not stream:
        vids = [f for f in info.get("formats", [])
                if f.get("vcodec") not in (None, "none") and f.get("url")]
        if vids:
            stream = sorted(vids, key=lambda f: f.get("height") or 0)[-1]["url"]
    if not stream:
        raise RuntimeError("Could not resolve a video stream URL for frame capture.")
    return stream, float(info.get("duration") or 0.0)


def _grab(stream_url: str, t: float, out_path: str, proxy: str | None) -> bool:
    """Grab one frame at time ``t`` via ffmpeg input-seek; return True on success."""
    cmd = ["ffmpeg", "-y", "-loglevel", "error", "-ss", f"{max(0.0, t):.2f}",
           "-i", stream_url, "-frames:v", "1", "-q:v", "2", out_path]
    env = dict(os.environ)
    if proxy:  # stream URL is bound to the extraction IP; route ffmpeg through it too
        env["http_proxy"] = env["https_proxy"] = proxy
    try:
        subprocess.run(cmd, capture_output=True, text=True, timeout=90, env=env)
    except subprocess.TimeoutExpired:
        return False
    return os.path.exists(out_path) and os.path.getsize(out_path) > 0


def _sharpness(path: str) -> float:
    """Variance of the Laplacian (higher = sharper). Used to reject blurry frames."""
    import numpy as np
    from PIL import Image

    im = np.asarray(Image.open(path).convert("L"), dtype=np.float64)
    if im.shape[0] < 3 or im.shape[1] < 3:
        return 0.0
    lap = (-4 * im[1:-1, 1:-1] + im[:-2, 1:-1] + im[2:, 1:-1]
           + im[1:-1, :-2] + im[1:-1, 2:])
    return float(lap.var())


def grab_best_frame(stream_url: str, t: float, out_dir: str, idx: int,
                    proxy: str | None = None) -> str | None:
    """Grab candidates around ``t`` and keep the sharpest; return its path or None."""
    cands = []
    for k, dt in enumerate((-0.5, 0.0, 0.5)):
        p = os.path.join(out_dir, f"shot_{idx}_{k}.jpg")
        if _grab(stream_url, t + dt, p, proxy):
            cands.append(p)
    if not cands:
        return None
    best = max(cands, key=_sharpness)
    for p in cands:
        if p != best:
            try:
                os.remove(p)
            except OSError:
                pass
    return best


def capture_shots(times: dict[int, float], stream_url: str, out_dir: str,
                  proxy: str | None = None, progress=None) -> dict[int, dict]:
    """Grab a frame for each ``{step_index: time}``; return ``{idx: {"time", "path"}}``."""
    os.makedirs(out_dir, exist_ok=True)
    out: dict[int, dict] = {}
    items = list(times.items())
    for n, (idx, t) in enumerate(items):
        if progress:
            progress((n + 1) / max(1, len(items)), desc=f"Screenshot {n + 1}/{len(items)}")
        path = grab_best_frame(stream_url, t, out_dir, idx, proxy)
        if path:
            out[idx] = {"time": t, "path": path}
    return out