# tools/automaster_app/pipeline/tp_limiter.py
"""Lookahead true-peak limiter (ITU-R BS.1770-4 style detection).

Chain: LUFS gain staging -> 4x oversample -> per-sample required gain vs
ceiling -> sliding-minimum lookahead -> Hann attack smoothing -> decimated
exponential release -> apply in OS domain -> downsample -> safety clamp ->
up to three corrective LUFS/limiting passes -> final true-peak trim to
ceiling -> TPDF dither.
"""
import numpy as np
import pyloudnorm as pyln
from scipy.signal import resample_poly
from scipy.ndimage import minimum_filter1d

OS = 4
CTRL_DECIM = 16          # release loop runs at OS-rate/16 (~11 kHz at 44.1k)


def true_peak_db(audio, sr):
    up = resample_poly(audio, OS, 1, axis=0)
    return 20 * np.log10(np.max(np.abs(up)) + 1e-12)


def _gain_curve(amp, ceiling, sr_os, lookahead_ms, release_ms):
    required = np.minimum(1.0, ceiling / np.maximum(amp, 1e-12))
    la = max(1, int(sr_os * lookahead_ms / 1000.0))
    g = minimum_filter1d(required, size=2 * la + 1)
    win = np.hanning(2 * la + 1)
    win /= win.sum()
    g = np.pad(g, la, mode="edge")
    g = np.convolve(g, win, mode="valid")
    g = np.minimum(g, required)                      # smoothing must not overshoot
    # exponential release at decimated control rate
    gd = g[::CTRL_DECIM].copy()
    alpha = np.exp(-CTRL_DECIM / (release_ms / 1000.0 * sr_os))
    prev = 1.0
    for i in range(len(gd)):
        prev = min(gd[i], 1.0 - (1.0 - prev) * alpha)
        gd[i] = prev
    g_rel = np.interp(np.arange(len(g)), np.arange(len(gd)) * CTRL_DECIM, gd)
    return np.minimum(g_rel, required)


def _tpdf_dither(audio, bit_depth, ceiling):
    lsb = 2.0 ** (1 - bit_depth)
    rng = np.random.default_rng(0xD17)
    noise = (rng.random(audio.shape) - rng.random(audio.shape)) * lsb
    return np.clip(audio + noise, -ceiling, ceiling)


def export_dither(audio, bit_depth=24, ceiling_dbtp=-1.0):
    """TPDF dither for PCM export; no-op when bit_depth is not 16 or 24."""
    if bit_depth not in (16, 24):
        return audio
    ceiling = 10 ** (ceiling_dbtp / 20.0)
    return _tpdf_dither(audio, bit_depth, ceiling)


def limit(audio, sr, target_lufs, ceiling_dbtp=-1.0, lookahead_ms=5.0,
          release_ms=None, bit_depth=24):
    meter = pyln.Meter(sr)
    ceiling = 10 ** (ceiling_dbtp / 20.0)
    report = {}

    x = audio
    for _ in range(3):                                # initial + up to two corrective passes
        loud = meter.integrated_loudness(x)
        delta = target_lufs - loud
        if abs(delta) <= 0.5 and _ > 0:
            break
        x = x * 10 ** (delta / 20.0)

        if release_ms is None:
            peak = np.max(np.abs(x)) + 1e-12
            rms = np.sqrt(np.mean(x ** 2)) + 1e-12
            crest_db = 20 * np.log10(peak / rms)
            rel = float(np.clip(50.0 + (crest_db - 10.0) * 15.0, 50.0, 200.0))
        else:
            rel = release_ms

        up = resample_poly(x, OS, 1, axis=0)
        amp = np.max(np.abs(up), axis=1)
        g = _gain_curve(amp, ceiling, sr * OS, lookahead_ms, rel)
        gr_db = float(-20 * np.log10(np.min(g) + 1e-12))
        report["max_gain_reduction_db"] = max(
            report.get("max_gain_reduction_db", 0.0), gr_db)   # max across passes, not last pass
        up *= g[:, None]
        x = resample_poly(up, 1, OS, axis=0)[: audio.shape[0]]
        np.clip(x, -ceiling, ceiling, out=x)          # safety clamp post-decimation

    tp = true_peak_db(x, sr)
    if tp > ceiling_dbtp:
        x *= 10 ** ((ceiling_dbtp - tp) / 20.0)
        np.clip(x, -ceiling, ceiling, out=x)

    if bit_depth in (16, 24):
        x = _tpdf_dither(x, bit_depth, ceiling)

    report["output_lufs"] = float(meter.integrated_loudness(x))
    report["output_true_peak_dbtp"] = float(true_peak_db(x, sr))
    return x, report