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

# ─── Hardcoded input (relative to this script's directory) ───
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   = 90  # 1.5 minutes in seconds
OUT_DIR    = os.path.join(os.path.dirname(os.path.abspath(__file__)), "extracted")

# ─── 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 = ["-c:v", "ffv1", "-level", "3", "-pix_fmt", "yuv444p"]
AUDIO_CODEC = ["-c:a", "pcm_f32le", "-ar", "48000", "-ac", "1"]
# ─── Check ffmpeg ───
if shutil.which(f"ffmpeg") is None:
    sys.exit(f"ERROR: 'ffmpeg' not found on PATH. Install it and try again.")

# ─── 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) ───
    cmd_video = [
        "ffmpeg",
        "-y",               # overwrite output
        "-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",
        "-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)