#!/usr/bin/env python3
"""Batch-compress every video in a folder to HEVC using the "1080P HEVC NVENC"
HandBrake preset (hevc_nvenc_1080p.json, must sit next to this script).

Drop videos into the folder, run the script, collect the *_hevc.mp4 outputs:

    python compress_hevc.py                      # uses E:\\Media\\Compress-Temp
    python compress_hevc.py D:\\some\\folder     # any other folder
    python compress_hevc.py --target-gb 3.0      # change target size (default 2.5)
    python compress_hevc.py --dry-run            # show what would run, encode nothing

For each file the script probes the duration with ffprobe and derives the
average video bitrate needed to land near the target size, so a 1-hour clip
and a 5-hour movie both come out roughly the same size. Files that are
already at/below the target size are skipped, as are *_hevc.mp4 outputs and
files whose output already exists. Encodes run sequentially (NVENC has a
single hardware encoder, so parallelism buys nothing).
"""

import argparse
import shutil
import subprocess
import sys
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
PRESET_FILE = SCRIPT_DIR / "hevc_nvenc_1080p.json"
PRESET_NAME = "1080P HEVC NVENC"

DEFAULT_FOLDER = r"E:\Media\Compress-Temp"
VIDEO_EXTS = {".mp4", ".mkv", ".avi", ".mov", ".m4v", ".ts", ".wmv"}
OUTPUT_SUFFIX = "_hevc"

AUDIO_KBPS = 160          # audio bitrate set by the preset
OVERSHOOT = 1.09          # NVENC delivers ~9% above the requested avg bitrate
MIN_VIDEO_KBPS = 800      # below this, 1080p HEVC falls apart visibly
MAX_VIDEO_KBPS = 8000     # no point going higher for re-compression


def find_tool(name, fallback):
    path = shutil.which(name)
    if path:
        return path
    if Path(fallback).exists():
        return fallback
    sys.exit(f"error: '{name}' not found on PATH or at {fallback}")


def probe_duration(ffprobe, path):
    out = subprocess.run(
        [ffprobe, "-v", "error", "-show_entries", "format=duration",
         "-of", "default=noprint_wrappers=1:nokey=1", str(path)],
        capture_output=True, text=True,
    )
    if out.returncode != 0:
        raise RuntimeError(f"ffprobe failed for {path}: {out.stderr.strip()}")
    return float(out.stdout.strip())


def calc_video_bitrate(duration_s, target_bytes):
    """Requested (-b) video kbps to land near target_bytes, given NVENC overshoot."""
    total_kbps = target_bytes * 8 / 1000 / duration_s
    video_kbps = (total_kbps - AUDIO_KBPS) / OVERSHOOT
    return int(round(max(MIN_VIDEO_KBPS, min(MAX_VIDEO_KBPS, video_kbps))))


def fmt_gib(nbytes):
    return f"{nbytes / 1024**3:.2f} GiB"


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("folder", nargs="?", default=DEFAULT_FOLDER,
                        help=f"folder to scan (default: {DEFAULT_FOLDER})")
    parser.add_argument("--target-gb", type=float, default=2.5,
                        help="target size per video in GiB (default: 2.5)")
    parser.add_argument("--dry-run", action="store_true",
                        help="print the planned encodes without running them")
    args = parser.parse_args()

    folder = Path(args.folder)
    if not folder.is_dir():
        sys.exit(f"error: folder does not exist: {folder}")
    if not PRESET_FILE.is_file():
        sys.exit(f"error: preset file not found: {PRESET_FILE}")

    handbrake = find_tool("HandBrakeCLI", r"C:\Program Files\HandBrake\HandBrakeCLI.exe")
    ffprobe = find_tool("ffprobe", r"C:\Program Files\ffmpeg\bin\ffprobe.exe")
    target_bytes = int(args.target_gb * 1024**3)

    videos = sorted(
        p for p in folder.iterdir()
        if p.is_file()
        and p.suffix.lower() in VIDEO_EXTS
        and not p.stem.endswith(OUTPUT_SUFFIX)
    )
    if not videos:
        print(f"no videos found in {folder}")
        return

    failures = 0
    for src in videos:
        out = src.with_name(src.stem + OUTPUT_SUFFIX + ".mp4")

        if out.exists():
            print(f"SKIP {src.name} (output already exists)")
            continue
        if src.stat().st_size <= target_bytes:
            print(f"SKIP {src.name} (already {fmt_gib(src.stat().st_size)}, at/below target)")
            continue

        try:
            duration = probe_duration(ffprobe, src)
        except RuntimeError as e:
            print(f"SKIP {src.name} ({e})")
            failures += 1
            continue

        bitrate = calc_video_bitrate(duration, target_bytes)
        cmd = [
            handbrake,
            "--preset-import-file", str(PRESET_FILE),
            "-Z", PRESET_NAME,
            "-b", str(bitrate),
            "-i", str(src),
            "-o", str(out),
        ]

        print(f"ENCODE {src.name}: {fmt_gib(src.stat().st_size)}, "
              f"{duration / 3600:.2f} h, -b {bitrate} -> {out.name}")
        if args.dry_run:
            print(f"       {' '.join(cmd)}")
            continue

        result = subprocess.run(cmd)
        if result.returncode != 0 or not out.exists():
            print(f"FAIL {src.name} (HandBrakeCLI exit {result.returncode})")
            out.unlink(missing_ok=True)
            failures += 1
            continue

        print(f"DONE {out.name}: {fmt_gib(out.stat().st_size)} "
              f"(was {fmt_gib(src.stat().st_size)})")

    sys.exit(1 if failures else 0)


if __name__ == "__main__":
    main()