#!/usr/bin/env python3
"""Batch-compress every video in a folder to HEVC, v2.

Same as compress_hevc.py but built for throughput on NVIDIA GPUs:
uses the "1080P HEVC NVENC v2" preset (hevc_nvenc_1080p_v2.json, must sit
next to this script), decodes on the GPU via NVDEC (--enable-hw-decoding),
and skips decomb entirely (preset has it off — progressive sources only!).

Measured on a GTX 1060: ~400 fps vs ~240 fps for the v1 pipeline (1.65x),
which fully saturates the single NVENC engine — so running 2 jobs in
parallel (-j 2) is supported but gains nothing; one job is already
GPU-bound.

    python compress_hevc_v2.py                      # E:\\Media\\Compress-Temp
    python compress_hevc_v2.py D:\\some\\folder     # any other folder
    python compress_hevc_v2.py -j 2                 # 2 parallel encodes
    python compress_hevc_v2.py --target-gb 3.0      # target size (default 2.5)
    python compress_hevc_v2.py --dry-run            # show plan, encode nothing

WARNING: the v2 preset disables deinterlacing. Only feed it progressive
sources (check with: ffprobe -show_entries stream=field_order ...).

Per-file video bitrate is derived from duration so every output lands near
the target size. Files at/below the target, *_hevc.mp4 outputs, and files
whose output already exists are skipped, so re-running resumes a batch.

Progress: single-job runs show HandBrake's live progress on the console;
with -j 2 each job writes to <output>.log in the same folder instead
(kept only on failure).
"""

import argparse
import shutil
import subprocess
import sys
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

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

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
MAX_JOBS = 2              # consumer Pascal NVENC allows 2 concurrent sessions

print_lock = threading.Lock()


def log(msg):
    with print_lock:
        print(msg, flush=True)


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 run_encode(handbrake, src, out, bitrate, dry_run, show_progress):
    cmd = [
        handbrake,
        "--enable-hw-decoding", "nvdec",
        "--preset-import-file", str(PRESET_FILE),
        "-Z", PRESET_NAME,
        "-b", str(bitrate),
        "-i", str(src),
        "-o", str(out),
    ]
    if dry_run:
        log(f"       {' '.join(cmd)}")
        return True

    log_path = out.with_suffix(".log")
    if show_progress:
        # single job: let HandBrake draw its live progress on the console
        result = subprocess.run(cmd)
    else:
        # parallel: keep consoles from garbling each other, log per job instead
        log(f"START {src.name} (progress: {log_path})")
        with open(log_path, "w") as log_file:
            result = subprocess.run(cmd, stdout=log_file,
                                    stderr=subprocess.STDOUT)

    if result.returncode != 0 or not out.exists():
        log(f"FAIL {src.name} (HandBrakeCLI exit {result.returncode})"
            + ("" if show_progress else f", see {log_path}"))
        out.unlink(missing_ok=True)
        return False

    log_path.unlink(missing_ok=True)
    log(f"DONE {out.name}: {fmt_gib(out.stat().st_size)} "
        f"(was {fmt_gib(src.stat().st_size)})")
    return True


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("-j", "--jobs", type=int, default=1, choices=[1, 2],
                        help="parallel encodes (default: 1, max: 2 — NVENC session limit)")
    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

    jobs = []
    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)
        print(f"QUEUE {src.name}: {fmt_gib(src.stat().st_size)}, "
              f"{duration / 3600:.2f} h, -b {bitrate} -> {out.name}")
        jobs.append((src, out, bitrate))

    with ThreadPoolExecutor(max_workers=args.jobs) as pool:
        futures = [
            pool.submit(run_encode, handbrake, src, out, bitrate,
                        args.dry_run, args.jobs == 1)
            for src, out, bitrate in jobs
        ]
        for future in as_completed(futures):
            if not future.result():
                failures += 1

    sys.exit(1 if failures else 0)


if __name__ == "__main__":
    main()