#!/usr/bin/env python3
"""Denoise a spoken take (and optionally normalize its loudness).

Talking-head audio is usually recorded in a room, not a booth — HVAC, traffic,
laptop fans. Clean it BEFORE tightening: tighten_vo.py finds pauses by audio
level, and a noise floor makes real silence look loud.

Engines, in quality order:
  deepfilternet  neural, best quality. Install: pip install deepfilternet
                 torch torchaudio — deepfilternet does NOT pull torch itself,
                 and the venv's bin/ must be on PATH for `deepFilter` to be
                 found. Use --strict to fail rather than quietly downgrade.
  arnndn         RNNoise via ffmpeg, needs a .rnnn model file
  afftdn         FFT denoise, built into ffmpeg — the default, no install

Loudness is NOT normalized by default. export_variants.py already normalizes
every platform master to -14 LUFS, and normalizing twice compounds pumping.
Use --loudnorm only when handing audio to something outside this pipeline.

Usage:
  clean_audio.py raw/take_01.mp4 --out work/extracts/take_01_clean.mp4
  clean_audio.py raw/take_01.mp4 --out work/extracts/t.mp4 --engine deepfilternet
  clean_audio.py raw/take_01.mp4 --out work/extracts/t.mp4 --loudnorm -16
"""
import argparse, json, shutil, subprocess, sys, tempfile
from pathlib import Path

# afftdn noise reduction in dB. Past ~20 speech starts sounding underwater.
DEFAULT_NR = 12
TRUE_PEAK = -1.5


def run(cmd, label, **kw):
    r = subprocess.run(cmd, capture_output=True, text=True, **kw)
    if r.returncode != 0:
        sys.exit(f"{label} failed:\n{r.stderr[-900:]}")
    return r


def measure_loudness(path):
    """Measure a file's loudness with loudnorm's analysis pass.

    Returns loudnorm's JSON. The measured loudness of `path` is `input_i` —
    `output_i` is loudnorm's prediction for a hypothetical second pass, not a
    reading of this file.
    """
    r = subprocess.run(
        ["ffmpeg", "-hide_banner", "-i", str(path), "-af",
         "loudnorm=print_format=json", "-f", "null", "-"],
        capture_output=True, text=True)
    start = r.stderr.rfind("{")
    end = r.stderr.rfind("}")
    if start == -1 or end == -1:
        return None
    try:
        return json.loads(r.stderr[start:end + 1])
    except json.JSONDecodeError:
        return None


def denoise_filter(engine, nr, model):
    if engine == "arnndn":
        if not model:
            sys.exit("--engine arnndn needs --rnnn-model pointing at a .rnnn "
                     "file (grab one from the RNNoise model repo), or use the "
                     "default afftdn engine which needs no model.")
        return f"arnndn=m='{model}'"
    return f"afftdn=nr={nr}:nf=-25"


def deepfilternet_wav(src, workdir):
    """Extract audio, run DeepFilterNet on it, return the cleaned wav."""
    exe = shutil.which("deepFilter")
    if not exe:
        return None
    raw = workdir / "in.wav"
    run(["ffmpeg", "-y", "-v", "error", "-i", str(src),
         "-ar", "48000", "-ac", "1", str(raw)], "audio extract")
    run([exe, str(raw), "-o", str(workdir)], "deepFilter")
    out = next(workdir.glob("*_DeepFilterNet*.wav"), None)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("source")
    ap.add_argument("--out", required=True)
    ap.add_argument("--engine", choices=("afftdn", "arnndn", "deepfilternet"),
                    default="afftdn")
    ap.add_argument("--nr", type=float, default=DEFAULT_NR,
                    help=f"afftdn noise reduction in dB (default {DEFAULT_NR}). "
                         f"Above ~20 speech starts sounding underwater.")
    ap.add_argument("--rnnn-model", help="Path to a .rnnn model for arnndn")
    ap.add_argument("--strict", action="store_true",
                    help="Fail instead of substituting a weaker denoiser.")
    ap.add_argument("--loudnorm", type=float, nargs="?", const=-14.0,
                    metavar="LUFS",
                    help="Also normalize loudness, two-pass (default -14 LUFS). "
                         "Skip it for takes staying in this pipeline — "
                         "export_variants.py already normalizes.")
    args = ap.parse_args()

    src = Path(args.source)
    if not src.exists():
        sys.exit(f"Source not found: {src}")
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)

    with tempfile.TemporaryDirectory() as td:
        td = Path(td)
        engine = args.engine
        clean_wav = None
        if engine == "deepfilternet":
            clean_wav = deepfilternet_wav(src, td)
            if clean_wav is None and args.strict:
                sys.exit(
                    "deepFilter not on PATH and --strict was given.\n"
                    "  pip install deepfilternet torch torchaudio\n"
                    "  (deepfilternet does NOT pull torch itself, and the "
                    "venv's bin/ must be on PATH for deepFilter to be found)")
            if clean_wav is None:
                print("WARNING: deepFilter not on PATH — falling back to "
                      "ffmpeg afftdn, which is audibly weaker on a noisy room. "
                      "This is NOT the quality you asked for; install "
                      "deepfilternet (plus torch torchaudio — it does not "
                      "pull them itself) or accept the weaker result "
                      "knowingly. Use --strict to make this an error.",
                      file=sys.stderr)
                engine = "afftdn"

        af = []
        if engine != "deepfilternet":
            af.append(denoise_filter(engine, args.nr, args.rnnn_model))

        # Pass 1 — denoise into a temp file. Loudness must be measured AFTER
        # denoising: removing a noise floor changes integrated loudness, so
        # measurements taken from the raw source would mis-target the second
        # pass.
        stage = td / ("denoised" + out.suffix)
        cmd = ["ffmpeg", "-y", "-v", "error", "-i", str(src)]
        if clean_wav is not None:
            cmd += ["-i", str(clean_wav), "-map", "0:v:0", "-map", "1:a:0"]
        else:
            cmd += ["-map", "0:v:0", "-map", "0:a:0"]
        if af:
            cmd += ["-af", ",".join(af)]
        cmd += ["-c:v", "copy", "-c:a", "aac", "-b:a", "192k", "-ar", "48000",
                str(stage)]
        run(cmd, "denoise")

        tmp = out.with_name(out.stem + ".tmp" + out.suffix)
        if args.loudnorm is None:
            run(["ffmpeg", "-y", "-v", "error", "-i", str(stage),
                 "-c", "copy", "-movflags", "+faststart", str(tmp)], "remux")
        else:
            # Pass 2 — normalize using the denoised file's own measurements.
            m = measure_loudness(stage)
            ln = f"loudnorm=I={args.loudnorm}:TP={TRUE_PEAK}:LRA=11"
            if m:
                ln += (f":measured_I={m['input_i']}:measured_TP={m['input_tp']}"
                       f":measured_LRA={m['input_lra']}"
                       f":measured_thresh={m['input_thresh']}"
                       f":offset={m['target_offset']}")
            run(["ffmpeg", "-y", "-v", "error", "-i", str(stage),
                 "-af", ln, "-c:v", "copy", "-c:a", "aac", "-b:a", "192k",
                 "-ar", "48000", "-movflags", "+faststart", str(tmp)],
                "loudnorm")
        tmp.rename(out)

    after = measure_loudness(out)
    print(json.dumps({
        "engine": engine, "engine_requested": args.engine,
        "substituted": engine != args.engine,
        "source": str(src), "output": str(out),
        "loudnorm_target": args.loudnorm,
        "integrated_lufs": after.get("input_i") if after else None,
    }, indent=2))
    print(f"\nClean {out} before tightening — tighten_vo.py reads audio level, "
          f"and a noise floor hides real pauses.", file=sys.stderr)


if __name__ == "__main__":
    main()