"""
Mastering chain for generated TTS/voice clips: fixes the "tinny" spectral
imbalance common to neural vocoder output (excess 3-6kHz presence energy,
starved low end relative to a naturally-recorded voice - see the project's
spectral comparison: raw recordings skew toward 200-1000Hz, generated clips
dump ~35% of their energy into 3-6kHz) and adds the density/space that
reads as "produced" rather than dry/synthetic.

Chain order matters: EQ before compression (so the compressor reacts to
the corrected tonal balance, not the raw tinny one), reverb after
compression (so the room tail doesn't get squashed), limiter last.
"""

import argparse

import numpy as np
import soundfile as sf
from pedalboard import (
    Compressor,
    HighpassFilter,
    HighShelfFilter,
    Limiter,
    LowShelfFilter,
    PeakFilter,
    Pedalboard,
    Reverb,
)


def build_chain(
    warmth_db: float,
    presence_cut_db: float,
    air_db: float,
    comp_ratio: float,
    reverb_wet: float,
) -> Pedalboard:
    return Pedalboard([
        HighpassFilter(cutoff_frequency_hz=60),
        LowShelfFilter(cutoff_frequency_hz=150, gain_db=warmth_db, q=0.7),
        PeakFilter(cutoff_frequency_hz=3800, gain_db=presence_cut_db, q=1.0),
        HighShelfFilter(cutoff_frequency_hz=6500, gain_db=air_db, q=0.7),
        Compressor(threshold_db=-20, ratio=comp_ratio, attack_ms=12, release_ms=120),
        Reverb(room_size=0.08, damping=0.85, wet_level=reverb_wet, dry_level=1.0 - reverb_wet),
        Limiter(threshold_db=-1.0, release_ms=100),
    ])


def main() -> None:
    parser = argparse.ArgumentParser(description="Apply a studio mastering chain to a voice clip")
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--warmth-db", type=float, default=5.5, help="Low-shelf boost @150Hz")
    parser.add_argument("--presence-cut-db", type=float, default=-5.0, help="Peak cut @3.8kHz (the tinny band)")
    parser.add_argument("--air-db", type=float, default=-4.5, help="High-shelf @6.5kHz, for softness")
    parser.add_argument("--comp-ratio", type=float, default=2.5)
    parser.add_argument("--reverb-wet", type=float, default=0.025)
    args = parser.parse_args()

    audio, sr = sf.read(args.input, always_2d=True)
    audio = audio.T.astype(np.float32)  # pedalboard wants (channels, samples)

    board = build_chain(
        args.warmth_db, args.presence_cut_db, args.air_db, args.comp_ratio, args.reverb_wet
    )
    processed = board(audio, sr)

    # pedalboard's Limiter is attack/release-based, not a true brick-wall
    # limiter - fast transients can still punch through its threshold (seen
    # in practice: in-memory peaks hit 0dBFS despite threshold_db=-1.0,
    # producing audible clipping distortion on playback). Hard safety
    # ceiling here guarantees the written file never clips regardless.
    target_peak = 10 ** (-1.0 / 20)
    peak = np.max(np.abs(processed))
    if peak > target_peak:
        processed = processed * (target_peak / peak)

    sf.write(args.output, processed.T, sr)
    print(f"[INFO] Saved {args.output}")


if __name__ == "__main__":
    main()