DV to OGM video converter

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

import argparse
import random
import struct
import subprocess
import sys
from pathlib import Path
from typing import Dict, Iterable, List, Tuple

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


# -----------------------------
# Basic helpers
# -----------------------------
def eprint(msg: str) -> None:
    print(msg, file=sys.stderr)


def run(cmd: List[str]) -> int:
    p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if p.returncode != 0:
        eprint("Command failed:\n  " + " ".join(cmd))
        if p.stdout.strip():
            eprint("stdout:\n" + p.stdout)
        if p.stderr.strip():
            eprint("stderr:\n" + p.stderr)
    return p.returncode


# -----------------------------
# DV!! extraction (same structure you already validated)
# -----------------------------
def parse_dvbang(dv_bytes: bytes) -> Tuple[Tuple[int, ...], bytes, bytes]:
    if len(dv_bytes) < 76:
        raise ValueError("File too small")

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

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

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

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

    data_off = table_off + table_len
    if len(dv_bytes) < data_off + preroll_len:
        raise ValueError("Truncated: preroll")

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

    video = bytearray()
    audio = bytearray(preroll)

    for i, e in enumerate(entries):
        size = e[0] & 0x7FFFFFFF
        audio_bytes = e[3]

        chunk = dv_bytes[pos : pos + size]
        if len(chunk) != size:
            raise ValueError(f"Truncated packet {i}: expected {size}, got {len(chunk)}")
        pos += size

        if audio_bytes:
            if audio_bytes > size:
                raise ValueError(f"Packet {i}: audio_bytes > size")
            vlen = size - audio_bytes
            video += chunk[:vlen]
            audio += chunk[vlen:]
        else:
            video += chunk

    return hdr, bytes(video), bytes(audio)


# -----------------------------
# Ogg primitives (page builder + CRC)
# -----------------------------
def _ogg_crc_table() -> List[int]:
    poly = 0x04C11DB7
    table: List[int] = []
    for i in range(256):
        r = i << 24
        for _ in range(8):
            if r & 0x80000000:
                r = ((r << 1) & 0xFFFFFFFF) ^ poly
            else:
                r = (r << 1) & 0xFFFFFFFF
        table.append(r)
    return table


CRC_TABLE = _ogg_crc_table()


def ogg_crc(page_bytes: bytes) -> int:
    crc = 0
    for b in page_bytes:
        crc = ((crc << 8) & 0xFFFFFFFF) ^ CRC_TABLE[((crc >> 24) & 0xFF) ^ b]
    return crc


def build_ogg_page(
    packets: List[bytes],
    serial: int,
    seq: int,
    granule: int,
    header_type: int,
) -> bytes:
    segs: List[int] = []
    body = b""

    for pkt in packets:
        body += pkt
        n = len(pkt)
        while n >= 255:
            segs.append(255)
            n -= 255
        segs.append(n)

    nseg = len(segs)
    if nseg > 255:
        raise ValueError("Too many segments in one page (split packets/pages)")

    header = bytearray()
    header += b"OggS"
    header += bytes([0])  # version
    header += bytes([header_type])
    header += struct.pack("<Q", granule)
    header += struct.pack("<I", serial)
    header += struct.pack("<I", seq)
    header += struct.pack("<I", 0)  # CRC placeholder
    header += bytes([nseg])
    header += bytes(segs)

    page = bytes(header) + body
    crc = ogg_crc(page)
    page = page[:22] + struct.pack("<I", crc) + page[26:]
    return page


# -----------------------------
# Parse existing Vorbis Ogg pages so we can copy them exactly
# -----------------------------
def ogg_iter_pages(data: bytes) -> Iterable[Dict]:
    i = 0
    while True:
        j = data.find(b"OggS", i)
        if j < 0:
            return
        if j + 27 > len(data):
            return
        header_type = data[j + 5]
        granule = struct.unpack_from("<Q", data, j + 6)[0]
        serial = struct.unpack_from("<I", data, j + 14)[0]
        seq = struct.unpack_from("<I", data, j + 18)[0]
        nseg = data[j + 26]
        segtable = data[j + 27 : j + 27 + nseg]
        body_len = sum(segtable)
        body_start = j + 27 + nseg
        body = data[body_start : body_start + body_len]
        page_bytes = data[j : body_start + body_len]
        yield {
            "header_type": header_type,
            "granule": granule,
            "serial": serial,
            "seq": seq,
            "segtable": segtable,
            "body": body,
            "bytes": page_bytes,
        }
        i = body_start + body_len


def vorbis_sample_rate_from_ogg_pages(pages: List[Dict]) -> int:
    # Reconstruct first packet (Vorbis identification): 0x01 "vorbis" ...
    cur = bytearray()
    for p in pages:
        off = 0
        for s in p["segtable"]:
            cur += p["body"][off : off + s]
            off += s
            if s < 255:
                pkt = bytes(cur)
                # Vorbis ID packet begins with 0x01 + "vorbis"
                if len(pkt) >= 16 and pkt[0] == 0x01 and pkt[1:7] == b"vorbis":
                    # version(4) at 7..10, channels(1) at 11, sample_rate(4) at 12..15
                    return struct.unpack_from("<I", pkt, 12)[0]
                return 0
    return 0


def vorbis_header_end_page_index(pages: List[Dict]) -> int:
    # Count packets across pages until we have the first 3 Vorbis packets (id/comment/setup)
    pkt_count = 0
    cur = bytearray()
    for idx, p in enumerate(pages):
        off = 0
        for s in p["segtable"]:
            cur += p["body"][off : off + s]
            off += s
            if s < 255:
                pkt_count += 1
                cur = bytearray()
                if pkt_count >= 3:
                    return idx
    return len(pages) - 1


# -----------------------------
# MPEG-4 Part2 elementary stream framing (split at VOP startcode)
# -----------------------------
def split_mpeg4_vop_frames(es: bytes) -> List[bytes]:
    vop = b"\x00\x00\x01\xb6"
    idxs: List[int] = []
    i = 0
    while True:
        j = es.find(vop, i)
        if j < 0:
            break
        idxs.append(j)
        i = j + 4

    if not idxs:
        raise ValueError("No MPEG-4 VOP start codes found (0x000001B6)")

    header = es[: idxs[0]]
    frames: List[bytes] = []
    for k, start in enumerate(idxs):
        end = idxs[k + 1] if k + 1 < len(idxs) else len(es)
        frames.append(es[start:end])

    # Put VOL/headers in the first frame packet
    frames[0] = header + frames[0]
    return frames


# -----------------------------
# OGM (OggDS) stream headers and packets
# Spec layout referenced in ogmstreams.h :contentReference[oaicite:1]{index=1}
# -----------------------------
def build_ogm_stream_header_video(width: int, height: int, fps: int, fourcc: bytes = b"XVID") -> bytes:
    # stream_header structure (new header format) + video union fields.
    # Layout (little-endian), size=52 bytes (see ogmstreams.h) :contentReference[oaicite:2]{index=2}
    streamtype = (b"video" + b"\x00" * 3)  # 8 bytes
    subtype = fourcc[:4].ljust(4, b"\x00")

    size = 52
    time_unit = int(round(10_000_000 / fps))  # "reference time" units
    samples_per_unit = 1
    default_len = 1
    buffersize = 0
    bits_per_sample = 0
    padding = 0

    sh_video = struct.pack("<ii", width, height)

    base = struct.pack(
        "<8s4siqqiih",
        streamtype,
        subtype,
        size,
        time_unit,
        samples_per_unit,
        default_len,
        buffersize,
        bits_per_sample,
    )
    return base + sh_video + struct.pack("<h", padding)


def build_vorbis_comment(vendor: bytes = b"dvbang") -> bytes:
    # Vorbis-style comment body: vendor_len + vendor + user_comment_list_len
    return struct.pack("<I", len(vendor)) + vendor + struct.pack("<I", 0)


def ogm_video_header_packet(width: int, height: int, fps: int) -> bytes:
    # First packet: 0x01 + stream_header :contentReference[oaicite:3]{index=3}
    return b"\x01" + build_ogm_stream_header_video(width, height, fps)


def ogm_video_comment_packet() -> bytes:
    # Second packet: 0x03 + vorbis comment payload :contentReference[oaicite:4]{index=4}
    return b"\x03" + build_vorbis_comment()


def ogm_video_data_packet(frame_bytes: bytes, samples: int = 1, keyframe: bool = True, lenbytes: int = 2) -> bytes:
    """
    Data packets:
      - byte0: flags, includes lenbytes encoding bits and keyframe bit :contentReference[oaicite:5]{index=5}
      - next lenbytes: packet duration in samples (frames for video)
      - then payload
    """
    lb = lenbytes
    if not (0 <= lb <= 7):
        raise ValueError("lenbytes must be 0..7")

    bit0 = (lb & 0x1)
    bit1 = (lb & 0x2) >> 1
    bit2 = (lb & 0x4) >> 2

    flags = 0
    if keyframe:
        flags |= 0x08
    # lenbytes mapping per ogmstreams.h bit description :contentReference[oaicite:6]{index=6}
    flags |= (bit0 << 6) | (bit1 << 7) | (bit2 << 1)

    return bytes([flags]) + samples.to_bytes(lb, "little") + frame_bytes


# -----------------------------
# Build multiplexed OGM (.ogv) without ffmpeg muxing
# -----------------------------
def mux_to_ogv(
    out_path: Path,
    *,
    width: int,
    height: int,
    fps: int,
    frames: List[bytes],
    audio_pages: List[Dict],
    audio_sample_rate: int,
) -> None:
    if not audio_pages:
        raise ValueError("No audio Ogg pages")

    audio_serial = audio_pages[0]["serial"]
    video_serial = random.randint(1, 0xFFFFFFFF)
    while video_serial == audio_serial:
        video_serial = random.randint(1, 0xFFFFFFFF)

    # Find where Vorbis headers end (first 3 packets)
    hdr_end = vorbis_header_end_page_index(audio_pages)

    # Required by many demuxers: all BOS pages first.
    # We output:
    #   - audio BOS (copied)
    #   - video BOS (our header packet)
    # then the rest of audio header pages, then video comment, then interleaved data.
    v_header = ogm_video_header_packet(width, height, fps)
    v_comment = ogm_video_comment_packet()

    with out_path.open("wb") as f:
        # 1) BOS pages
        f.write(audio_pages[0]["bytes"])  # audio BOS page (unchanged)
        f.write(build_ogg_page([v_header], video_serial, 0, 0, 0x02))  # video BOS

        vseq = 1

        # 2) video comment header page
        f.write(build_ogg_page([v_comment], video_serial, vseq, 0, 0x00))
        vseq += 1

        # 3) remaining audio header pages (excluding BOS page already written)
        for p in audio_pages[1 : hdr_end + 1]:
            f.write(p["bytes"])

        # 4) interleave data pages by time
        v_i = 0

        for ap in audio_pages[hdr_end + 1 :]:
            at = (ap["granule"] / audio_sample_rate) if audio_sample_rate else 0.0

            # Emit video frames up to this audio timestamp (with a small half-frame slack)
            while v_i < len(frames) and ((v_i + 1) / fps) <= at + (0.5 / fps):
                pkt = ogm_video_data_packet(frames[v_i], samples=1, keyframe=True, lenbytes=2)
                gran = v_i + 1  # granulepos = frames decoded so far
                header_type = 0x04 if v_i == (len(frames) - 1) else 0x00
                f.write(build_ogg_page([pkt], video_serial, vseq, gran, header_type))
                vseq += 1
                v_i += 1

            # Then copy this audio page unchanged
            f.write(ap["bytes"])

        # If audio ended early, flush remaining video
        while v_i < len(frames):
            pkt = ogm_video_data_packet(frames[v_i], samples=1, keyframe=True, lenbytes=2)
            gran = v_i + 1
            header_type = 0x04 if v_i == (len(frames) - 1) else 0x00
            f.write(build_ogg_page([pkt], video_serial, vseq, gran, header_type))
            vseq += 1
            v_i += 1


# -----------------------------
# Main
# -----------------------------
def main() -> int:
    ap = argparse.ArgumentParser(description="Convert DV!! to .ogv (OGM/OggDS) in pure Python (no ffmpeg muxing).")
    ap.add_argument("input", type=Path, help="Input DV!! file (e.g. kleo19.dv)")
    ap.add_argument("--out", type=Path, default=None, help="Output .ogv path (default: <input>.ogv)")
    ap.add_argument("--fps", type=int, default=25, help="FPS (kleo19.dv: 25)")
    ap.add_argument("--validate", action="store_true", help="Do ffmpeg decode validation step")
    args = ap.parse_args()

    dv_path: Path = args.input
    if not dv_path.exists():
        eprint(f"Not found: {dv_path}")
        return 2

    out_path = args.out or dv_path.with_suffix(".ogv")

    hdr, video_es, audio_ogg = parse_dvbang(dv_path.read_bytes())
    width, height = hdr[4], hdr[5]

    # Parse audio ogg pages and get sample rate
    audio_pages = list(ogg_iter_pages(audio_ogg))
    sr = vorbis_sample_rate_from_ogg_pages(audio_pages)
    if sr <= 0:
        eprint("WARNING: could not detect Vorbis sample rate; interleaving may be less accurate.")
        sr = 44100  # safe fallback for typical files

    # Split MPEG-4 Part2 elementary stream into frames
    frames = split_mpeg4_vop_frames(video_es)

    mux_to_ogv(
        out_path,
        width=width,
        height=height,
        fps=args.fps,
        frames=frames,
        audio_pages=audio_pages,
        audio_sample_rate=sr,
    )

    print(f"ogv: {out_path}")
    print(f"video: MPEG-4 Part 2 (frames={len(frames)} {width}x{height} fps={args.fps})")
    print(f"audio: Vorbis (sr={sr} pages={len(audio_pages)})")

    # Validation (decode test only; not muxing)
    if args.validate:
        rc = run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-i", str(out_path), "-f", "null", "-"])
        if rc != 0:
            eprint("WARNING: ffmpeg decode test reported errors (file was still written).")

    return 0


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