#!/usr/bin/env python3
"""Voice treatments and audio accents, applied to intervals, video untouched.

Two kinds of thing live here:

  * **voice treatments** — a whole chain that changes how the speaker sounds:
    a clean broadcast chain, a telephone, a megaphone, a whisper. These are applied
    over an interval and crossfaded at the edges, because a chain that switches
    instantly is heard as a glitch rather than as an effect.

  * **accents** — one-shot gestures at a timestamp: a reversed tail before a cut, a
    pitch drop on a punchline, a stutter. These are mixed in, not substituted.

The video stream is always copied. Everything here is an audio decision, and
re-encoding the picture to change the sound is how a master loses a generation of
quality for nothing.

Loudness is measured after the fact and reported against the platform target, so a
chain that quietly cost 4 dB does not ship unnoticed.
"""
from __future__ import annotations

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

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

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

# Chains are written as ffmpeg audio filter strings. Frequencies and ratios are
# chosen for a spoken voice at 48 kHz, not for music.
VOICE_TREATMENTS: dict[str, dict] = {
    "clean": {
        "use": "База для любой речи: убирает гул, поднимает разборчивость, ровняет уровень.",
        "chain": (
            "highpass=f=80,"
            "equalizer=f=240:t=q:w=1.2:g=-2,"      # boxiness out
            "equalizer=f=3000:t=q:w=1.1:g=2,"      # consonants in
            "equalizer=f=8000:t=q:w=1.0:g=1,"      # air
            "acompressor=threshold=-20dB:ratio=3:attack=8:release=140:makeup=2,"
            "alimiter=limit=0.95"
        ),
    },
    "podcast_tight": {
        "use": "Плотный разговорный звук: ближе, ровнее, без динамических провалов.",
        "chain": (
            "highpass=f=90,lowpass=f=16000,"
            "afftdn=nf=-26,"
            "equalizer=f=180:t=q:w=1.0:g=-3,"
            "equalizer=f=2600:t=q:w=1.2:g=2.5,"
            "acompressor=threshold=-24dB:ratio=4:attack=5:release=90:makeup=3,"
            "acompressor=threshold=-14dB:ratio=2.5:attack=20:release=200:makeup=1,"
            "alimiter=limit=0.94"
        ),
    },
    "warm_radio": {
        "use": "Тёплый эфирный тон: мягкий верх, подчёркнутая грудь голоса.",
        "chain": (
            "highpass=f=95,lowpass=f=12000,"
            "equalizer=f=150:t=q:w=1.0:g=2.5,"
            "equalizer=f=5000:t=q:w=1.2:g=-2,"
            "acompressor=threshold=-22dB:ratio=4:attack=10:release=160:makeup=3,"
            "aecho=0.85:0.5:12:0.08"
        ),
    },
    "telephone": {
        "use": "Телефон или голосовое сообщение: узкая полоса и лёгкое искажение.",
        "chain": (
            "highpass=f=320,lowpass=f=3200,"
            "equalizer=f=1800:t=q:w=1.4:g=4,"
            "acompressor=threshold=-18dB:ratio=6:attack=3:release=60:makeup=4,"
            "alimiter=limit=0.9"
        ),
    },
    "megaphone": {
        "use": "Громкоговоритель, объявление, крик в толпу.",
        "chain": (
            "highpass=f=420,lowpass=f=4000,"
            "equalizer=f=1200:t=q:w=1.0:g=6,"
            "acompressor=threshold=-14dB:ratio=8:attack=2:release=40:makeup=5,"
            "aecho=0.8:0.6:40|70:0.3|0.2"
        ),
    },
    "whisper_intimate": {
        "use": "Интимный шёпот: очень близко, тише, с воздухом.",
        "chain": (
            "highpass=f=120,"
            "equalizer=f=6000:t=q:w=1.2:g=3,"
            "equalizer=f=300:t=q:w=1.0:g=-2,"
            "acompressor=threshold=-28dB:ratio=2:attack=15:release=250:makeup=4,"
            "volume=-2dB"
        ),
    },
    "vinyl_lofi": {
        "use": "Ретро-подача: узкая полоса, шум пластинки, лёгкая нестабильность.",
        "chain": (
            "highpass=f=180,lowpass=f=7000,"
            "equalizer=f=900:t=q:w=0.8:g=2,"
            "tremolo=f=1.6:d=0.06,"
            "acompressor=threshold=-20dB:ratio=4:attack=12:release=180:makeup=2"
        ),
    },
    "robot": {
        "use": "Механический голос: для шутки или для цитаты машины.",
        "chain": (
            "highpass=f=200,lowpass=f=6000,"
            "afftfilt=real='hypot(re,im)*cos(0)':imag='hypot(re,im)*sin(0)':win_size=512,"
            "acompressor=threshold=-18dB:ratio=6:attack=4:release=70:makeup=3"
        ),
    },
    "underwater": {
        "use": "Под водой, за стеной, «выключенное» состояние.",
        "chain": (
            "lowpass=f=900,"
            "equalizer=f=400:t=q:w=0.7:g=3,"
            "aecho=0.8:0.7:60:0.4,"
            "volume=-3dB"
        ),
    },
}

# One-shot gestures mixed on top of the voice rather than replacing it.
ACCENTS: dict[str, dict] = {
    "pitch_drop": {
        "use": "Резкое падение тона на выводе или на шутке.",
        "length": 0.6,
        "chain": "asetrate=48000*0.72,aresample=48000,volume=-3dB",
    },
    "pitch_rise": {
        "use": "Быстрый подъём тона: удивление, вопрос.",
        "length": 0.5,
        "chain": "asetrate=48000*1.22,aresample=48000,volume=-3dB",
    },
    "reverse_tail": {
        "use": "Обратный хвост перед склейкой: тянет внимание к переходу.",
        "length": 0.8,
        "chain": "areverse,volume=-8dB,afade=t=in:d=0.6",
    },
    "stutter": {
        "use": "Заикание на слове: акцент в клиповом монтаже.",
        "length": 0.4,
        "chain": "atempo=1.0,aecho=0.9:0.9:35|70|105:0.6|0.4|0.25,volume=-4dB",
    },
    "hall_hit": {
        "use": "Зальный отзвук на одной фразе: значимость.",
        "length": 1.6,
        "chain": "aecho=0.8:0.85:250|420:0.5|0.3,volume=-9dB",
    },
    "radio_squelch": {
        "use": "Обрыв связи: конец мысли, переход в другую сцену.",
        "length": 0.5,
        "chain": "highpass=f=500,lowpass=f=2600,volume=-5dB,tremolo=f=22:d=0.6",
    },
}


def ff(value: float) -> str:
    return f"{value:.6f}".rstrip("0").rstrip(".") or "0"


def main() -> None:
    utf8_stdout()
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("input", nargs="?")
    parser.add_argument("output", nargs="?")
    parser.add_argument("--treatment", default="clean",
                        help="Voice chain for the whole track, or 'none'")
    parser.add_argument("--plan", help="JSON list of {start, end, treatment} intervals")
    parser.add_argument("--accents", help="JSON list of {at, accent} one-shots")
    parser.add_argument("--crossfade", type=float, default=0.12,
                        help="Seconds of crossfade at each treatment boundary")
    parser.add_argument("--loudness", default="reels", choices=list(LOUDNESS))
    parser.add_argument("--no-normalise", action="store_true")
    parser.add_argument("--list", action="store_true", help="Show treatments and accents")
    parser.add_argument("--filter-report")
    args = parser.parse_args()

    if args.list:
        emit({
            "treatments": {name: spec["use"] for name, spec in VOICE_TREATMENTS.items()},
            "accents": {name: {"use": spec["use"], "length": spec["length"]}
                        for name, spec in ACCENTS.items()},
            "how": (
                "Обработка речи применяется к интервалам и сшивается кроссфейдом; "
                "акценты подмешиваются поверх. Видео всегда копируется."
            ),
        })
        return
    if not args.input or not args.output:
        parser.error("input and output are required unless --list is used")

    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")
    # Trim to the video stream when there is one. Using the container duration
    # (the max of the streams) stretches the audio a couple of frames past the
    # picture on every pass, and after a few steps the file fails an A/V gate.
    video_duration = (info.get("video") or {}).get("duration")
    duration = float(video_duration or info.get("duration") or 0.0)

    intervals: list[dict] = []
    if args.plan:
        payload = json.loads(Path(args.plan).read_text(encoding="utf-8"))
        entries = payload.get("intervals") if isinstance(payload, dict) else payload
        for entry in entries or []:
            name = entry.get("treatment", "clean")
            if name not in VOICE_TREATMENTS and name != "none":
                raise SystemExit(f"Unknown treatment '{name}'. See --list")
            intervals.append({
                "start": max(0.0, float(entry.get("start", 0.0))),
                "end": min(duration, float(entry.get("end", duration))),
                "treatment": name,
            })
        intervals.sort(key=lambda item: item["start"])
    else:
        if args.treatment not in VOICE_TREATMENTS and args.treatment != "none":
            raise SystemExit(f"Unknown treatment '{args.treatment}'. See --list")
        intervals = [{"start": 0.0, "end": duration, "treatment": args.treatment}]

    # Fill the gaps between planned intervals with the clean chain so the whole
    # track is covered exactly once.
    covered: list[dict] = []
    cursor = 0.0
    for item in intervals:
        if item["start"] > cursor + 0.01:
            covered.append({"start": cursor, "end": item["start"], "treatment": "clean"})
        covered.append(item)
        cursor = max(cursor, item["end"])
    if cursor < duration - 0.01:
        covered.append({"start": cursor, "end": duration, "treatment": "clean"})
    covered = [item for item in covered if item["end"] - item["start"] > 0.02]

    accents: list[dict] = []
    if args.accents:
        payload = json.loads(Path(args.accents).read_text(encoding="utf-8"))
        entries = payload.get("accents") if isinstance(payload, dict) else payload
        for entry in entries or []:
            name = entry.get("accent")
            if name not in ACCENTS:
                raise SystemExit(f"Unknown accent '{name}'. See --list")
            at = float(entry.get("at", 0.0))
            if 0 <= at < duration:
                accents.append({"at": at, "accent": name,
                                "gain_db": float(entry.get("gain_db", 0.0))})

    graph: list[str] = []
    splits = len(covered) + (len(accents) or 0) + 1
    graph.append(
        f"[0:a]aformat=sample_rates=48000:channel_layouts=stereo,asplit={splits}"
        + "".join(f"[src{i}]" for i in range(splits))
    )

    # Treated segments, each crossfaded into the next so the chain change is not
    # heard as a switch.
    fade = max(0.0, min(args.crossfade, 0.4))
    segment_labels = []
    for index, item in enumerate(covered):
        chain = ("anull" if item["treatment"] == "none"
                 else VOICE_TREATMENTS[item["treatment"]]["chain"])
        length = item["end"] - item["start"]
        pieces = [f"atrim=start={ff(item['start'])}:end={ff(item['end'])}",
                  "asetpts=PTS-STARTPTS", chain]
        if fade > 0 and length > fade * 2.5:
            if index > 0:
                pieces.append(f"afade=t=in:st=0:d={ff(fade)}:curve=tri")
            if index < len(covered) - 1:
                pieces.append(f"afade=t=out:st={ff(length - fade)}:d={ff(fade)}:curve=tri")
        graph.append(f"[src{index}]" + ",".join(pieces) + f"[seg{index}]")
        segment_labels.append(f"[seg{index}]")

    if len(segment_labels) == 1:
        graph.append(f"{segment_labels[0]}anull[voice]")
    else:
        graph.append("".join(segment_labels) + f"concat=n={len(segment_labels)}:v=0:a=1[voice]")

    mix_labels = ["[voice]"]
    for order, accent in enumerate(accents):
        spec = ACCENTS[accent["accent"]]
        source_index = len(covered) + order
        end = min(duration, accent["at"] + spec["length"])
        delay = int(round(accent["at"] * 1000))
        graph.append(
            f"[src{source_index}]atrim=start={ff(accent['at'])}:end={ff(end)},"
            f"asetpts=PTS-STARTPTS,{spec['chain']},"
            f"volume={accent['gain_db']:.2f}dB,"
            f"adelay={delay}|{delay}:all=1[acc{order}]"
        )
        mix_labels.append(f"[acc{order}]")

    if len(mix_labels) > 1:
        graph.append(
            "".join(mix_labels)
            + f"amix=inputs={len(mix_labels)}:normalize=0:dropout_transition=0[mixed]"
        )
    else:
        graph.append("[voice]anull[mixed]")

    target = LOUDNESS[args.loudness]
    # Ceiling matched to the platform's true-peak target: a hardcoded limiter after
    # loudnorm overrides loudnorm's own true-peak control and lands the master above
    # the target it was asked for.
    #
    # `alimiter` measures *sample* peaks, and a signal can sit under the ceiling on
    # every sample while the reconstructed waveform between them goes over — which
    # is exactly what a delivery gate measuring dBTP reports, and what a platform
    # transcoder clips. Limiting at 4x the sample rate puts three intermediate
    # points between every original pair, so the limiter sees most of the
    # overshoot; the extra 0.3 dB of ceiling covers what four times oversampling
    # still misses. Resampling back afterwards can only lower the peak further.
    #
    # The resampler is left at ffmpeg's default: `resampler=soxr` is not compiled
    # into every build (including the common gyan.dev Windows one) and asking for
    # it fails the whole graph with "Requested resampling engine is unavailable"
    # rather than falling back.
    ceiling = 10 ** ((float(target["tp"]) - 0.3) / 20.0)
    limiter = (f"aresample={4 * 48000},"
               f"alimiter=limit={ceiling:.4f}:level=disabled,"
               f"aresample=48000")
    if args.no_normalise:
        graph.append(
            f"[mixed]atrim=duration={ff(duration)},{limiter}[aout]"
        )
    else:
        graph.append(
            f"[mixed]atrim=duration={ff(duration)},"
            f"loudnorm=I={target['i']}:TP={target['tp']}:LRA={target['lra']},"
            f"{limiter}[aout]"
        )

    # The spare split output is unused when there are no accents; drop it.
    used = len(covered) + len(accents)
    if used < splits:
        graph[0] = (
            f"[0:a]aformat=sample_rates=48000:channel_layouts=stereo,asplit={used}"
            + "".join(f"[src{i}]" for i in range(used))
        ) if used > 1 else (
            "[0:a]aformat=sample_rates=48000:channel_layouts=stereo[src0]"
        )

    filter_graph = ";\n".join(graph)
    report_path = (Path(args.filter_report) if args.filter_report
                   else Path(args.output).with_suffix(".audio_filter.txt"))
    report_path.parent.mkdir(parents=True, exist_ok=True)
    report_path.write_text(filter_graph + "\n", encoding="utf-8")

    output = Path(args.output).resolve()
    output.parent.mkdir(parents=True, exist_ok=True)
    command = [
        "ffmpeg", "-y", "-v", "error", "-nostdin", "-i", str(source),
        "-filter_complex_script", str(report_path),
        "-map", "0:v?", "-map", "[aout]",
        "-c:v", "copy", "-c:a", "aac", "-b:a", "192k", "-ar", "48000", "-ac", "2",
        # The AAC encoder pads to a whole frame, so a trimmed audio branch still
        # comes out a few frames longer than the copied video. `-shortest` ends the
        # file with the video stream, which is what keeps A/V equal through a chain
        # of steps instead of growing by two frames each pass.
        "-shortest",
        "-movflags", "+faststart", str(output),
    ]
    completed = subprocess.run(command, capture_output=True, text=True, errors="replace")
    if completed.returncode:
        emit({"status": "failed", "filter_report": str(report_path),
              "tail": (completed.stderr or "").strip()[-1500:]})
        raise SystemExit(completed.returncode)

    before = loudness(source)
    after = loudness(output)
    emit({
        "status": "ok",
        "input": str(source),
        "output": str(output),
        "video": "copied without re-encoding",
        "segments": [
            {"start": round(item["start"], 3), "end": round(item["end"], 3),
             "treatment": item["treatment"],
             "use": VOICE_TREATMENTS.get(item["treatment"], {}).get("use")}
            for item in covered
        ],
        "accents": accents,
        "crossfade_seconds": fade,
        "loudness_before": before,
        "loudness_after": after,
        "loudness_target": target,
        "on_target": (
            after.get("integrated_lufs") is not None
            and abs(after["integrated_lufs"] - target["i"]) <= 1.0
        ),
        "filter_report": str(report_path),
    })


if __name__ == "__main__":
    main()