#!/usr/bin/env python3
"""Benchmark Video Concat's pipeline stages on real disks.

Runs ffmpeg with exactly the flags src-tauri/src/concat.rs uses, so the numbers
mean what the app will do. Nothing here touches the app itself.

    python tools/bench_concat.py --ffmpeg-dir "C:/Users/me/Documents/VideoConcat" --make 6 --src D:/bench --dst D:/bench/out
    python tools/bench_concat.py --ffmpeg-dir ... --src D:/recordings/session --dst E:/out --reps 3

Stages:
  probe_serial    one ffprobe per file, one after another (the pre-5.1 behaviour)
  probe_parallel  bounded parallel probing
  mp4_faststart   -c copy -movflags +faststart -> .mp4 (the pre-5.1 default)
  mp4_plain       -c copy -> .mp4 (the 5.1 default)
  mkv             -c copy -> .mkv
  raw_copy        byte-for-byte copy of the same input bytes: the I/O floor

Measured 2026-09-25, 6 x 60 s H.264/AAC FLV segments, 1.7 GB:
  local NVMe -> NVMe   faststart 1.77 s   plain 1.47 s   raw 1.39 s
  NAS -> NAS (10GbE)   faststart 27.4 s   plain 10.9 s   raw 6.9 s

Mind the OS file cache: a second run over the same inputs reads them from
RAM. Compare variants within one run, not across runs.
"""
from __future__ import annotations

import argparse
import os
import shutil
import statistics
import subprocess
import threading
import time
from pathlib import Path

NOWIN = 0x08000000 if os.name == "nt" else 0


def tool(ffdir: str | None, name: str) -> str:
    exe = name + (".exe" if os.name == "nt" else "")
    if ffdir and (Path(ffdir) / exe).is_file():
        return str(Path(ffdir) / exe)
    found = shutil.which(name)
    if not found:
        raise SystemExit(f"{name} not found - pass --ffmpeg-dir")
    return found


def make_segments(ff: str, dst: Path, count: int, seconds: int) -> None:
    """High-bitrate H.264/AAC FLV segments, shaped like a BiliBili recording."""
    dst.mkdir(parents=True, exist_ok=True)
    seed = dst / "_seed.flv"
    subprocess.run([ff, "-hide_banner", "-loglevel", "error", "-y",
                    "-f", "lavfi", "-i", "testsrc2=size=1280x720:rate=30",
                    "-f", "lavfi", "-i", "sine=frequency=440:sample_rate=48000", "-t", "10",
                    "-c:v", "libx264", "-preset", "ultrafast", "-b:v", "40M", "-minrate", "40M",
                    "-maxrate", "40M", "-bufsize", "40M", "-x264-params", "nal-hrd=cbr:force-cfr=1",
                    "-g", "60", "-c:a", "aac", "-b:a", "128k", "-f", "flv", str(seed)],
                   check=True, creationflags=NOWIN)
    reps = max(1, seconds // 10)
    lst = dst / "_seed.txt"
    lst.write_text("".join(f"file '{seed.as_posix()}'\n" for _ in range(reps)), encoding="utf-8")
    first = dst / "seg1.flv"
    subprocess.run([ff, "-hide_banner", "-loglevel", "error", "-y", "-f", "concat", "-safe", "0",
                    "-i", str(lst), "-c", "copy", "-f", "flv", str(first)], check=True, creationflags=NOWIN)
    for i in range(2, count + 1):
        shutil.copyfile(first, dst / f"seg{i}.flv")
    seed.unlink()
    lst.unlink()


def probe(fp: str, path: Path) -> float:
    out = subprocess.run([fp, "-v", "error", "-show_entries", "format=duration",
                          "-of", "default=noprint_wrappers=1:nokey=1", str(path)],
                         capture_output=True, text=True, creationflags=NOWIN)
    try:
        return float(out.stdout.strip())
    except ValueError:
        return 0.0


def probe_parallel(fp: str, files: list[Path]) -> None:
    ts = [threading.Thread(target=probe, args=(fp, f)) for f in files]
    for t in ts:
        t.start()
    for t in ts:
        t.join()


def concat(ff: str, files: list[Path], out: Path, faststart: bool) -> float:
    lst = out.parent / f".bench_list_{os.getpid()}.txt"
    lst.write_text("".join(f"file '{f.resolve().as_posix()}'\n" for f in files), encoding="utf-8")
    args = [ff, "-y", "-f", "concat", "-safe", "0", "-i", str(lst), "-c:v", "copy", "-c:a", "copy"]
    if faststart:
        args += ["-movflags", "+faststart"]
    args += ["-nostdin", "-hide_banner", "-loglevel", "error", str(out)]
    t0 = time.perf_counter()
    r = subprocess.run(args, capture_output=True, text=True, creationflags=NOWIN)
    dt = time.perf_counter() - t0
    lst.unlink(missing_ok=True)
    if r.returncode:
        raise RuntimeError(r.stderr)
    return dt


def raw_copy(files: list[Path], out: Path) -> float:
    t0 = time.perf_counter()
    with open(out, "wb") as o:
        for f in files:
            with open(f, "rb") as i:
                shutil.copyfileobj(i, o, 8 * 1024 * 1024)
        o.flush()
        os.fsync(o.fileno())
    return time.perf_counter() - t0


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--ffmpeg-dir")
    ap.add_argument("--src", type=Path, required=True, help="folder with seg*.flv (or where --make writes them)")
    ap.add_argument("--dst", type=Path, required=True, help="where outputs are written (deleted after each run)")
    ap.add_argument("--make", type=int, default=0, metavar="N", help="first generate N segments in --src")
    ap.add_argument("--seconds", type=int, default=60, help="length of each generated segment")
    ap.add_argument("--reps", type=int, default=3)
    ap.add_argument("--only", help="comma list of stages to run")
    a = ap.parse_args()

    ff, fp = tool(a.ffmpeg_dir, "ffmpeg"), tool(a.ffmpeg_dir, "ffprobe")
    if a.make:
        print(f"generating {a.make} x {a.seconds}s segments in {a.src} ...")
        make_segments(ff, a.src, a.make, a.seconds)
    files = sorted(a.src.glob("seg*.flv"))
    if not files:
        raise SystemExit(f"no seg*.flv in {a.src} - use --make")
    only = set(a.only.split(",")) if a.only else None
    want = lambda s: only is None or s in only  # noqa: E731
    a.dst.mkdir(parents=True, exist_ok=True)
    total_mb = sum(f.stat().st_size for f in files) / 1048576
    print(f"inputs {len(files)} files, {total_mb:.0f} MB  |  {a.src} -> {a.dst}  |  median of {a.reps}\n")

    def timed(fn):
        t0 = time.perf_counter()
        fn()
        return time.perf_counter() - t0

    if want("probe_serial"):
        t = statistics.median(timed(lambda: [probe(fp, f) for f in files]) for _ in range(a.reps))
        print(f"{'probe_serial':15s} {t * 1000:8.0f} ms")
    if want("probe_parallel"):
        t = statistics.median(timed(lambda: probe_parallel(fp, files)) for _ in range(a.reps))
        print(f"{'probe_parallel':15s} {t * 1000:8.0f} ms")

    results = {}
    for name, ext, fs in (("mp4_faststart", ".mp4", True), ("mp4_plain", ".mp4", False), ("mkv", ".mkv", False)):
        if not want(name):
            continue
        runs = []
        for _ in range(a.reps):
            out = a.dst / f"bench_{name}{ext}"
            runs.append(concat(ff, files, out, fs))
            out.unlink()
        results[name] = t = statistics.median(runs)
        print(f"{name:15s} {t:8.2f} s  {total_mb / t:6.0f} MB/s")
    if want("raw_copy"):
        runs = []
        for _ in range(a.reps):
            out = a.dst / "bench_raw.bin"
            runs.append(raw_copy(files, out))
            out.unlink()
        t = statistics.median(runs)
        print(f"{'raw_copy':15s} {t:8.2f} s  {total_mb / t:6.0f} MB/s  <- I/O floor")

    if "mp4_faststart" in results and "mp4_plain" in results:
        d = results["mp4_faststart"] - results["mp4_plain"]
        print(f"\n+faststart costs {d:.2f} s = {100 * d / results['mp4_faststart']:.0f}% of a faststart run")


if __name__ == "__main__":
    main()