#!/usr/bin/env python3
"""Restore the speech that is already in the source, from measurements.

`audio_fx.py` applies a *style* — a named chain that makes a voice sound like a
podcast, a telephone, a megaphone. This script does the step before that: it
looks at what is actually wrong with the recording and repairs only that.

A phone clip recorded in a room has four problems, and they are all measurable:

  * **hum** — mains at 50 or 60 Hz plus harmonics, a needle in the noise
    spectrum rather than a hill. Notch it; a high-pass wide enough to remove
    the third harmonic would take the voice with it.
  * **broadband noise** — air conditioning, street, preamp hiss. Its level and
    its spectral shape both come from the quietest windows in the file, so the
    denoiser is told the real noise floor instead of a guessed `nf=-25`.
  * **rumble and boxiness** — the LF energy under the voice, and the 200–350 Hz
    build-up of a small room. Both are ratios against the speech band.
  * **sibilance** — 5.5–9 kHz against 300–3400 Hz. Cheap microphones and close
    speaking push this ratio up, and a de-esser applied to a voice that does not
    have the problem just makes it dull.

Nothing is applied because it is "good practice". Every stage in the emitted
chain carries the number that triggered it, and a stage whose number is inside
the normal band is reported as *skipped* — with the measurement — rather than
silently applied anyway.

The video stream is copied. Loudness is not normalised here: that happens once,
at the end of the pipeline, in `audio_fx.py`.
"""
from __future__ import annotations

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

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))

from media_probe import loudness, probe  # noqa: E402
from skill_config import emit, require_binary, utf8_stdout  # noqa: E402

RATE = 48000
WINDOW = 1024          # ~21 ms at 48 kHz — level envelope, syllable resolution
HOP = 512
# The spectrum needs a different window from the envelope. At 1024 points a bin
# is 47 Hz wide, so the 125 Hz third-octave band (111–140 Hz) contains no bin at
# all and reads as silence, and a 50 Hz mains peak cannot be told apart from its
# neighbours. 8192 points gives 5.9 Hz bins, which resolves both.
SPEC_WINDOW = 8192
SPEC_HOP = 4096

# Bands the decisions are made in.
BAND_RUMBLE = (20.0, 90.0)
BAND_MUD = (180.0, 380.0)
BAND_SPEECH = (300.0, 3400.0)
BAND_PRESENCE = (2000.0, 5000.0)
BAND_SIBILANCE = (5500.0, 9000.0)
BAND_AIR = (9000.0, 16000.0)

# Long-term average speech spectrum, third-octave centres, relative dB.
# The shape — flat 200–500 Hz, then roughly −3 dB per octave — is what makes a
# recording sound like speech. Comparing a file against it turns "sounds boxy"
# into a number, and it is the reason the EQ moves below are deviations from a
# reference rather than habits. Absolute level is irrelevant: the curve is
# fitted to the file over 315–3150 Hz and only the residual is read.
LTASS: tuple[tuple[float, float], ...] = (
    (63, 38.6), (80, 43.5), (100, 54.4), (125, 57.7), (160, 56.8), (200, 58.2),
    (250, 59.0), (315, 60.3), (400, 60.1), (500, 60.0), (630, 58.5), (800, 57.0),
    (1000, 55.0), (1250, 53.5), (1600, 52.0), (2000, 51.0), (2500, 50.0),
    (3150, 48.0), (4000, 46.5), (5000, 45.0), (6300, 43.0), (8000, 41.0),
    (10000, 39.0), (12500, 36.0), (16000, 33.0),
)


def decode(path: Path, seconds: float | None = None) -> np.ndarray:
    command = ["ffmpeg", "-v", "error", "-nostdin", "-i", str(path)]
    if seconds:
        command += ["-t", f"{seconds:.3f}"]
    command += ["-map", "0:a:0", "-f", "f32le", "-ac", "1", "-ar", str(RATE), "-"]
    done = subprocess.run(command, capture_output=True)
    if done.returncode:
        raise SystemExit((done.stderr or b"").decode("utf-8", "replace")[-1200:])
    return np.frombuffer(done.stdout, dtype="<f4").astype(np.float64)


def db(value: float) -> float:
    return float(20.0 * np.log10(max(value, 1e-12)))


def band_density(spectrum: np.ndarray, freqs: np.ndarray, band: tuple[float, float]) -> float:
    """Mean power per bin, not the sum.

    Summing makes every ratio a statement about bandwidth: 300–3400 Hz holds
    fifteen times as many bins as 180–380 Hz, so a sum-based "mud ratio" reports
    the same number for a boxy room and a dry booth.
    """
    mask = (freqs >= band[0]) & (freqs < band[1])
    if not mask.any():
        # Narrower than one bin: read the nearest bin rather than reporting
        # silence, which would show up as a −200 dB "hole" in the residual.
        return float(spectrum[int(np.argmin(np.abs(freqs - (band[0] + band[1]) / 2)))])
    return float(spectrum[mask].mean())


def third_octave_levels(spectrum: np.ndarray, freqs: np.ndarray) -> dict[float, float]:
    """Power density per third-octave band, in dB."""
    levels: dict[float, float] = {}
    for centre, _ in LTASS:
        low, high = centre / 2 ** (1 / 6), centre * 2 ** (1 / 6)
        if low >= freqs[-1]:
            continue
        levels[centre] = db(np.sqrt(band_density(spectrum, freqs, (low, high))))
    return levels


def ltass_residual(levels: dict[float, float]) -> tuple[dict[float, float], float]:
    """Per-band deviation from the speech reference after a best-fit offset.

    The fit uses 315–3150 Hz — the range every voice occupies and no defect
    lives in exclusively — so the residual outside it is the actual colouration.
    """
    reference = dict(LTASS)
    fit_bands = [c for c in levels if 315 <= c <= 3150 and c in reference]
    if not fit_bands:
        return {c: 0.0 for c in levels}, 0.0
    offset = float(np.mean([levels[c] - reference[c] for c in fit_bands]))
    return ({c: round(levels[c] - reference[c] - offset, 1)
             for c in levels if c in reference}, offset)


def analyse(signal: np.ndarray) -> dict:
    """Everything the chain is built from, in one pass over the file."""
    if signal.size < WINDOW * 4:
        raise SystemExit("Audio too short to measure")

    frames = 1 + (signal.size - WINDOW) // HOP
    window = np.hanning(WINDOW)
    strided = np.lib.stride_tricks.as_strided(
        signal, shape=(frames, WINDOW),
        strides=(signal.strides[0] * HOP, signal.strides[0]),
    )
    rms = np.sqrt((strided ** 2).mean(axis=1) + 1e-20)
    rms_db = 20.0 * np.log10(rms)

    # Speech vs noise by percentile, not by a fixed threshold: a quiet recording
    # and a loud one have completely different absolute levels, and the split
    # that matters is the one inside this file.
    speech_db = float(np.percentile(rms_db, 92))
    noise_db = float(np.percentile(rms_db, 8))
    median_db = float(np.percentile(rms_db, 50))
    snr = speech_db - noise_db

    # Spectral pass, on its own longer framing.
    spec_frames = 1 + max(0, (signal.size - SPEC_WINDOW)) // SPEC_HOP
    if spec_frames < 4:
        raise SystemExit("Audio too short for a spectral estimate")
    spec_window = np.hanning(SPEC_WINDOW)
    spec_strided = np.lib.stride_tricks.as_strided(
        signal, shape=(spec_frames, SPEC_WINDOW),
        strides=(signal.strides[0] * SPEC_HOP, signal.strides[0]),
    )
    spec_rms_db = 20.0 * np.log10(np.sqrt((spec_strided ** 2).mean(axis=1) + 1e-20))
    freqs = np.fft.rfftfreq(SPEC_WINDOW, 1.0 / RATE)
    quiet_idx = np.argsort(spec_rms_db)[:max(3, spec_frames // 8)]
    loud_idx = np.argsort(spec_rms_db)[-max(3, spec_frames // 4):]

    def mean_spectrum(indices: np.ndarray) -> np.ndarray:
        take = indices[:: max(1, indices.size // 200)]
        block = spec_strided[take] * spec_window
        return (np.abs(np.fft.rfft(block, axis=1)) ** 2).mean(axis=0)

    noise_spec = mean_spectrum(quiet_idx)
    speech_spec = mean_spectrum(loud_idx)

    # Hum: a mains harmonic is a spike against its own neighbourhood, so compare
    # each candidate bin to the median of the surrounding 40 Hz. A high-pass
    # cannot do this job — the 150 Hz harmonic lives inside a male voice.
    hum = []
    for base in (50.0, 60.0):
        excess = []
        for harmonic in range(1, 5):
            centre = base * harmonic
            if centre > 400:
                break
            peak_mask = np.abs(freqs - centre) <= 2.5
            around = (np.abs(freqs - centre) <= 25) & ~peak_mask
            if not peak_mask.any() or not around.any():
                continue
            excess.append(db(np.sqrt(noise_spec[peak_mask].max()))
                          - db(np.sqrt(np.median(noise_spec[around]))))
        if excess:
            hum.append({"base_hz": base, "harmonics_db": [round(v, 1) for v in excess],
                        "worst_db": round(max(excess), 1)})
    hum.sort(key=lambda item: item["worst_db"], reverse=True)

    levels = third_octave_levels(speech_spec, freqs)
    residual, _fit_offset = ltass_residual(levels)

    def residual_mean(low: float, high: float) -> float:
        picked = [value for centre, value in residual.items() if low <= centre <= high]
        return round(float(np.mean(picked)), 1) if picked else 0.0

    # Every figure is "dB away from what speech normally does here", so 0 means
    # correct and the sign says which way to move.
    ratios = {
        "rumble_db": residual_mean(*BAND_RUMBLE),
        "mud_db": residual_mean(*BAND_MUD),
        "presence_db": residual_mean(*BAND_PRESENCE),
        "sibilance_db": residual_mean(*BAND_SIBILANCE),
        "air_db": residual_mean(*BAND_AIR),
    }

    # Usable bandwidth: the highest third-octave band still within 12 dB of the
    # reference. A clip band-limited by a codec must not be "opened up" with an
    # exciter — there is nothing up there to lift. Reading it off a cumulative
    # energy percentile instead reports ~2 kHz for every voice, because that is
    # simply where speech power lives.
    bandwidth = 0.0
    for centre in sorted(residual):
        if residual[centre] > -12.0:
            bandwidth = centre * 2 ** (1 / 6)

    peak = float(np.max(np.abs(signal)))
    clipped = float(np.mean(np.abs(signal) > 0.985))

    # Reverb: how fast energy falls in the 250 ms after a speech offset. A dry
    # close mic drops fast; a reflective room holds up.
    decays = []
    speech_mask = rms_db > (noise_db + snr * 0.55)
    for index in range(1, frames - 14):
        if speech_mask[index] and not speech_mask[index + 1]:
            tail = rms_db[index + 1:index + 13]
            if tail.size == 12:
                decays.append(float(rms_db[index] - tail.min()))
    decay_db = round(float(np.median(decays)), 1) if len(decays) >= 5 else None

    return {
        "frames": int(frames),
        "speech_level_db": round(speech_db, 1),
        "noise_floor_db": round(noise_db, 1),
        "median_level_db": round(median_db, 1),
        "snr_db": round(snr, 1),
        "peak_dbfs": round(db(peak), 2),
        "clipped_fraction": round(clipped, 6),
        "usable_bandwidth_hz": round(bandwidth),
        "offset_decay_db_250ms": decay_db,
        "hum": hum,
        "balance": ratios,
        "ltass_residual_db": {str(int(centre)): value for centre, value in sorted(residual.items())},
    }


def build_chain(measured: dict, strength: float, keep_bandwidth: bool) -> tuple[list[str], list[dict]]:
    """Turn measurements into filters. Each entry says which number justified it."""
    chain: list[str] = []
    decisions: list[dict] = []
    balance = measured["balance"]

    def decide(stage: str, applied: bool, because: str, filters: list[str] | None = None):
        decisions.append({"stage": stage, "applied": applied, "because": because,
                          "filters": filters or []})
        if applied and filters:
            chain.extend(filters)

    if measured["clipped_fraction"] > 0.0002:
        decide("declip", True,
               f"{measured['clipped_fraction'] * 100:.3f}% отсчётов на потолке — восстановление формы до всего остального",
               ["adeclip=window=55:overlap=75:arorder=8:threshold=10"])
    else:
        decide("declip", False, f"клиппинга нет ({measured['clipped_fraction'] * 100:.4f}%)")

    # Rumble. Cut where the voice is not: the corner follows the measured LF
    # excess over the speech reference instead of a habitual 80 Hz.
    if balance["rumble_db"] > 4.0:
        corner = 110 if balance["rumble_db"] > 10.0 else 85
        decide("highpass", True,
               f"низ на {balance['rumble_db']:+.1f} дБ выше нормы речи — срез {corner} Гц",
               [f"highpass=f={corner}:poles=2"])
    else:
        decide("highpass", True,
               f"низ в норме ({balance['rumble_db']:+.1f} дБ), только защитный срез 60 Гц",
               ["highpass=f=60:poles=1"])

    worst_hum = measured["hum"][0] if measured["hum"] else None
    if worst_hum and worst_hum["worst_db"] >= 6.0:
        base = worst_hum["base_hz"]
        notches = [f"equalizer=f={base * n:.0f}:t=q:w=18:g=-{min(24, 6 + worst_hum['worst_db']):.0f}"
                   for n in (1, 2, 3) if base * n <= 400]
        decide("hum_notch", True,
               f"сеть {base:.0f} Гц выступает на {worst_hum['worst_db']} дБ над своим окружением",
               notches)
    else:
        decide("hum_notch", False,
               f"сетевой наводки нет (максимум {worst_hum['worst_db'] if worst_hum else 0} дБ над фоном)")

    # Broadband denoise driven by the measured floor. `afftdn` wants the noise
    # floor in dBFS; handing it the real number is the difference between a
    # transparent pass and a voice that swims.
    snr = measured["snr_db"]
    if snr < 34:
        reduction = float(np.clip((34 - snr) * 0.9, 5, 22)) * strength
        floor = float(np.clip(measured["noise_floor_db"] + 4, -80, -20))
        decide("denoise", True,
               f"SNR {snr} дБ, шумовой пол {measured['noise_floor_db']} дБFS — "
               f"снижение {reduction:.0f} дБ по измеренному профилю",
               [f"afftdn=nr={reduction:.1f}:nf={floor:.0f}:nt=w:tn=1"])
    else:
        decide("denoise", False, f"SNR {snr} дБ — шум ниже порога слышимости в миксе")

    # Only correct what deviates. A ±2.5 dB residual is normal voice-to-voice
    # variation, not a defect, and equalising it flattens the speaker's timbre.
    if balance["mud_db"] > 2.5:
        gain = float(np.clip(balance["mud_db"] * 0.7, 1.0, 5.0)) * strength
        decide("de_box", True,
               f"200–380 Гц на {balance['mud_db']:+.1f} дБ выше нормы — комната бубнит, снимаем {gain:.1f} дБ",
               [f"equalizer=f=270:t=q:w=1.1:g=-{gain:.1f}"])
    else:
        decide("de_box", False, f"середина в норме ({balance['mud_db']:+.1f} дБ)")

    if balance["presence_db"] < -2.5:
        gain = float(np.clip(-balance["presence_db"] * 0.7, 1.0, 4.5)) * strength
        decide("presence", True,
               f"полоса согласных на {balance['presence_db']:+.1f} дБ ниже нормы — глухо, "
               f"поднимаем {gain:.1f} дБ на 3 кГц",
               [f"equalizer=f=3000:t=q:w=1.0:g={gain:.1f}"])
    elif balance["presence_db"] > 4.0:
        gain = float(np.clip((balance["presence_db"] - 4.0) * 0.6, 1.0, 3.0)) * strength
        decide("presence", True,
               f"полоса согласных на {balance['presence_db']:+.1f} дБ выше нормы — резко, "
               f"снимаем {gain:.1f} дБ на 3.2 кГц",
               [f"equalizer=f=3200:t=q:w=1.2:g=-{gain:.1f}"])
    else:
        decide("presence", False, f"разборчивость в норме ({balance['presence_db']:+.1f} дБ)")

    # De-esser only if the file actually hisses. Applied blind it dulls a voice
    # that was fine.
    if balance["sibilance_db"] > 3.0:
        intensity = float(np.clip((balance["sibilance_db"] - 3.0) / 8.0, 0.1, 0.45)) * strength
        decide("deesser", True,
               f"шипящие на {balance['sibilance_db']:+.1f} дБ выше нормы — интенсивность {intensity:.2f}",
               [f"deesser=i={intensity:.2f}:m=0.5:f=0.35"])
    else:
        decide("deesser", False, f"шипящих в избытке нет ({balance['sibilance_db']:+.1f} дБ)")

    bandwidth = measured["usable_bandwidth_hz"]
    if keep_bandwidth:
        decide("air", False, "--keep-bandwidth: верх не трогаем")
    elif bandwidth >= 11000 and balance["air_db"] < -3.0:
        gain = float(np.clip(-balance["air_db"] * 0.5, 1.0, 3.0)) * strength
        decide("air", True,
               f"верх есть до {bandwidth:.0f} Гц, но он на {balance['air_db']:+.1f} дБ ниже нормы — "
               f"открываем на {gain:.1f} дБ",
               [f"equalizer=f=10000:t=q:w=0.8:g={gain:.1f}"])
    elif bandwidth >= 11000:
        decide("air", False, f"верх до {bandwidth:.0f} Гц и в норме ({balance['air_db']:+.1f} дБ)")
    else:
        decide("air", False,
               f"полоса обрывается на {bandwidth:.0f} Гц — записи там нечего поднимать, "
               "подъём дал бы только шум")

    # Room tail. A gentle downward expander below the speech threshold pulls the
    # reverb out of the gaps without gating the speech itself.
    decay = measured["offset_decay_db_250ms"]
    if decay is not None and decay < 12.0 and snr > 12:
        threshold = 10 ** ((measured["noise_floor_db"] + snr * 0.35) / 20.0)
        decide("room_expander", True,
               f"спад после фразы {decay} дБ за 250 мс — хвост комнаты, мягкая экспансия пауз",
               [f"agate=threshold={threshold:.5f}:ratio=1.6:attack=8:release=180:knee=6:makeup=1"])
    else:
        decide("room_expander", False,
               f"спад после фразы {decay} дБ — сухо, экспандер не нужен"
               if decay is not None else "переходов речь/тишина мало для оценки хвоста")

    # Levelling last. Two gentle stages beat one aggressive one: the first
    # catches syllables, the second holds the phrase.
    decide("level", True,
           f"разброс уровня {measured['speech_level_db'] - measured['median_level_db']:.1f} дБ "
           "между пиком фразы и медианой — двухступенчатое выравнивание",
           ["acompressor=threshold=-22dB:ratio=2.6:attack=6:release=110:makeup=1.5",
            "acompressor=threshold=-13dB:ratio=2.0:attack=24:release=240:makeup=1"])

    decide("safety", True, "потолок против межсэмплового выброса перед следующим шагом",
           ["alimiter=limit=0.94:level=disabled"])
    return chain, decisions


def main() -> None:
    utf8_stdout()
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("input")
    parser.add_argument("output", nargs="?", help="Omit with --measure-only")
    parser.add_argument("--measure-only", action="store_true",
                        help="Report the diagnosis and the chain it would build, render nothing")
    parser.add_argument("--strength", type=float, default=1.0,
                        help="Scale every corrective gain (0.5 = half, 1.5 = harder)")
    parser.add_argument("--keep-bandwidth", action="store_true",
                        help="Never touch the top octave even if it measures dull")
    parser.add_argument("--report")
    args = parser.parse_args()

    require_binary("ffmpeg")
    source = Path(args.input).resolve()
    if not source.is_file():
        raise SystemExit(f"Not found: {source}")
    info = probe(source)
    if not info.get("has_audio"):
        raise SystemExit(f"{source} has no audio track")

    signal = decode(source)
    measured = analyse(signal)
    strength = float(np.clip(args.strength, 0.2, 2.0))
    chain, decisions = build_chain(measured, strength, args.keep_bandwidth)
    applied = [item for item in decisions if item["applied"]]

    payload = {
        "status": "measured",
        "input": str(source),
        "diagnosis": measured,
        "strength": strength,
        "decisions": decisions,
        "applied_stages": [item["stage"] for item in applied],
        "skipped_stages": [item["stage"] for item in decisions if not item["applied"]],
        "chain": ",".join(chain),
    }

    if args.measure_only or not args.output:
        payload["note"] = ("Ничего не отрендерено. Каждое решение подписано измерением — "
                           "проверь их, прежде чем запускать без --measure-only.")
        emit(payload)
        return

    output = Path(args.output).resolve()
    output.parent.mkdir(parents=True, exist_ok=True)
    video_duration = (info.get("video") or {}).get("duration")
    duration = float(video_duration or info.get("duration") or 0.0)

    graph = (f"[0:a]aformat=sample_rates={RATE}:channel_layouts=stereo,"
             + ",".join(chain)
             + f",atrim=duration={duration:.6f}[aout]")
    command = [
        "ffmpeg", "-y", "-v", "error", "-nostdin", "-i", str(source),
        "-filter_complex", graph,
        "-map", "0:v?", "-map", "[aout]",
        "-c:v", "copy", "-c:a", "aac", "-b:a", "224k", "-ar", str(RATE), "-ac", "2",
        "-shortest", "-movflags", "+faststart", str(output),
    ]
    done = subprocess.run(command, capture_output=True, text=True, errors="replace")
    if done.returncode:
        payload.update({"status": "failed", "tail": (done.stderr or "").strip()[-1500:]})
        emit(payload)
        raise SystemExit(done.returncode)

    after = analyse(decode(output))
    payload.update({
        "status": "ok",
        "output": str(output),
        "video": "copied without re-encoding",
        "after": after,
        "improvement": {
            "snr_db": round(after["snr_db"] - measured["snr_db"], 1),
            "noise_floor_db": round(after["noise_floor_db"] - measured["noise_floor_db"], 1),
            "sibilance_db": round(after["balance"]["sibilance_db"] - measured["balance"]["sibilance_db"], 1),
            "mud_db": round(after["balance"]["mud_db"] - measured["balance"]["mud_db"], 1),
        },
        "loudness_before": loudness(source),
        "loudness_after": loudness(output),
        "loudness_note": "Громкость здесь не нормализуется — это делает audio_fx.py один раз в конце.",
    })
    if after["snr_db"] < measured["snr_db"] - 0.5:
        payload["warning"] = ("SNR не вырос. Проверь, не был ли шум частью речи "
                              "(дыхание, близкий микрофон) — тогда снизь --strength.")
    report = Path(args.report) if args.report else output.with_suffix(".voice_restore.json")
    report.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    payload["report"] = str(report)
    emit(payload)


if __name__ == "__main__":
    main()