movie_sync/extract.py

#!/usr/bin/env python3
"""
Extracts the first 1.5 minutes of a video into:
  - a video-only file (no audio)
  - an audio-only file (no video)
Output goes into an 'extracted' folder next to this script.
Requires: ffmpeg installed and on PATH.
"""

import subprocess
import os
import sys
import shutil
import json

# ─── Input / config (config.toml is base config, fallback to hardcoded) ───
CONFIG_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.toml")
_cfg_inputs = None
_cfg_duration = 90
_cfg_start = 0.0
_cfg_out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "extracted")
_cfg_target_fps = "24000/1001"
_cfg_eps = 0.02
_cfg_normalize = True
_cfg_target_lufs = -14
_cfg_target_tp = -1.0
_cfg_target_lra = 11
_cfg_audio_stream = None  # None = prompt, int = index, "skip" = no audio extraction
try:
    try:
        import tomllib  # py3.11+
    except ModuleNotFoundError:
        import tomli as tomllib  # type: ignore
    if os.path.exists(CONFIG_PATH):
        with open(CONFIG_PATH, "rb") as _f:
            _cfg = tomllib.load(_f)
        _ex = _cfg.get("extract", {})
        if _ex.get("file_a") or _ex.get("file_b"):
            _cfg_inputs = [p for p in [_ex.get("file_a"), _ex.get("file_b")] if p]
        _cfg_duration = int(_ex.get("duration", _cfg_duration))
        _cfg_start = float(_ex.get("start", _cfg_start))
        _cfg_out_dir = _ex.get("out_dir", _cfg_out_dir)
        if _ex.get("target_fps"):
            _cfg_target_fps = str(_ex["target_fps"])
        if _ex.get("target_fps_eps") is not None:
            try:
                _cfg_eps = float(_ex["target_fps_eps"])
            except Exception:
                pass
        _au = _cfg.get("audio", {})
        if _au.get("normalize") is not None:
            _cfg_normalize = bool(_au["normalize"])
        if _au.get("target_lufs") is not None:
            try:
                _cfg_target_lufs = float(_au["target_lufs"])
            except Exception:
                pass
        if _au.get("target_tp") is not None:
            try:
                _cfg_target_tp = float(_au["target_tp"])
            except Exception:
                pass
        if _au.get("target_lra") is not None:
            try:
                _cfg_target_lra = float(_au["target_lra"])
            except Exception:
                pass
        if _au.get("stream") is not None:
            val = _au["stream"]
            if isinstance(val, str) and val == "skip":
                _cfg_audio_stream = "skip"
            elif isinstance(val, (int, float)):
                _cfg_audio_stream = int(val)
        if _cfg_inputs:
            print(f"[config] loaded {CONFIG_PATH}: duration={_cfg_duration}s start={_cfg_start}s target_fps={_cfg_target_fps}")
except Exception as _e:
    print(f"[config] warn: failed to load {CONFIG_PATH}: {_e}", file=sys.stderr)

if _cfg_inputs is not None:
    inputs = _cfg_inputs
else:
    inputs = [
        "/path/to/movie_b.mp4",
        "/path/to/movie_a.mkv"
    ]
INPUT_FILES = []
for input_file in inputs:
    INPUT_FILES.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), input_file))
DURATION   = _cfg_duration
START_SEC  = _cfg_start  # for ffmpeg -ss (0 = file start)
OUT_DIR    = os.path.join(os.path.dirname(os.path.abspath(__file__)), _cfg_out_dir) if not os.path.isabs(_cfg_out_dir) else _cfg_out_dir

# ─── Analysis-friendly codecs (NOT GPU) ───
# Video: FFV1 (lossless, intra-only) in MKV. Every frame stands alone, no
#   P/B-frame prediction, no blocking artifacts -> exact frames for similarity
#   (SSIM/pHash/features) + fast random access. Decoded with -pix_fmt yuv444p
#   to avoid 4:2:0 chroma-subsampling noise.
# Audio: PCM float32 WAV, 48 kHz mono. Bit-exact waveform preserves phase for
#   cross-correlation / GCC-PHAT. Never MP3/AAC/Opus here (MDCT + padding
#   smears transients and shifts correlation peaks). FLAC is the only
#   acceptable compressed alternative (lossless).
# GPU (RX 7900 XTX via VAAPI/AMF: h264_vaapi/hevc_vaapi/h264_amf): deliberately
#   NOT used. Those encoders are lossy delivery codecs, lower quality per bit,
#   inter-frame only, and FFV1/PCM/ProRes/DNxHR have no GPU implementation.
#   GPU only makes sense for preview proxies, not for measurement. CPU encode
#   of 90 s clips is fast enough and bit-exact.
VIDEO_CODEC_BASE = ["-c:v", "ffv1", "-level", "3", "-pix_fmt", "yuv444p"]
AUDIO_CODEC = ["-c:a", "pcm_f32le", "-ar", "48000", "-ac", "1"]
# target fps may be overridden via config.toml
TARGET_FPS = _cfg_target_fps  # default "24000/1001"
try:
    _n, _d = TARGET_FPS.split("/") if "/" in TARGET_FPS else (TARGET_FPS, "1")
    TARGET_FPS_FLOAT = float(_n) / float(_d)
except Exception:
    TARGET_FPS_FLOAT = 24000 / 1001
FPS_EPS = _cfg_eps  # tolerance for float comparison
NORMALIZE = _cfg_normalize
TARGET_LUFS = _cfg_target_lufs
TARGET_TP = _cfg_target_tp
TARGET_LRA = _cfg_target_lra

def _parse_rate(s: str) -> float | None:
    s = s.strip()
    if not s or s in ("0/0", "0"):
        return None
    if "/" in s:
        try:
            n, d = s.split("/")
            return float(n) / float(d) if float(d) != 0 else None
        except ValueError:
            return None
    try:
        return float(s)
    except ValueError:
        return None

def probe_fps(path: str) -> float | None:
    """Return probed video fps via ffprobe, or None if probe fails."""
    if shutil.which("ffprobe") is None:
        return None
    # avg_frame_rate is usually canonical; fall back to r_frame_rate
    for entry in ("avg_frame_rate", "r_frame_rate"):
        try:
            r = subprocess.run(
                ["ffprobe", "-v", "error", "-select_streams", "v:0",
                 "-show_entries", f"stream={entry}", "-of", "default=nw=1:nk=1", path],
                capture_output=True, text=True, timeout=10
            )
            for line in r.stdout.strip().splitlines():
                fps = _parse_rate(line)
                if fps and fps > 0 and fps < 1000:
                    return fps
        except Exception:
            continue
    return None

def fps_matches(a: float | None, b: float, eps: float = FPS_EPS) -> bool:
    return a is not None and abs(a - b) < eps

def probe_audio_streams(path: str) -> list[dict]:
    """Probe all audio streams in a file. Returns list of dicts matching ~/.bin/ff style."""
    if shutil.which("ffprobe") is None:
        return []
    try:
        r = subprocess.run(
            ["ffprobe", "-v", "error", "-select_streams", "a",
             "-show_entries",
             "stream=index,codec_name,channels,sample_rate,bit_rate,bits_per_sample,duration,"
             "stream_disposition=forced,default:"
             "stream_tags=language,title",
             "-of", "json", path],
            capture_output=True, text=True, timeout=10
        )
        if r.returncode != 0:
            return []
        data = json.loads(r.stdout)
        streams = []
        for s in data.get("streams", []):
            tags = {k.lower(): v for k, v in (s.get("tags") or {}).items()}
            dispo = s.get("disposition", {})
            sr = s.get("sample_rate")
            br = s.get("bit_rate")
            depth = s.get("bits_per_sample") or s.get("bits_per_raw_sample")
            dur = s.get("duration")
            streams.append({
                "index": s.get("index", 0),
                "codec": s.get("codec_name", ""),
                "channels": s.get("channels", 0),
                "sample_rate": int(sr) if sr else 0,
                "sample_rate_khz": f"{int(sr)/1000:.1f} kHz" if sr else "",
                "bit_rate": f"{int(br)//1000} kbps" if br else "",
                "bit_depth": f"{depth}-bit" if depth else "",
                "duration": dur,
                "language": _normalize_audio_language(tags.get("language", "und")),
                "title": tags.get("title", ""),
                "is_default": dispo.get("default", 0) == 1,
                "is_forced": dispo.get("forced", 0) == 1,
            })
        return streams
    except Exception:
        return []

def _normalize_audio_language(raw: str) -> str:
    """Map common ISO-639 codes to readable names."""
    lang_map = {
        "eng": "English", "ger": "German", "deu": "German",
        "fre": "French", "fra": "French", "spa": "Spanish",
        "ita": "Italian", "jpn": "Japanese", "chi": "Chinese",
        "zho": "Chinese", "kor": "Korean", "por": "Portuguese",
        "rus": "Russian", "ara": "Arabic", "hin": "Hindi",
        "dut": "Dutch", "nld": "Dutch", "swe": "Swedish",
        "nor": "Norwegian", "dan": "Danish", "fin": "Finnish",
        "pol": "Polish", "tur": "Turkish", "tha": "Thai",
        "vie": "Vietnamese", "und": "Undetermined", "mis": "Misnamed",
        "inl": "Latin American Spanish", "qaa": "Reserved",
    }
    raw_lower = raw.lower().strip()
    return lang_map.get(raw_lower, raw_lower.upper())

def select_audio_stream(path: str, cfg_hint: int | str | None = None) -> int | None:
    """
    Probe audio streams and let user pick one.
    Returns 0-based audio stream index for ffmpeg -map 0:a:N, or None (no audio).
    cfg_hint: int = use that index directly, "skip" = None, None = prompt.
    """
    streams = probe_audio_streams(path)

    if not streams:
        print(f"[audio] {os.path.basename(path)}: no audio streams found")
        return None

    if len(streams) == 1:
        s = streams[0]
        label = _format_stream_label(s)
        print(f"[audio] {os.path.basename(path)}: only 1 stream → {label}")
        return 0  # -map 0:a:0

    # Multiple streams — show menu
    print(f"\n{'=' * 60}")
    print(f"Audio streams in: {os.path.basename(path)}")
    print(f"{'=' * 60}")
    for i, s in enumerate(streams):
        label = _format_stream_label(s)
        marker = " *" if s.get("is_default") else ""
        print(f"  [{i}] {label}{marker}")
    print(f"  [s] skip audio extraction")
    print(f"{'=' * 60}")

    while True:
        choice = input("Select audio stream [0]: ").strip().lower()
        if choice == "" or choice == "0":
            return 0
        if choice == "s":
            return None
        try:
            idx = int(choice)
            if 0 <= idx < len(streams):
                return idx
        except ValueError:
            pass
        print(f"  Invalid choice, enter 0-{len(streams)-1} or 's'")

def _format_stream_label(s: dict) -> str:
    parts = []
    # Language + title first (most important)
    lang = s.get("language", "")
    title = s.get("title", "")
    if lang:
        parts.append(lang)
    if title:
        parts.append(f'"{title}"')
    # Codec
    codec = s.get("codec", "")
    if codec:
        parts.append(codec.upper())
    # Specs: channels, sample rate, bit depth
    specs = []
    ch = s.get("channels", 0)
    if ch:
        specs.append(f"{ch}ch")
    if s.get("sample_rate_khz"):
        specs.append(s["sample_rate_khz"])
    if s.get("bit_depth"):
        specs.append(s["bit_depth"])
    if specs:
        parts.append(f"({', '.join(specs)})")
    # Bitrate
    if s.get("bit_rate"):
        parts.append(f"@{s['bit_rate']}")
    # Flags
    flags = []
    if s.get("is_default"):
        flags.append("Default")
    if s.get("is_forced"):
        flags.append("Forced")
    if flags:
        parts.append(f"[{'/'.join(flags)}]")
    return " ".join(parts)

def _parse_loudnorm_json(stderr: str) -> dict | None:
    """Extract JSON block from loudnorm stderr output."""
    try:
        start = stderr.rindex("{")
        end = stderr.rindex("}") + 1
        return json.loads(stderr[start:end])
    except (ValueError, json.JSONDecodeError):
        return None

def extract_audio_loudnorm(input_file: str, duration: int, start_sec: float,
                           out_path: str, stem: str, stream_idx: int | None = None) -> None:
    """Two-pass EBU R128 loudnorm extraction: measure then apply."""
    ss_args = ["-ss", str(start_sec)] if start_sec and start_sec != 0 else []
    map_args = ["-map", f"0:a:{stream_idx}"] if stream_idx is not None else []
    base = ["ffmpeg", "-y", *ss_args, "-i", input_file, "-t", str(duration), "-vn", *map_args]
    measure_af = f"loudnorm=I={TARGET_LUFS}:TP={TARGET_TP}:LRA={TARGET_LRA}:print_format=json"

    # ─── Pass 1: measure ───
    print(f"[loudness] {stem}: pass 1/2 measuring…")
    cmd_measure = base + ["-af", measure_af, "-f", "null", "-"]
    r = subprocess.run(cmd_measure, capture_output=True, text=True)
    if r.returncode != 0:
        print(f"[loudness] {stem}: pass 1 failed ({r.stderr.strip()[-80:]}), falling back to simple extract", file=sys.stderr)
        cmd_fallback = base + [*AUDIO_CODEC, out_path]
        r2 = subprocess.run(cmd_fallback, capture_output=True, text=True)
        if r2.returncode != 0:
            sys.exit(f"ffmpeg (audio) failed:\n{r2.stderr}")
        return

    measured = _parse_loudnorm_json(r.stderr)
    if measured is None:
        print(f"[loudness] {stem}: could not parse measurement, falling back to simple extract", file=sys.stderr)
        cmd_fallback = base + [*AUDIO_CODEC, out_path]
        r2 = subprocess.run(cmd_fallback, capture_output=True, text=True)
        if r2.returncode != 0:
            sys.exit(f"ffmpeg (audio) failed:\n{r2.stderr}")
        return

    mi = measured.get("input_i", "-24")
    mtp = measured.get("input_tp", "-2")
    mlra = measured.get("input_lra", "7")
    mthresh = measured.get("input_thresh", "-34")
    offset = measured.get("target_offset", "0")
    print(f"[loudness] {stem}: measured I={mi} LUFS, TP={mtp} dBTP, LRA={mlra}, offset={offset} dB")

    # ─── Pass 2: apply with measured values ───
    print(f"[loudness] {stem}: pass 2/2 applying loudnorm → {TARGET_LUFS} LUFS…")
    apply_af = (f"loudnorm=I={TARGET_LUFS}:TP={TARGET_TP}:LRA={TARGET_LRA}"
                f":measured_I={mi}:measured_TP={mtp}:measured_LRA={mlra}"
                f":measured_thresh={mthresh}:offset={offset}:linear=true")
    cmd_apply = base + ["-af", apply_af, *AUDIO_CODEC, out_path]
    r = subprocess.run(cmd_apply, capture_output=True, text=True)
    if r.returncode != 0:
        sys.exit(f"ffmpeg (audio) pass 2 failed:\n{r.stderr}")
    print(f"[loudness] {stem}: normalized to {TARGET_LUFS} LUFS ✔")

# ─── Check ffmpeg ───
if shutil.which(f"ffmpeg") is None:
    sys.exit(f"ERROR: 'ffmpeg' not found on PATH. Install it and try again.")

# ─── Pre-probe all inputs to decide fps handling ───
probed = {}
for p in INPUT_FILES:
    fps = probe_fps(p) if os.path.exists(p) else None
    probed[p] = fps
    tag = f"{fps:.3f}" if fps else "unknown (ffprobe missing or file not found)"
    print(f"[probe] {os.path.basename(p)}: {tag} fps" + (f" {'≈ target ' + TARGET_FPS if fps_matches(fps, TARGET_FPS_FLOAT) else '≠ target ' + TARGET_FPS}" if fps else ""))

# sanity check: both mismatch target and differ from each other → warn (neither is canonical)
valid_fps = [v for v in probed.values() if v is not None]
if len(valid_fps) >= 2:
    all_match = all(fps_matches(v, TARGET_FPS_FLOAT) for v in valid_fps)
    any_mismatch = any(not fps_matches(v, TARGET_FPS_FLOAT) for v in valid_fps)
    if any_mismatch:
        distinct = len({round(v, 2) for v in valid_fps}) > 1
        if not all_match and distinct:
            print(f"WARNING: inputs have different non-target fps {valid_fps} — both will be normalized to {TARGET_FPS} ({TARGET_FPS_FLOAT:.6f} fps). "
                  f"If this is PAL speedup (25) vs NTSC film (24000/1001), consider also compensating audio tempo.")
        elif not all_match:
            print(f"NOTE: inputs share fps {valid_fps[0]:.3f} but differ from target {TARGET_FPS} — both will be normalized to {TARGET_FPS}.")

# ─── Derive output stems from path divergence ───
def _derive_stems(paths: list[str]) -> list[str]:
    """For each path, give a stem: if filenames differ, just use filename stem.
    If filenames match, walk up dirs until parents diverge and use the remaining relative path."""
    stems = []
    # first strip extensions, get bare names
    basenames = [os.path.splitext(os.path.basename(p))[0] for p in paths]
    if len(paths) == 1 or basenames[0] != basenames[1]:
        # different filenames — just use them directly
        return basenames
    # same filename — walk up until parent dirs differ
    parts_a = os.path.normpath(paths[0]).split(os.sep)
    parts_b = os.path.normpath(paths[1]).split(os.sep)
    # find divergence point from the end (walk up from file)
    common = 0
    for a, b in zip(reversed(parts_a), reversed(parts_b)):
        if a == b:
            common += 1
        else:
            break
    # remaining parts after divergence = relative unique path (dirs + filename)
    rel_a = parts_a[len(parts_a) - common:]   # includes filename
    rel_b = parts_b[len(parts_b) - common:]
    stem_a = os.path.join(*rel_a) if len(rel_a) > 1 else rel_a[0]
    stem_b = os.path.join(*rel_b) if len(rel_b) > 1 else rel_b[0]
    return [stem_a, stem_b]

_stems = _derive_stems(INPUT_FILES)

# ─── Prepare output folder ───
os.makedirs(OUT_DIR, exist_ok=True)

_audio_stream_selections = []  # track selected stream indices for config

for _file_idx, INPUT_FILE in enumerate(INPUT_FILES):

    # ─── Derive output filenames ───
    stem = _stems[_file_idx]
    video_out = os.path.join(OUT_DIR, f"{stem}_video.mkv")  # FFV1
    audio_out = os.path.join(OUT_DIR, f"{stem}_audio.wav")  # PCM f32

    print(f"Input : {INPUT_FILE}")
    print(f"Video → {video_out}")
    print(f"Audio → {audio_out}")
    print(f"-" * 50)

    # ─── Probe and select audio stream ───
    audio_stream = None
    if not (_cfg_audio_stream == "skip"):
        audio_stream = select_audio_stream(INPUT_FILE, cfg_hint=_cfg_audio_stream)
        if audio_stream is not None:
            print(f"[audio] {stem}: selected stream {audio_stream}")
        else:
            print(f"[audio] {stem}: no audio selected, skipping audio extraction")
    _audio_stream_selections.append(audio_stream)

    # ─── Extract video-only (duration s, no audio, CPU FFV1 re-encode) ───
    # Only apply fps filter if probed fps differs from target (avoid needless re-timing)
    fps = probed.get(INPUT_FILE)
    needs_fps = not fps_matches(fps, TARGET_FPS_FLOAT)
    # fps filter alone forces CFR; use fps_mode cfr (vsync removed in ffmpeg >=7)
    video_codec = VIDEO_CODEC_BASE + (["-vf", f"fps={TARGET_FPS}:round=near", "-fps_mode", "cfr"] if needs_fps else [])
    if fps is not None:
        print(f"[fps] {stem}: probed {fps:.3f} → {'convert to ' + TARGET_FPS if needs_fps else 'already ' + TARGET_FPS + ', keeping native timing'}")
    else:
        print(f"[fps] {stem}: unknown fps → forcing {TARGET_FPS} (safe fallback)")
        # ensure fallback forces target if we couldn't probe
        if not needs_fps:
            video_codec = VIDEO_CODEC_BASE + ["-vf", f"fps={TARGET_FPS}:round=near", "-fps_mode", "cfr"]

    ss_args = ["-ss", str(START_SEC)] if START_SEC and START_SEC != 0 else []
    cmd_video = [
        "ffmpeg",
        "-y",               # overwrite output
        *ss_args,           # seek before -i (fast, frame-accurate with re-encode)
        "-i", INPUT_FILE,
        "-t", str(DURATION),
        "-an",              # no audio
        *video_codec,       # ffv1 lossless intra-only, CPU (no GPU equiv.)
        video_out
    ]
    print(f"[video: {stem}] running ffmpeg…")
    r = subprocess.run(cmd_video, capture_output=True, text=True)
    if r.returncode != 0:
        sys.exit(f"ffmpeg (video) failed:\n{r.stderr}")
    print(f"[video: {stem}] done ✔")

    # ─── Extract audio-only (duration s, no video, CPU PCM re-encode) ───
    if audio_stream is None:
        print(f"[audio: {stem}] skipped (no stream selected)")
    elif NORMALIZE:
        print(f"[audio: {stem}] extracting stream {audio_stream} with two-pass loudnorm…")
        extract_audio_loudnorm(INPUT_FILE, DURATION, START_SEC, audio_out, stem, stream_idx=audio_stream)
    else:
        map_args = ["-map", f"0:a:{audio_stream}"]
        cmd_audio = [
            "ffmpeg",
            "-y",
            *ss_args,
            "-i", INPUT_FILE,
            "-t", str(DURATION),
            "-vn",              # no video
            *map_args,          # selected audio stream
            *AUDIO_CODEC,       # pcm_f32le 48kHz mono, CPU (no GPU equiv.)
            audio_out
        ]
        print(f"[audio: {stem}] running ffmpeg (no normalization)…")
        r = subprocess.run(cmd_audio, capture_output=True, text=True)
        if r.returncode != 0:
            sys.exit(f"ffmpeg (audio) failed:\n{r.stderr}")
        print(f"[audio: {stem}] done ✔")

    print(f"-" * 50)
print(f"All done. Files are in:", OUT_DIR)

# ─── Save selected audio stream indices to config.toml ───
if len(_audio_stream_selections) >= 2 and CONFIG_PATH and os.path.exists(CONFIG_PATH):
    try:
        with open(CONFIG_PATH, "r") as f:
            cfg_text = f.read()
        # Update or add audio_stream_a and audio_stream_b in [extract] section
        import re
        def _cfg_set(key: str, val) -> str:
            nonlocal cfg_text
            pat = re.compile(rf'(^{key}\s*=\s*).*$')
            replacement = f'{key} = {val}'
            if pat.search(cfg_text):
                cfg_text = pat.sub(replacement, cfg_text, count=1)
            else:
                # Insert before the next section header or at end of [extract]
                ins_pat = re.compile(r'(\[audio\])')
                if ins_pat.search(cfg_text):
                    cfg_text = ins_pat.sub(f'{replacement}\n\n\\1', cfg_text, count=1)
                else:
                    cfg_text += f"\n{replacement}\n"
        _cfg_set("audio_stream_a", _audio_stream_selections[0]
                 if _audio_stream_selections[0] is not None else "null")
        _cfg_set("audio_stream_b", _audio_stream_selections[1]
                 if _audio_stream_selections[1] is not None else "null")
        with open(CONFIG_PATH, "w") as f:
            f.write(cfg_text)
        print(f"[config] saved audio_stream_a={_audio_stream_selections[0]}, "
              f"audio_stream_b={_audio_stream_selections[1]} → {CONFIG_PATH}")
    except Exception as e:
        print(f"[config] warn: could not save stream indices: {e}", file=sys.stderr)