DV video convertor

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import math
import struct
import subprocess
import sys
from pathlib import Path
from typing import List, Tuple

MAGIC_DV = 0x21215644  # "DV!!" little-endian

def run(cmd: List[str], *, check: bool = True) -> subprocess.CompletedProcess:
    p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if check and p.returncode != 0:
        raise RuntimeError(
            "Command failed:\n  {}\n\nstdout:\n{}\n\nstderr:\n{}".format(
                " ".join(cmd), p.stdout, p.stderr
            )
        )
    return p


def ffprobe_json(path: Path) -> dict:
    p = run(
        [
            "ffprobe",
            "-hide_banner",
            "-loglevel",
            "error",
            "-print_format",
            "json",
            "-show_format",
            "-show_streams",
            "-show_packets",
            "-select_streams",
            "v:0",
            str(path),
        ],
        check=True,
    )
    return json.loads(p.stdout)


def ffprobe_format_duration(path: Path) -> float:
    p = run(
        [
            "ffprobe",
            "-hide_banner",
            "-loglevel",
            "error",
            "-show_entries",
            "format=duration",
            "-of",
            "default=nw=1:nk=1",
            str(path),
        ],
        check=True,
    )
    return float(p.stdout.strip())


def parse_header_and_table(dv_bytes: bytes) -> Tuple[Tuple[int, ...], List[Tuple[int, int, int, int, int]]]:
    if len(dv_bytes) < 76:
        raise ValueError("File too small to be DV!!")

    hdr = struct.unpack_from("<19I", dv_bytes, 0)
    if hdr[0] != MAGIC_DV:
        raise ValueError("Bad magic: not DV!!")

    n = hdr[7]
    table_off = 76
    table_len = 20 * n
    if len(dv_bytes) < table_off + table_len:
        raise ValueError("Truncated: packet table incomplete")

    entries = [struct.unpack_from("<5I", dv_bytes, table_off + 20 * i) for i in range(n)]
    return hdr, entries


def extract_streams(dv_path: Path, out_dir: Path) -> Tuple[Path, Path, int]:
    """
    Extract:
      - video: raw MPEG4 elementary stream (.m4v)
      - audio: Ogg/Vorbis (.ogg) = preroll + concatenated audio tails
    """
    dv_bytes = dv_path.read_bytes()
    hdr, entries = parse_header_and_table(dv_bytes)

    n = hdr[7]
    preroll_len = hdr[9]

    table_off = 76
    data_off = table_off + 20 * n

    if len(dv_bytes) < data_off + preroll_len:
        raise ValueError("Truncated: missing audio preroll")

    preroll = dv_bytes[data_off : data_off + preroll_len]
    packet_region_off = data_off + preroll_len

    sizes = [(e[0] & 0x7FFFFFFF) for e in entries]
    total_packets = sum(sizes)

    if len(dv_bytes) < packet_region_off + total_packets:
        raise ValueError("Truncated: packet region shorter than expected")

    video = bytearray()
    audio = bytearray(preroll)

    pos = packet_region_off
    for i in range(n):
        sz = sizes[i]
        audio_bytes = entries[i][3]
        chunk = dv_bytes[pos : pos + sz]
        if len(chunk) != sz:
            raise ValueError(f"Truncated packet {i}: expected {sz}, got {len(chunk)}")
        if audio_bytes:
            if audio_bytes > sz:
                raise ValueError(f"Packet {i}: audio_bytes ({audio_bytes}) > size ({sz})")
            vlen = sz - audio_bytes
            video += chunk[:vlen]
            audio += chunk[vlen:]
        else:
            video += chunk
        pos += sz

    out_dir.mkdir(parents=True, exist_ok=True)
    video_path = out_dir / f"{dv_path.stem}.video.m4v"
    audio_path = out_dir / f"{dv_path.stem}.audio.ogg"
    video_path.write_bytes(video)
    audio_path.write_bytes(audio)

    return video_path, audio_path, n


def mux_mkv_with_fixed_timestamps(
    video_path: Path,
    audio_path: Path,
    out_path: Path,
    *,
    fps: int,
) -> None:
    """
    Stream-copy mux, but inject correct timestamps into the raw MPEG4 video
    using setts bitstream filter.

    We use time_base=1/1000 and 25fps => 40ms per frame:
      pts = 40*N, dts = 40*N, duration = 40
    """
    frame_ms = int(round(1000 / fps))  # 25 -> 40
    bsf = f"setts=pts={frame_ms}*N:dts={frame_ms}*N:duration={frame_ms}:time_base=1/1000"

    run(
        [
            "ffmpeg",
            "-hide_banner",
            "-y",
            "-loglevel",
            "error",
            "-i",
            str(video_path),
            "-i",
            str(audio_path),
            "-map",
            "0:v:0",
            "-map",
            "1:a:0",
            "-c",
            "copy",
            "-bsf:v",
            bsf,
            str(out_path),
        ],
        check=True,
    )


def validate_output(out_path: Path, *, expected_frames: int, audio_path: Path, fps: int) -> None:
    # 1) Ensure output is decodable end-to-end (fatal errors => non-zero)
    run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-i", str(out_path), "-f", "null", "-"], check=True)

    # 2) Video packet count == expected frame count
    j = ffprobe_json(out_path)
    packets = j.get("packets", [])
    if len(packets) != expected_frames:
        raise RuntimeError(f"Bad mux: video packet count {len(packets)} != expected {expected_frames}")

    # 3) Duration match (audio duration is the reliable one here)
    out_dur = ffprobe_format_duration(out_path)
    aud_dur = ffprobe_format_duration(audio_path)

    # Allow small container rounding error (Matroska timebase granularity)
    if abs(out_dur - aud_dur) > 0.10:
        raise RuntimeError(f"Bad mux: duration mismatch. out={out_dur:.3f}s audio={aud_dur:.3f}s")

    # 4) Sanity: expected duration from frames/fps roughly matches too
    expected_dur = expected_frames / fps
    if abs(out_dur - expected_dur) > 0.25:
        raise RuntimeError(
            f"Bad mux: duration vs frames/fps mismatch. out={out_dur:.3f}s expected~{expected_dur:.3f}s"
        )


def main() -> int:
    ap = argparse.ArgumentParser(description="Convert DV!! .dv to a playable MKV without re-encoding.")
    ap.add_argument("input", type=Path, help="Input DV!! file (e.g. kleo19.dv)")
    ap.add_argument("--out-dir", type=Path, default=Path("out"), help="Output directory")
    ap.add_argument("--fps", type=int, default=25, help="Video FPS (kleo19.dv is 25)")
    args = ap.parse_args()

    dv_path: Path = args.input
    if not dv_path.exists():
        print(f"Not found: {dv_path}", file=sys.stderr)
        return 2

    out_dir: Path = args.out_dir
    video_path, audio_path, frames = extract_streams(dv_path, out_dir)

    out_mkv = out_dir / f"{dv_path.stem}.mkv"
    mux_mkv_with_fixed_timestamps(video_path, audio_path, out_mkv, fps=args.fps)

    validate_output(out_mkv, expected_frames=frames, audio_path=audio_path, fps=args.fps)

    print("OK")
    print(f"  video: {video_path}")
    print(f"  audio: {audio_path}")
    print(f"  mkv:   {out_mkv}")
    print(f"  frames: {frames}, fps: {args.fps}")
    return 0


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