WeAreSyncing/stream_edges.py

"""Offset estimation from where the video and audio tracks start and end,
as an optional alternative to marking a sync event by hand.

Premise: a recorder that stops both tracks at the same real-world moment,
but whose audio started capturing late (or early) while still being
written as starting at t=0, produces a file where the audio runs out
before (or after) the video by exactly the sync error. So

    delay_ms = (video_end - audio_end) * 1000

using the same sign convention as sync_logic.compute_offset (positive:
audio is ahead and gets delayed; negative: video gets delayed).

This is only an *estimate*: a file whose audio genuinely stopped early
(a mic cutting out before the camera did) looks identical but has nothing
to do with sync. The start-based figure is reported alongside it as a
cross-check -- it's what carries the offset when the container encodes it
as a stream start_time instead, and a disagreement between the two is
worth surfacing.
"""
import subprocess
from dataclasses import dataclass
from typing import List, Optional, Tuple

# How far before the end of the file to start reading packets. ffprobe
# seeks there instead of scanning the whole file, which keeps this fast on
# multi-hour recordings. If a track ends earlier than this window (so no
# packets for it turn up), the whole track is scanned instead.
TAIL_WINDOW_SEC = 30.0

# Start- and end-based estimates are compared against this. AAC encoder
# delay alone typically puts the start ~20-45ms off, so anything much
# tighter would warn on perfectly ordinary files.
DISAGREEMENT_TOLERANCE_MS = 50.0


class StreamEdgesError(RuntimeError):
    pass


@dataclass(frozen=True)
class StreamEdges:
    video_start: float
    video_end: float
    audio_start: float
    audio_end: float

    @property
    def end_delay_ms(self) -> float:
        return (self.video_end - self.audio_end) * 1000.0

    @property
    def start_delay_ms(self) -> float:
        return (self.video_start - self.audio_start) * 1000.0

    @property
    def estimates_disagree(self) -> bool:
        return abs(self.end_delay_ms - self.start_delay_ms) > DISAGREEMENT_TOLERANCE_MS


def _ffprobe(args: List[str]) -> str:
    result = subprocess.run(["ffprobe", "-v", "error", *args], capture_output=True, text=True)
    if result.returncode != 0:
        raise StreamEdgesError(f"ffprobe failed: {result.stderr.strip()}")
    return result.stdout


def _parse_float(value: str) -> Optional[float]:
    try:
        return float(value)
    except ValueError:
        return None  # "N/A"


def _format_duration(input_path: str) -> Optional[float]:
    out = _ffprobe(["-show_entries", "format=duration", "-of", "csv=p=0", input_path])
    return _parse_float(out.strip())


def _stream_start(input_path: str, stream_specifier: str) -> float:
    out = _ffprobe([
        "-select_streams", stream_specifier,
        "-show_entries", "stream=start_time", "-of", "csv=p=0", input_path,
    ])
    start = _parse_float(out.strip())
    return start if start is not None else 0.0


def _packets(input_path: str, stream_specifier: str, read_from: Optional[float]) -> List[Tuple[float, Optional[float]]]:
    args = ["-select_streams", stream_specifier]
    if read_from is not None:
        args += ["-read_intervals", f"{read_from:.3f}%"]
    args += ["-show_entries", "packet=pts_time,duration_time", "-of", "csv=p=0", input_path]
    packets = []
    for line in _ffprobe(args).splitlines():
        fields = line.strip().rstrip(",").split(",")
        pts = _parse_float(fields[0]) if fields and fields[0] else None
        if pts is None:
            continue
        duration = _parse_float(fields[1]) if len(fields) > 1 else None
        packets.append((pts, duration))
    return packets


def _stream_end(input_path: str, stream_specifier: str, file_duration: Optional[float]) -> float:
    """End of the last packet (pts + duration), not the container's stored
    duration -- that's frequently missing (e.g. per-stream in MKV) or
    rounded, and this needs to be accurate to well under a frame."""
    read_from = None
    if file_duration is not None and file_duration > TAIL_WINDOW_SEC:
        read_from = file_duration - TAIL_WINDOW_SEC
    packets = _packets(input_path, stream_specifier, read_from)
    if not packets and read_from is not None:
        packets = _packets(input_path, stream_specifier, None)
    if not packets:
        raise StreamEdgesError(f"no packets found for stream {stream_specifier}")

    # Max rather than last: video with B-frames isn't stored in pts order.
    ends = [pts + duration for pts, duration in packets if duration is not None]
    if ends:
        return max(ends)
    # No per-packet durations (some containers): fall back to the gap
    # between the last two packets as the final packet's duration.
    pts_sorted = sorted(pts for pts, _ in packets)
    step = pts_sorted[-1] - pts_sorted[-2] if len(pts_sorted) > 1 else 0.0
    return pts_sorted[-1] + step


def probe_stream_edges(
    input_path: str,
    video_stream_index: Optional[int] = None,
    audio_stream_index: Optional[int] = None,
) -> StreamEdges:
    """Stream indices are ffmpeg global indices (mpv's `ff-index`), so the
    estimate uses the same tracks the player is showing -- see
    clap_detection.py for why that matters on multi-track files."""
    video_spec = f"{video_stream_index}" if video_stream_index is not None else "v:0"
    audio_spec = f"{audio_stream_index}" if audio_stream_index is not None else "a:0"
    file_duration = _format_duration(input_path)
    return StreamEdges(
        video_start=_stream_start(input_path, video_spec),
        video_end=_stream_end(input_path, video_spec, file_duration),
        audio_start=_stream_start(input_path, audio_spec),
        audio_end=_stream_end(input_path, audio_spec, file_duration),
    )