media-batch-tools/tools/video_timeline.py

from __future__ import annotations

import statistics
from dataclasses import dataclass

from tools.avisynth_validate import FrameIdentity
from tools.video_probe import VideoProbe


@dataclass(frozen=True)
class OutputFrameTiming:
    output_index: int
    source_id: int
    source_frame: int
    start: float
    duration: float


@dataclass(frozen=True)
class OutputTimeline:
    frames: list[OutputFrameTiming]
    audio_segments: list[tuple[float, float]]
    consumed_source_frames: list[int]
    duration_seconds: float
    timing_kind: str
    beginning_trimmed: bool
    ending_trimmed: bool
    dropped_frame_count: int
    source_frame_count: int
    matrix: int | None


def build_output_timeline(
    identities: list[FrameIdentity],
    *,
    probe: VideoProbe,
) -> OutputTimeline:
    source_durations = _source_frame_durations(probe)
    kept_frames: list[OutputFrameTiming] = []
    audio_segments: list[tuple[float, float]] = []
    consumed_source_frames: list[int] = []
    matrices: set[int] = set()
    dropped_count = 0
    leading_trimmed = False
    last_consumed_source_frame: int | None = None

    for frame in sorted(identities, key=lambda item: item.output_frame):
        if frame.source_id is None or frame.source_frame is None:
            continue
        if frame.source_frame >= len(source_durations):
            raise RuntimeError(
                f"source frame {frame.source_frame} is outside probed frame table "
                f"for {probe.path}"
            )

        duration = source_durations[frame.source_frame]
        if frame.drop_frame:
            dropped_count += 1
            if kept_frames:
                consumed_source_frames.append(frame.source_frame)
                _append_audio_segment(
                    audio_segments,
                    start=probe.frame_timestamps[frame.source_frame],
                    duration=duration,
                )
                previous = kept_frames[-1]
                kept_frames[-1] = OutputFrameTiming(
                    output_index=previous.output_index,
                    source_id=previous.source_id,
                    source_frame=previous.source_frame,
                    start=previous.start,
                    duration=previous.duration + duration,
                )
                last_consumed_source_frame = frame.source_frame
            else:
                leading_trimmed = True
            continue

        consumed_source_frames.append(frame.source_frame)
        if frame.matrix is not None:
            matrices.add(frame.matrix)
        _append_audio_segment(
            audio_segments,
            start=probe.frame_timestamps[frame.source_frame],
            duration=duration,
        )
        kept_frames.append(
            OutputFrameTiming(
                output_index=len(kept_frames),
                source_id=frame.source_id,
                source_frame=frame.source_frame,
                start=probe.frame_timestamps[frame.source_frame],
                duration=duration,
            )
        )
        last_consumed_source_frame = frame.source_frame

    if not kept_frames:
        raise RuntimeError(f"No kept output frames after validation for {probe.path}")
    if len(matrices) > 1:
        raise RuntimeError(
            f"Output frames have mixed _Matrix values for {probe.path}: "
            + ", ".join(str(value) for value in sorted(matrices))
        )

    source_count = len(probe.frame_timestamps)
    beginning_trimmed = leading_trimmed or kept_frames[0].source_frame > 0
    if last_consumed_source_frame is None:
        last_consumed_source_frame = kept_frames[-1].source_frame
    ending_trimmed = last_consumed_source_frame < source_count - 1
    durations = [frame.duration for frame in kept_frames]
    return OutputTimeline(
        frames=kept_frames,
        audio_segments=audio_segments,
        consumed_source_frames=consumed_source_frames,
        duration_seconds=sum(durations),
        timing_kind=_classify_durations(durations),
        beginning_trimmed=beginning_trimmed,
        ending_trimmed=ending_trimmed,
        dropped_frame_count=dropped_count,
        source_frame_count=source_count,
        matrix=next(iter(matrices)) if matrices else None,
    )


def _append_audio_segment(
    segments: list[tuple[float, float]],
    *,
    start: float,
    duration: float,
) -> None:
    if not segments:
        segments.append((start, duration))
        return

    previous_start, previous_duration = segments[-1]
    previous_end = previous_start + previous_duration
    if abs(previous_end - start) <= 0.001:
        segments[-1] = (previous_start, previous_duration + duration)
    else:
        segments.append((start, duration))


def _source_frame_durations(probe: VideoProbe) -> list[float]:
    timestamps = probe.frame_timestamps
    if not timestamps:
        raise RuntimeError(f"No probed frame timestamps for {probe.path}")

    durations: list[float] = []
    for current, next_timestamp in zip(timestamps, timestamps[1:]):
        durations.append(max(0.0, next_timestamp - current))

    fallback = statistics.median(durations) if durations else 0.0
    if probe.duration_seconds is not None and len(timestamps) > 1:
        last_duration = probe.duration_seconds - timestamps[-1]
        if last_duration <= 0:
            last_duration = fallback
    else:
        last_duration = fallback
    durations.append(max(0.0, last_duration))
    return durations


def _classify_durations(durations: list[float]) -> str:
    if len(durations) < 2:
        return "unknown"
    median = statistics.median(durations)
    max_deviation = max(abs(duration - median) for duration in durations)
    return "vfr" if max_deviation > 0.002 else "cfr"