#!/usr/bin/env python3
"""
Batch LUFS normalization using ffmpeg (two-pass loudnorm), enforcing a peak
limit (default linear 0.93) and converting to 44.1 kHz, mono, PCM 16-bit.
Recursively processes all .wav files under a directory, atomically overwrites
originals, and deletes any files that fail processing. Logs results to CSV.

Important:
 - Requires ffmpeg in PATH (ffmpeg loudnorm). Recommended ffmpeg >= 4.x.
 - Requires Python packages: soundfile, numpy
   Install with: pip install soundfile numpy

Usage examples:
  python luf_norm_ffmpeg_no_quarantine.py --input ../Dataset
  python luf_norm_ffmpeg_no_quarantine.py --input ../Dataset --target -18.0 --peak 0.93
  python luf_norm_ffmpeg_no_quarantine.py --input ../Dataset --workers 4
"""
from __future__ import annotations

import argparse
import csv
import json
import os
import re
import shutil
import subprocess
import tempfile
import uuid
from functools import partial
from multiprocessing import Pool
from pathlib import Path
from typing import Dict, Optional

import numpy as np
import soundfile as sf

# Regex used to extract JSON printed by ffmpeg loudnorm (it prints a JSON block to stderr)
_JSON_RE = re.compile(r"\{[\s\S]*?\}", re.MULTILINE)


def find_wav_files(root: Path):
    return sorted(root.rglob("*.wav"))


def _extract_loudnorm_json(stderr_text: str) -> Dict:
    matches = list(_JSON_RE.finditer(stderr_text))
    if not matches:
        raise RuntimeError("ffmpeg loudnorm did not produce JSON measurement in stderr")
    # Use last JSON block
    js = matches[-1].group(0)
    return json.loads(js)


def ffmpeg_measure(in_path: Path, target_lufs: float, tp_db: float) -> Dict:
    """
    First-pass: run ffmpeg loudnorm with print_format=json to get measurement.
    Returns parsed measurement dict.
    """
    cmd = [
        "ffmpeg",
        "-hide_banner",
        "-nostats",
        "-y",
        "-i",
        str(in_path),
        "-af",
        f"loudnorm=I={target_lufs}:TP={tp_db}:LRA=7:print_format=json",
        "-ar",
        "44100",
        "-ac",
        "1",
        "-f",
        "null",
        "-",
    ]
    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if proc.returncode != 0:
        raise RuntimeError(f"ffmpeg measure failed (rc={proc.returncode}): {proc.stderr.splitlines()[-20:]}")
    return _extract_loudnorm_json(proc.stderr)


def ffmpeg_apply_second_pass(in_path: Path, out_path: Path, target_lufs: float, tp_db: float, measured: Dict, use_limiter: bool = False, limiter_db: Optional[float] = None):
    """
    Second pass: apply loudnorm with measured_* values to produce normalized output.
    If use_limiter True, append an alimiter with limit=limiter_db (dB) after loudnorm.
    Writes to out_path (wav).
    """
    # Required measured fields for the second pass
    measured_I = measured.get("input_i")
    measured_TP = measured.get("input_tp")
    measured_LRA = measured.get("input_lra")
    measured_thresh = measured.get("input_thresh")
    target_offset = measured.get("target_offset", 0.0)

    if None in (measured_I, measured_TP, measured_LRA, measured_thresh):
        raise RuntimeError(f"Missing measured fields for second pass: {measured}")

    # Build audio filter string
    af = (
        f"loudnorm=I={target_lufs}:TP={tp_db}:LRA=7:"
        f"measured_I={measured_I}:measured_TP={measured_TP}:measured_LRA={measured_LRA}:"
        f"measured_thresh={measured_thresh}:offset={target_offset}"
    )
    if use_limiter:
        if limiter_db is None:
            raise RuntimeError("Limiter requested but limiter_db is None")
        # alimiter.limit expects a linear amplitude in [0.0625, 1], not dB.
        linear_limit = float(10 ** (limiter_db / 20.0))
        # Clamp to valid range to avoid ffmpeg errors
        linear_limit = max(0.0625, min(1.0, linear_limit))
        af = f"{af},alimiter=limit={linear_limit}"
    cmd = [
        "ffmpeg",
        "-hide_banner",
        "-nostats",
        "-y",
        "-i",
        str(in_path),
        "-af",
        af,
        "-ar",
        "44100",
        "-ac",
        "1",
        "-c:a",
        "pcm_s16le",
        str(out_path),
    ]
    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if proc.returncode != 0:
        # include some stderr context for debugging
        raise RuntimeError(f"ffmpeg apply failed (rc={proc.returncode}). stderr excerpt: {proc.stderr.splitlines()[-40:]}")


def _unique_dest(dest: Path) -> Path:
    """Return a unique path by appending a short uuid if dest exists."""
    if not dest.exists():
        return dest
    stem = dest.stem
    suffix = dest.suffix
    parent = dest.parent
    while True:
        new_name = f"{stem}_{uuid.uuid4().hex[:8]}{suffix}"
        candidate = parent / new_name
        if not candidate.exists():
            return candidate


def atomic_replace(tmp_path: Path, dest: Path):
    """
    Atomically replace dest with tmp_path. tmp_path should exist.
    """
    try:
        os.replace(str(tmp_path), str(dest))
    except Exception:
        # fallback to move
        shutil.move(str(tmp_path), str(dest))


def validate_audio_and_get_peak(path: Path) -> Dict:
    """
    Read audio file and return dict with sr, channels, peak, and whether samples are finite.
    """
    data, sr = sf.read(str(path), dtype="float32", always_2d=False)
    channels = 1 if data.ndim == 1 else int(data.shape[1])
    mono = data if data.ndim == 1 else np.mean(data, axis=1)
    finite = bool(np.all(np.isfinite(mono)))
    peak = float(np.max(np.abs(mono))) if mono.size else 0.0
    return {"sr": int(sr), "channels": channels, "finite": finite, "peak": peak}


def process_one(in_path: Path, target_lufs: float, peak_linear: float, workers_idx: int = 0) -> Dict:
    """
    Process a single file:
      - 1st pass measure (ffmpeg loudnorm print json)
      - 2nd pass apply loudnorm using measured values -> tmp out
      - validate tmp out; if peak > peak_linear then re-run second pass with limiter
      - if still invalid or non-finite, delete original and record error
      - otherwise atomically overwrite original with tmp out
    Returns a record dict for CSV logging.
    """
    rec = {
        "path": str(in_path),
        "input_i": None,
        "input_tp": None,
        "input_lra": None,
        "input_thresh": None,
        "target_offset": None,
        "final_lufs": None,   # may be filled from measured output_i if present
        "final_peak": None,
        "sr": None,
        "channels": None,
        "deleted": False,
        "error": None,
    }

    try:
        # compute dBTP from linear peak
        tp_db = 20.0 * np.log10(peak_linear) if peak_linear > 0 else -999.0
        limiter_db = tp_db  # limiter in dB to use if reprocessing needed

        # 1) Measure
        measured = ffmpeg_measure(in_path, target_lufs, tp_db)

        # record measured fields if available
        rec["input_i"] = measured.get("input_i")
        rec["input_tp"] = measured.get("input_tp")
        rec["input_lra"] = measured.get("input_lra")
        rec["input_thresh"] = measured.get("input_thresh")
        rec["target_offset"] = measured.get("target_offset")

        # Use a temp file in same directory to write output, then atomic replace
        out_tmp = Path(in_path.parent) / (in_path.stem + f".tmp_{uuid.uuid4().hex[:8]}.wav")

        # 2) Second pass: apply loudnorm (no limiter first)
        try:
            ffmpeg_apply_second_pass(in_path, out_tmp, target_lufs, tp_db, measured, use_limiter=False)
        except Exception as e:
            # second pass failed -> delete the original and record error
            rec["error"] = f"ffmpeg second pass failed: {e}"
            try:
                in_path.unlink()
                rec["deleted"] = True
            except Exception as e2:
                rec["error"] += f"; failed deleting original: {e2}"
            # cleanup tmp if present
            if out_tmp.exists():
                try:
                    out_tmp.unlink()
                except Exception:
                    pass
            return rec

        # 3) validate output
        try:
            info = validate_audio_and_get_peak(out_tmp)
            rec["sr"] = info["sr"]
            rec["channels"] = info["channels"]
            rec["final_peak"] = float(info["peak"])
            if not info["finite"]:
                raise RuntimeError("Non-finite samples in ffmpeg output")
        except Exception as e:
            rec["error"] = (rec["error"] + "; " if rec["error"] else "") + f"Validation failed after second pass: {e}"
            try:
                in_path.unlink()
                rec["deleted"] = True
            except Exception as e2:
                rec["error"] += f"; failed deleting original: {e2}"
            if out_tmp.exists():
                try:
                    out_tmp.unlink()
                except Exception:
                    pass
            return rec

        # 4) If peak > peak_linear (allow tiny epsilon), re-run with limiter
        eps = 1e-6
        if rec["final_peak"] > peak_linear + eps:
            # Run second pass again but append alimiter to ensure peaks are clamped
            # write to a new temporary file
            out_tmp2 = Path(in_path.parent) / (in_path.stem + f".tmp_lim_{uuid.uuid4().hex[:8]}.wav")
            try:
                ffmpeg_apply_second_pass(in_path, out_tmp2, target_lufs, tp_db, measured, use_limiter=True, limiter_db=limiter_db)
            except Exception as e:
                rec["error"] = (rec["error"] + "; " if rec["error"] else "") + f"ffmpeg limiter pass failed: {e}"
                # delete original and cleanup
                try:
                    in_path.unlink()
                    rec["deleted"] = True
                except Exception as e2:
                    rec["error"] += f"; failed deleting original: {e2}"
                # cleanup temp files
                for t in (out_tmp, out_tmp2):
                    if t.exists():
                        try:
                            t.unlink()
                        except Exception:
                            pass
                return rec

            # validate out_tmp2
            try:
                info2 = validate_audio_and_get_peak(out_tmp2)
                rec["sr"] = info2["sr"]
                rec["channels"] = info2["channels"]
                rec["final_peak"] = float(info2["peak"])
                if not info2["finite"]:
                    raise RuntimeError("Non-finite samples in limiter output")
                # if still above limit, treat as error
                if rec["final_peak"] > peak_linear + eps:
                    raise RuntimeError(f"Peak still above limit after limiter: {rec['final_peak']:.6f} > {peak_linear:.6f}")
            except Exception as e:
                rec["error"] = (rec["error"] + "; " if rec["error"] else "") + f"Validation failed after limiter: {e}"
                try:
                    in_path.unlink()
                    rec["deleted"] = True
                except Exception as e2:
                    rec["error"] += f"; failed deleting original: {e2}"
                # cleanup temps
                for t in (out_tmp, out_tmp2):
                    if t.exists():
                        try:
                            t.unlink()
                        except Exception:
                            pass
                return rec

            # limiter output is good — replace the tmp pointer to limiter output for final replace
            # remove the first tmp (if exists)
            if out_tmp.exists():
                try:
                    out_tmp.unlink()
                except Exception:
                    pass
            out_tmp = out_tmp2

        # 5) All good: atomically overwrite original with out_tmp
        try:
            atomic_replace(out_tmp, in_path)
        except Exception as e:
            rec["error"] = (rec["error"] + "; " if rec["error"] else "") + f"Atomic replace failed: {e}"
            # in case of a failure, try to delete original to avoid leaving inconsistent data (as requested)
            try:
                in_path.unlink()
                rec["deleted"] = True
            except Exception as e2:
                rec["error"] += f"; failed deleting original after replace failure: {e2}"
            if out_tmp.exists():
                try:
                    out_tmp.unlink()
                except Exception:
                    pass
            return rec

        # Optionally, ffmpeg's JSON sometimes includes output_i (final loudness).
        # We didn't capture the apply-pass stderr here. If you want exact final LUFS, we could
        # measure it using pyloudnorm or another ffmpeg measurement pass. For now, leave final_lufs empty.
        rec["final_lufs"] = None

        return rec

    except Exception as e:
        rec["error"] = str(e)
        # On unexpected error, delete the original as user requested no quarantine
        try:
            in_path.unlink()
            rec["deleted"] = True
        except Exception as e2:
            rec["error"] += f"; failed deleting original during error cleanup: {e2}"
        return rec


def write_csv(logfile: Path, records, fieldnames):
    with open(logfile, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for r in records:
            writer.writerow({k: r.get(k) for k in fieldnames})


def main():
    p = argparse.ArgumentParser(description="Batch LUFS normalization using ffmpeg loudnorm (two-pass), delete bad files.")
    p.add_argument("--input", "-i", default="../Dataset", help="Input root directory (recursive). Default ../Dataset")
    p.add_argument("--target", "-t", type=float, default=-16.0, help="Target integrated LUFS (default -16.0)")
    p.add_argument("--peak", type=float, default=0.93, help="Maximum linear peak allowed after processing (default 0.93)")
    p.add_argument("--log", default="normalize_log.csv", help="CSV log filename")
    p.add_argument("--workers", "-w", type=int, default=1, help="Number of parallel workers (default 1)")
    args = p.parse_args()

    input_root = Path(args.input).expanduser().resolve()
    if not input_root.exists():
        print(f"Input directory does not exist: {input_root}")
        return

    wav_files = find_wav_files(input_root)
    if not wav_files:
        print(f"No WAV files found under {input_root}")
        return

    print(f"Found {len(wav_files)} WAV files.")
    print(f"Target LUFS = {args.target} dB. Peak limit (linear) = {args.peak} (TP ≈ {20.0 * np.log10(args.peak):.3f} dB).")
    print("Converting to 44.1 kHz, mono, PCM_16. Bad files will be deleted (no quarantine).")
    print(f"Workers: {args.workers}")

    worker = partial(process_one, target_lufs=args.target, peak_linear=args.peak)

    records = []
    if args.workers == 1:
        for pth in wav_files:
            rec = worker(pth)
            records.append(rec)
            status = "OK" if not rec.get("error") else "ERR"
            print(f"{pth}: {status} {'(deleted)' if rec.get('deleted') else ''} {rec.get('error') or ''}")
    else:
        with Pool(processes=args.workers) as pool:
            for rec in pool.imap_unordered(worker, wav_files):
                records.append(rec)
                status = "OK" if not rec.get("error") else "ERR"
                print(f"{rec.get('path')}: {status} {'(deleted)' if rec.get('deleted') else ''} {rec.get('error') or ''}")

    # Write CSV
    fieldnames = [
        "path",
        "input_i",
        "input_tp",
        "input_lra",
        "input_thresh",
        "target_offset",
        "final_lufs",
        "final_peak",
        "sr",
        "channels",
        "deleted",
        "error",
    ]
    write_csv(Path(args.log), records, fieldnames)

    total = len(records)
    errors = sum(1 for r in records if r.get("error"))
    deleted = sum(1 for r in records if r.get("deleted"))
    print("\nSummary:")
    print(f"Total files processed: {total}")
    print(f"Files with error: {errors}")
    print(f"Files deleted: {deleted}")
    print(f"CSV log written to: {args.log}")


if __name__ == "__main__":
    main()