#!/usr/bin/env python3
"""Trim silence and normalize generated audio candidates with ffmpeg.

Raw generations stay untouched under `raw/`; processed files land beside them.
Deterministic, no network, no key access.
"""
from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path

TARGET_PEAK_DB = -1.5


def probe_duration(path: Path) -> float:
    out = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "json", str(path)],
        capture_output=True, text=True, check=True,
    )
    return float(json.loads(out.stdout)["format"]["duration"])


def measure_peak(path: Path) -> float:
    out = subprocess.run(
        ["ffmpeg", "-hide_banner", "-i", str(path), "-af", "volumedetect", "-f", "null", "-"],
        capture_output=True, text=True, check=False,
    )
    for line in out.stderr.splitlines():
        if "max_volume:" in line:
            return float(line.split("max_volume:")[1].strip().split(" ")[0])
    return 0.0


def process(source: Path, target: Path) -> dict:
    peak = measure_peak(source)
    gain = TARGET_PEAK_DB - peak if peak < 0 else 0.0
    filters = (
        "silenceremove=start_periods=1:start_duration=0.002:start_threshold=-50dB,"
        "areverse,"
        "silenceremove=start_periods=1:start_duration=0.002:start_threshold=-50dB,"
        "areverse,"
        "apad=pad_dur=0.08,"
        f"volume={gain:.2f}dB"
    )
    subprocess.run(
        ["ffmpeg", "-y", "-v", "error", "-i", str(source), "-af", filters,
         "-c:a", "libmp3lame", "-q:a", "2", str(target)],
        check=True,
    )
    # mp3 encoding can overshoot the requested gain; pull it back once if needed.
    rendered_peak = measure_peak(target)
    if rendered_peak > -1.0:
        correction = -1.5 - rendered_peak
        corrected = target.with_name(target.stem + ".corrected.mp3")
        subprocess.run(
            ["ffmpeg", "-y", "-v", "error", "-i", str(target), "-af", f"volume={correction:.2f}dB",
             "-c:a", "libmp3lame", "-q:a", "2", str(corrected)],
            check=True,
        )
        corrected.replace(target)
        gain += correction
    return {
        "source": str(source),
        "output": str(target),
        "source_peak_dbfs": round(peak, 2),
        "applied_gain_db": round(gain, 2),
        "duration_seconds": round(probe_duration(target), 3),
    }


def process_music(source: Path, target: Path) -> dict:
    """Loudness-normalize a music bed without touching its rests or loop length."""
    subprocess.run(
        ["ffmpeg", "-y", "-v", "error", "-i", str(source),
         "-af", "loudnorm=I=-18:TP=-1.5:LRA=11",
         "-c:a", "libmp3lame", "-q:a", "2", str(target)],
        check=True,
    )
    return {
        "source": str(source),
        "output": str(target),
        "mode": "music_loudnorm",
        "duration_seconds": round(probe_duration(target), 3),
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("paths", nargs="+")
    parser.add_argument("--music", action="store_true", help="loudness-normalize without trimming silence")
    parser.add_argument("--report", default=None)
    args = parser.parse_args()
    if args.music:
        raw_dir = Path(args.paths[0]).parent / "raw"
        raw_dir.mkdir(parents=True, exist_ok=True)
        records = []
        for raw in args.paths:
            source = Path(raw)
            archived = raw_dir / source.name
            if not archived.exists():
                archived.write_bytes(source.read_bytes())
            record = process_music(archived, source)
            records.append(record)
            print(f"OK {source.name} dur {record['duration_seconds']}")
        if args.report:
            Path(args.report).write_text(json.dumps({"items": records}, ensure_ascii=False, indent=1))
        return 0
    raw_dir = Path(args.paths[0]).parent / "raw"
    raw_dir.mkdir(parents=True, exist_ok=True)
    records = []
    for raw in args.paths:
        source = Path(raw)
        archived = raw_dir / source.name
        if not archived.exists():
            archived.write_bytes(source.read_bytes())
        record = process(archived, source)
        records.append(record)
        print(f"OK {source.name} peak {record['source_peak_dbfs']} gain {record['applied_gain_db']} dur {record['duration_seconds']}")
    if args.report:
        Path(args.report).write_text(json.dumps({"items": records}, ensure_ascii=False, indent=1))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())