#!/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

# ─── 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
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
        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 = [
        "/home/honney/Videos/Drachenzähmen/Drachenzaehmen.mp4",
        "/home/honney/Videos/Drachenzähmen/How to Train Your Dragon.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

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

# ─── 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}.")

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

for INPUT_FILE in INPUT_FILES:

    # ─── Derive output filenames ───
    stem = os.path.splitext(os.path.basename(INPUT_FILE))[0]
    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)

    # ─── Extract video-only (first 90s, 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 (first 90s, no video, CPU PCM re-encode) ───
    cmd_audio = [
        "ffmpeg",
        "-y",
        *ss_args,
        "-i", INPUT_FILE,
        "-t", str(DURATION),
        "-vn",              # no video
        *AUDIO_CODEC,       # pcm_f32le 48kHz mono, CPU (no GPU equiv.)
        audio_out
    ]
    print(f"[audio: {stem}] running ffmpeg…")
    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)