dv_timecode_subtitle.py

#!/usr/bin/env python3
"""Generate an ASS subtitle that overlays the real recording date/time
from a tape-captured DV file, using dvrescue for analysis.

Usage:
    python dv_timecode_subtitle.py /path/to/capture.mov
    python dv_timecode_subtitle.py capture.dv -o /tmp/out.ass
"""

from __future__ import annotations

import argparse
import shutil
import subprocess
import sys
import tempfile
import xml.etree.ElementTree as ET
from dataclasses import dataclass
from datetime import datetime, timedelta
from pathlib import Path
from typing import Iterable, Iterator

DVRESCUE_NS = "https://mediaarea.net/dvrescue"


@dataclass
class FrameMeta:
    n: int
    rdt: datetime


def _find(elem, tag):
    found = elem.find(f"{{{DVRESCUE_NS}}}{tag}")
    if found is None:
        found = elem.find(tag)
    return found


def _findall(elem, tag):
    res = elem.findall(f"{{{DVRESCUE_NS}}}{tag}")
    if not res:
        res = elem.findall(tag)
    return res


def parse_rdt(rdt: str) -> datetime:
    for fmt in ("%Y-%m-%d %H:%M:%S.%f", "%Y-%m-%d %H:%M:%S"):
        try:
            return datetime.strptime(rdt, fmt)
        except ValueError:
            pass
    raise ValueError(f"unrecognized rdt format: {rdt!r}")


def parse_rate(rate: str) -> float:
    if "/" in rate:
        num, den = rate.split("/", 1)
        return float(num) / float(den)
    return float(rate)


def run_dvrescue(input_path: Path, xml_path: Path) -> None:
    if shutil.which("dvrescue") is None:
        sys.exit(
            "Error: dvrescue not found in PATH. Install from https://mediaarea.net/dvrescue")
    try:
        subprocess.run(
            ["dvrescue", str(input_path), "-x", str(xml_path)],
            check=True,
        )
    except subprocess.CalledProcessError as e:
        sys.exit(f"Error: dvrescue exited with code {e.returncode}")


def parse_dvrescue_xml(xml_path: Path) -> tuple[float, int, list[FrameMeta]]:
    tree = ET.parse(xml_path)
    root = tree.getroot()
    media = _find(root, "media")
    if media is None:
        sys.exit("Error: dvrescue XML has no <media> element.")
    frames = _find(media, "frames")
    if frames is None:
        sys.exit("Error: dvrescue XML has no <frames> element.")

    rate_str = frames.get("video_rate") or media.get("video_rate")
    if not rate_str:
        sys.exit("Error: video_rate not found in dvrescue XML.")
    rate = parse_rate(rate_str)

    total = int(frames.get("count") or 0)

    metas: list[FrameMeta] = []
    for f in _findall(frames, "frame"):
        rdt = f.get("rdt")
        n = f.get("n")
        if not rdt or n is None:
            continue
        try:
            metas.append(FrameMeta(int(n), parse_rdt(rdt)))
        except ValueError:
            continue

    if not metas:
        sys.exit(
            "Error: no <frame rdt=...> entries found — this DV has no recorded date/time.")

    metas.sort(key=lambda m: m.n)
    if total <= metas[-1].n:
        total = metas[-1].n + 1
    return rate, total, metas


def build_segments(
    rate: float, total: int, metas: list[FrameMeta]
) -> Iterator[tuple[int, int, datetime]]:
    """Yield (start_frame, end_frame_exclusive, displayed_datetime) where the
    displayed datetime (truncated to whole seconds) stays constant.

    dvrescue only emits a <frame> on changes / discontinuities, so we treat
    each emitted entry as an anchor and assume rdt advances at one wall-clock
    second per video second within the span until the next anchor.
    """
    # Drop any anchors past the declared total.
    anchors = [m for m in metas if m.n < total]
    if not anchors:
        return

    cur_start: int | None = None
    cur_label: datetime | None = None

    for i, anchor in enumerate(anchors):
        end_n = anchors[i + 1].n if i + 1 < len(anchors) else total
        # Step through frames in this span; small enough for typical DV files.
        for n in range(anchor.n, end_n):
            dt = anchor.rdt + timedelta(seconds=(n - anchor.n) / rate)
            label = dt.replace(microsecond=0)
            if label != cur_label:
                if cur_label is not None and cur_start is not None:
                    yield cur_start, n, cur_label
                cur_start = n
                cur_label = label

    if cur_label is not None and cur_start is not None:
        yield cur_start, total, cur_label


def frame_to_ass_time(n: int, rate: float) -> str:
    t = n / rate
    if t < 0:
        t = 0.0
    h = int(t // 3600)
    m = int((t % 3600) // 60)
    s = t - h * 3600 - m * 60
    return f"{h:d}:{m:02d}:{s:05.2f}"


ASS_HEADER_TPL = """[Script Info]
ScriptType: v4.00+
PlayResX: {res_x}
PlayResY: {res_y}
ScaledBorderAndShadow: yes
YCbCr Matrix: TV.601

[V4+ Styles]
Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding
Style: Default,{font},{font_size},&H00FFFFFF,&H000000FF,&HFF000000,&HFF000000,1,0,0,0,100,100,0,0,0,0,1,3,20,{margin_r},{margin_v},1

[Events]
Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text
"""


def write_ass(
    out_path: Path,
    segments: Iterable[tuple[int, int, datetime]],
    rate: float,
    res_x: int,
    res_y: int,
    font: str,
    font_size: int,
    margin_r: int,
    margin_v: int,
    time_format: str,
) -> int:
    count = 0
    with out_path.open("w", encoding="utf-8") as fp:
        fp.write(ASS_HEADER_TPL.format(
            res_x=res_x, res_y=res_y,
            font=font, font_size=font_size,
            margin_r=margin_r, margin_v=margin_v,
        ))
        for start_n, end_n, dt in segments:
            t0 = frame_to_ass_time(start_n, rate)
            t1 = frame_to_ass_time(end_n, rate)
            text = dt.strftime(time_format)
            fp.write(f"Dialogue: 0,{t0},{t1},Default,,0,0,0,,{text}\n")
            count += 1
    return count


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Generate an ASS subtitle with the real recording time of a DV capture, "
                    "via dvrescue. The subtitle is anchored bottom-right.",
    )
    parser.add_argument("input", type=Path,
                        help="DV capture file (.dv, .mov, .avi ...)")
    parser.add_argument(
        "-o", "--output", type=Path,
        help="Output .ass path. Defaults to the input file with the suffix replaced by .ass, "
             "saved in the same directory as the video.",
    )
    parser.add_argument(
        "--keep-xml", type=Path,
        help="Also write the intermediate dvrescue XML to this path "
             "(useful for debugging).",
    )
    parser.add_argument(
        "--xml", type=Path,
        help="Skip running dvrescue and parse this existing dvrescue XML instead.",
    )
    parser.add_argument(
        "--res", default="720x576",
        help="ASS PlayRes WxH. Default 720x576 (PAL DV). Use 720x480 for NTSC.",
    )
    parser.add_argument("--font", default="Arial",
                        help="Subtitle font (default: Arial).")
    parser.add_argument("--font-size", type=int, default=18,
                        help="Subtitle font size in PlayRes units (default: 28).")
    parser.add_argument("--margin-r", type=int, default=20,
                        help="Right margin in PlayRes units (default: 20).")
    parser.add_argument("--margin-v", type=int, default=20,
                        help="Bottom margin in PlayRes units (default: 20).")
    parser.add_argument(
        "--time-format", default="%Y-%m-%d %H:%M:%S",
        help="strftime format for the displayed timestamp (default: '%%Y-%%m-%%d %%H:%%M:%%S').",
    )
    args = parser.parse_args()

    if not args.input.exists():
        sys.exit(f"Error: input file not found: {args.input}")

    out_path = args.output or args.input.with_suffix(".ass")
    out_path.parent.mkdir(parents=True, exist_ok=True)

    try:
        res_x_str, res_y_str = args.res.lower().split("x", 1)
        res_x, res_y = int(res_x_str), int(res_y_str)
    except ValueError:
        sys.exit(f"Error: --res must look like WIDTHxHEIGHT, got {args.res!r}")

    if args.xml:
        if not args.xml.exists():
            sys.exit(f"Error: --xml file not found: {args.xml}")
        rate, total, metas = parse_dvrescue_xml(args.xml)
    else:
        with tempfile.TemporaryDirectory() as tmpdir:
            xml_path = args.keep_xml or Path(tmpdir) / "dvrescue.xml"
            if args.keep_xml:
                xml_path.parent.mkdir(parents=True, exist_ok=True)
            run_dvrescue(args.input, xml_path)
            rate, total, metas = parse_dvrescue_xml(xml_path)

    segments = list(build_segments(rate, total, metas))
    count = write_ass(
        out_path, segments, rate,
        res_x=res_x, res_y=res_y,
        font=args.font, font_size=args.font_size,
        margin_r=args.margin_r, margin_v=args.margin_v,
        time_format=args.time_format,
    )

    duration = total / rate if rate else 0.0
    print(
        f"frame rate: {rate:.3f} fps | total frames: {total} | duration: {duration:.1f}s\n"
        f"anchors with rdt: {len(metas)} | subtitle events: {count}\n"
        f"wrote: {out_path}"
    )


if __name__ == "__main__":
    main()