#!/usr/bin/env python3
"""Mixa a trilha contínua por baixo da voz, com duck (sidechain), no vídeo JÁ pronto (references/som-e-trilha.md).

A trilha entra depois do corte seco, por cima de tudo, sem pular nos cortes. O script:
  1. mede a voz do vídeo (LUFS integrado);
  2. acerta o ganho da trilha pra ela ficar N dB abaixo da voz (medido na trilha como ela sai, já com o duck);
  3. abaixa a trilha quando a voz entra (sidechaincompress guiado pela voz) e devolve nas pausas longas;
  4. opcional: entra num segundo escolhido (--entrada-em) com um respiro mais alto no começo (--respiro);
  5. fecha com fade no fim e deixa o vídeo em -14 LUFS (loudnorm em 2 passadas, linear); o fade do fim nunca passa de
     um terço do tempo em que a trilha toca (com --entrada-em perto do fim, 1,5 s de fade engoliria a trilha inteira:
     ele diminui sozinho e avisa);
  6. confere: dB abaixo da voz, quanto o duck abaixou, duração, voz no lugar (desvio em ms), volume final.

Uso (<SKILL> é a pasta desta skill; o Python dela tem numpy):
    uv run --directory "<SKILL>" python helpers/mixar_trilha.py --video final.mp4 --trilha musica.wav
        [--abaixo-da-voz 22] [--entrada-em S] [--respiro [SEG]] [--respiro-db 10] [--duck 8] [--fade-fim 1.5]
        [--inicio-trilha 0] [--saida saida.mp4]

Pegadinha do ffmpeg 8: o sidechaincompress corta a saída quando uma das entradas acaba antes. Por isso todas as
entradas passam por apad + atrim na duração exata do vídeo.
"""
from __future__ import annotations

import argparse
import json
import math
import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

try:
    import numpy as np
except ImportError:  # pragma: no cover
    sys.exit('Falta o numpy. Rode com o Python da skill: uv run --directory "<SKILL>" python helpers/mixar_trilha.py ...')

XFADE_LACO = 1.0   # segundos de crossfade quando a trilha é mais curta que o vídeo e precisa repetir
FADE_FIM_PADRAO = 1.5     # segundos de fade da trilha no fim
FADE_FIM_FRACAO = 1 / 3   # o fade do fim ocupa no máximo isso do tempo em que a trilha toca (quando não foi pedido)


def erro(msg: str, code: int = 2):
    print("ERRO: " + msg, file=sys.stderr)
    sys.exit(code)


def run(cmd: list, **kw):
    r = subprocess.run(cmd, capture_output=True, text=True, **kw)
    if r.returncode:
        erro("o ffmpeg falhou:\n" + r.stderr[-2000:])
    return r


def sonda(p: str) -> dict:
    r = subprocess.run(["ffprobe", "-v", "error", "-show_entries",
                        "stream=codec_type,duration,sample_rate:format=duration", "-of", "json", p],
                       capture_output=True, text=True)
    if r.returncode:
        erro(f"não consegui abrir {p}")
    d = json.loads(r.stdout)
    out = dict(formato=float(d.get("format", {}).get("duration", 0) or 0), video=None, audio=None)
    for s in d.get("streams", []):
        if s.get("codec_type") in ("video", "audio") and out[s["codec_type"]] is None:
            out[s["codec_type"]] = float(s.get("duration") or out["formato"])
    return out


def loudness(p: str, af: str = "") -> dict:
    r = subprocess.run(["ffmpeg", "-hide_banner", "-nostats", "-i", p, "-vn", "-af",
                        af + "loudnorm=I=-14:TP=-1:LRA=11:print_format=json", "-f", "null", "-"],
                       capture_output=True, text=True).stderr
    try:
        return json.loads(r[r.rindex("{"): r.rindex("}") + 1])
    except ValueError:
        erro(f"não consegui medir o volume de {p}")


def cadeia_loudnorm(wav: str, alvo: float = -14.0, tp: float = -1.0) -> tuple[str, dict]:
    """-14 LUFS em 2 passadas, sempre linear. Se o pico real passaria de -1 dBTP, um limitador com o atraso
    compensado (latency=1, não desloca o áudio) segura só os picos antes."""
    j = loudness(wav)
    ganho = alvo - float(j["input_i"])
    pre = ""
    if float(j["input_tp"]) + ganho > tp - 0.2:
        lim = 10 ** ((tp - ganho - 1.0) / 20)
        pre = f"alimiter=limit={max(0.0625, min(1.0, lim)):.5f}:attack=5:release=60:level=false:latency=1,"
        j = loudness(wav, pre)
    af = (pre + f"loudnorm=I={alvo}:TP={tp}:LRA=11:measured_I={j['input_i']}:measured_TP={j['input_tp']}"
          f":measured_LRA={j['input_lra']}:measured_thresh={j['input_thresh']}:offset={j['target_offset']}"
          ":linear=true:print_format=json,aresample=48000")
    return af, j

def finalizar(video: str, wav: str, saida: str) -> tuple[dict, str, dict]:
    """junta a imagem (copiada) com o áudio em -14 LUFS e confere o pico real DEPOIS do AAC: o codificador pode
    passar do teto em até ~2 dB. Se passou de -0,5 dBTP, refaz com o teto mais baixo (um limitador leve só nos
    picos; o volume integrado continua -14 LUFS)."""
    tp = -1.0
    for tentativa in range(2):
        af, j = cadeia_loudnorm(wav, tp=tp)
        r = subprocess.run(["ffmpeg", "-hide_banner", "-nostats", "-y", "-i", video, "-i", wav, "-map", "0:v:0",
                            "-map", "1:a:0", "-c:v", "copy", "-af", af, "-c:a", "aac", "-b:a", "256k",
                            "-movflags", "+faststart", saida], capture_output=True, text=True)
        if r.returncode:
            erro("o ffmpeg falhou ao gravar o vídeo final:\n" + r.stderr[-1500:])
        tipo = "linear"
        try:
            tipo = json.loads(r.stderr[r.stderr.rindex("{"): r.stderr.rindex("}") + 1]).get("normalization_type",
                                                                                           "linear")
        except ValueError:
            pass
        fim = loudness(saida)
        if float(fim["input_tp"]) <= -0.5:
            break
        tp = round(tp - (float(fim["input_tp"]) + 0.8), 2)
    return j, tipo, fim



def audio_mono(p: str, sr: int = 16000) -> "np.ndarray":
    raw = subprocess.run(["ffmpeg", "-v", "error", "-i", p, "-vn", "-ac", "1", "-ar", str(sr), "-f", "f32le",
                          "-acodec", "pcm_f32le", "-"], capture_output=True).stdout
    return np.frombuffer(raw, dtype=np.float32)


def desvio_sincronia(ref: str, out: str, janelas: int = 10, sr: int = 16000) -> float:
    """maior deslocamento (ms) do áudio de saída em relação ao de entrada, medido em janelas de 1 s com fala"""
    a, b = audio_mono(ref, sr), audio_mono(out, sr)
    n = min(len(a), len(b))
    w, L = sr, int(0.05 * sr)
    if n < 3 * w:
        return 0.0
    en = [float((a[i:i + w] ** 2).mean()) for i in range(L, n - w - L, w)]
    idx = sorted(range(len(en)), key=lambda k: -en[k])[:janelas]
    pior = 0.0
    for k in idx:
        i = L + k * w
        c = np.correlate(b[i - L:i + w + L], a[i:i + w], mode="valid")
        pior = max(pior, abs((int(c.argmax()) - L) / sr * 1000))
    return round(pior, 2)


def env_db(x: "np.ndarray", sr: int, hop: float) -> "np.ndarray":
    w = max(1, int(sr * hop))
    n = len(x) // w
    return 10 * np.log10((x[: n * w].astype(np.float64).reshape(n, w) ** 2).mean(axis=1) + 1e-12)


def ajustar_fade_fim(tocando: float, fade_fim: float) -> tuple[float, str | None]:
    """(fade do fim que cabe, aviso). A trilha precisa soar no nível certo por um tempo: o fade do fim fica com no
    máximo um terço do tempo em que ela toca. Visto em 29/09/2026: --entrada-em 7,95 num vídeo de 9,3 s deixava
    1,35 s de trilha, e o fade padrão de 1,5 s engolia tudo."""
    maximo = round(max(0.1, tocando * FADE_FIM_FRACAO), 2)
    if fade_fim <= maximo:
        return fade_fim, None
    return maximo, (f"a trilha toca só {tocando:.2f} s; um fade de {fade_fim:g} s no fim engoliria quase tudo. Usei "
                    f"{maximo:.2f} s (um terço). Para outro valor, passe --fade-fim.")


# ------------------------------------------------------------------ filtro
def cadeia_trilha(dur_trilha: float, inicio_trilha: float, precisa: float, entrada: float, dur_video: float,
                  ganho_db: float, respiro: float, respiro_db: float, fade_fim: float) -> tuple[str, int]:
    """filtro da trilha (entrada 1) até o rótulo [mus]: laço se precisar, ganho, respiro, fades, atraso,
    apad+atrim na duração do vídeo. Devolve (filtro, quantas cópias da trilha)."""
    util = dur_trilha - inicio_trilha
    partes = []
    if util >= precisa + 0.05:
        partes.append(f"[1:a]atrim=start={inicio_trilha:.4f}:duration={precisa:.4f},asetpts=PTS-STARTPTS[t0];")
        copias = 1
    else:   # repete a trilha com crossfade até cobrir o vídeo
        extra = precisa - util
        copias = 1 + math.ceil(extra / max(1.0, dur_trilha - XFADE_LACO)) + 1
        labs = "".join(f"[c{k}]" for k in range(copias))
        partes.append(f"[1:a]asplit={copias}{labs};")
        partes.append(f"[c0]atrim=start={inicio_trilha:.4f},asetpts=PTS-STARTPTS[x0];")
        for k in range(1, copias):
            partes.append(f"[x{k - 1}][c{k}]acrossfade=d={XFADE_LACO}:c1=tri:c2=tri[x{k}];")
        partes.append(f"[x{copias - 1}]atrim=duration={precisa:.4f},asetpts=PTS-STARTPTS[t0];")
    cad = [f"aformat=sample_fmts=fltp:sample_rates=48000:channel_layouts=stereo", f"volume={ganho_db:.3f}dB"]
    if respiro > 0:
        rg = 10 ** (respiro_db / 20)
        rampa = 0.6
        cad.append(f"volume='if(lt(t,{respiro:.3f}),{rg:.5f},if(lt(t,{respiro + rampa:.3f}),"
                   f"{rg:.5f}+(1-{rg:.5f})*(t-{respiro:.3f})/{rampa},1))':eval=frame")
    cad.append(f"afade=t=in:st=0:d={0.4 if entrada > 0 else 0.15}")
    if fade_fim > 0:
        cad.append(f"afade=t=out:st={max(0.0, precisa - fade_fim):.4f}:d={min(fade_fim, precisa):.4f}")
    fim = ["apad", f"atrim=end={dur_video:.6f}", "asetpts=PTS-STARTPTS"]
    if entrada > 0:
        # silêncio concatenado antes da música. No ffmpeg 8 o adelay seguido de asetpts zera o atraso
        # e a música entrava no segundo 0 (visto em 27/09/2026 com --entrada-em 20.3).
        partes.append("[t0]" + ",".join(cad) + "[tpre];")
        partes.append(f"anullsrc=r=48000:cl=stereo,atrim=duration={entrada:.4f},"
                      "aformat=sample_fmts=fltp:sample_rates=48000:channel_layouts=stereo[sil];")
        partes.append("[sil][tpre]concat=n=2:v=0:a=1," + ",".join(fim) + "[mus];")
    else:
        partes.append("[t0]" + ",".join(cad + fim) + "[mus];")
    return "".join(partes), copias


def filtro_completo(args, dur_video: float, dur_trilha: float, ganho_db: float, limiar: float,
                    duck_db: float | None = None) -> str:
    precisa = dur_video - args.entrada_em
    trilha, _ = cadeia_trilha(dur_trilha, args.inicio_trilha, precisa, args.entrada_em, dur_video, ganho_db,
                              args.respiro, args.respiro_db, args.fade_fim)
    mix = 1 - 10 ** (-(args.duck if duck_db is None else duck_db) / 20)
    return (trilha +
            "[0:a]aformat=sample_fmts=fltp:sample_rates=48000:channel_layouts=stereo,"
            f"apad,atrim=end={dur_video:.6f},asetpts=PTS-STARTPTS,asplit=2[voz][sc];"
            f"[mus][sc]sidechaincompress=threshold={limiar:.6f}:ratio=20:attack=15:release=350:knee=2:"
            f"detection=rms:link=maximum:mix={mix:.4f}[duck];"
            "[duck]asplit=2[d1][d2];"
            "[voz][d1]amix=inputs=2:normalize=0:duration=first[mix]")


def renderizar(args, dur_video, dur_trilha, ganho_db, limiar, mix_wav, stem_wav, duck_db=None):
    fc = filtro_completo(args, dur_video, dur_trilha, ganho_db, limiar, duck_db)
    run(["ffmpeg", "-v", "error", "-y", "-i", args.video, "-i", args.trilha, "-filter_complex", fc,
         "-map", "[mix]", "-c:a", "pcm_f32le", mix_wav, "-map", "[d2]", "-c:a", "pcm_f32le", stem_wav])


def medir_duck(voz_path: str, stem_path: str, seco_path: str, depois_de: float) -> dict:
    """compara a trilha com duck e a mesma trilha sem duck, em janelas de 100 ms: quanto ela abaixa quando a voz
    fala, quanto tempo fica abaixada, e voz x trilha durante a fala"""
    sr = 16000
    v, m, s = audio_mono(voz_path, sr), audio_mono(stem_path, sr), audio_mono(seco_path, sr)
    n = min(len(v), len(m), len(s))
    ev, em, es = env_db(v[:n], sr, 0.1), env_db(m[:n], sr, 0.1), env_db(s[:n], sr, 0.1)
    k0 = int(depois_de / 0.1) + 1
    ev, em, es = ev[k0:], em[k0:], es[k0:]
    viva = es > -80     # só onde a trilha toca
    if viva.sum() < 20:
        return {}
    fala = (ev > np.percentile(ev, 90) - 18) & viva
    res = {}
    if fala.sum() > 5:
        red = es[fala] - em[fala]
        res["duck_na_fala_db"] = round(float(np.median(red)), 1)
        mf = 10 * np.log10((10 ** (em[fala] / 10)).mean())
        vf = 10 * np.log10((10 ** (ev[fala] / 10)).mean())
        res["voz_menos_trilha_na_fala_db"] = round(float(vf - mf), 1)
    res["tempo_abaixada_pct"] = round(float(((es - em) > 3)[viva].mean() * 100), 0)
    return res


def main(argv=None):
    try:
        sys.stdout.reconfigure(line_buffering=True)   # mensagens aparecem na hora, mesmo em log
    except AttributeError:
        pass
    ap = argparse.ArgumentParser(description="Trilha contínua por baixo da voz, com duck, medida em dB abaixo da voz.")
    ap.add_argument("--video", required=True, help="vídeo pronto (sem trilha; pode ter efeitos sonoros)")
    ap.add_argument("--trilha", required=True, help="arquivo da música (wav, mp3, m4a...)")
    ap.add_argument("--abaixo-da-voz", type=float, default=22.0,
                    help="quantos dB a trilha fica abaixo da voz (padrão 22; 28 ou mais some no celular)")
    ap.add_argument("--entrada-em", type=float, default=0.0, help="segundo do vídeo em que a trilha entra (padrão 0)")
    ap.add_argument("--respiro", type=float, nargs="?", const=2.0, default=0.0,
                    help="a trilha entra mais alta por SEG segundos (padrão 2) e depois assenta")
    ap.add_argument("--respiro-db", type=float, default=10.0, help="quanto o respiro fica acima do nível normal (10)")
    ap.add_argument("--duck", type=float, default=8.0, help="quanto a trilha abaixa quando a voz fala, no máximo (8 dB)")
    ap.add_argument("--fade-fim", type=float, default=None,
                    help="fade da trilha no fim do vídeo em segundos (1.5; com pouca trilha, um terço do que ela toca)")
    ap.add_argument("--inicio-trilha", type=float, default=0.0, help="de que segundo da música começar (0)")
    ap.add_argument("--saida", help="vídeo de saída (padrão: <video>_trilha.mp4 ao lado)")
    a = ap.parse_args(argv)

    for b in ("ffmpeg", "ffprobe"):
        if not shutil.which(b):
            erro(f"não achei o {b}. Instale o ffmpeg (no Mac: brew install ffmpeg).")
    for p in (a.video, a.trilha):
        if not Path(p).exists():
            erro(f"não achei {p}")
    saida = a.saida or str(Path(a.video).with_name(Path(a.video).stem + "_trilha.mp4"))
    if Path(saida).resolve() in (Path(a.video).resolve(), Path(a.trilha).resolve()):
        erro("a saída não pode ter o mesmo nome da entrada")
    if not 6 <= a.abaixo_da_voz <= 40:
        erro("--abaixo-da-voz fica entre 6 e 40 dB (o normal é 20 a 24)")
    sv, st = sonda(a.video), sonda(a.trilha)
    if sv["audio"] is None:
        erro("o vídeo não tem áudio (a voz). A trilha é medida contra a voz.")
    if st["audio"] is None:
        erro("o arquivo da trilha não tem áudio")
    D = sv["video"] or sv["formato"]
    Dm = st["audio"] or st["formato"]
    if a.entrada_em >= D - 1:
        erro(f"--entrada-em {a.entrada_em} passa do fim do vídeo ({D:.2f} s)")
    if a.inicio_trilha >= Dm - 1:
        erro(f"--inicio-trilha {a.inicio_trilha} passa do fim da música ({Dm:.2f} s)")
    if a.fade_fim is None:
        a.fade_fim, aviso_fade = ajustar_fade_fim(D - a.entrada_em, FADE_FIM_PADRAO)
        if aviso_fade:
            print("aviso: " + aviso_fade)
    elif ajustar_fade_fim(D - a.entrada_em, a.fade_fim)[1]:
        print(f"aviso: o --fade-fim {a.fade_fim:g} s ocupa mais de um terço dos {D - a.entrada_em:.2f} s em que a trilha "
              "toca; ficou como você pediu.")
    if a.respiro and a.respiro + 0.6 > D - a.entrada_em - a.fade_fim:
        print(f"aviso: o respiro de {a.respiro:g} s não cabe nos {D - a.entrada_em:.2f} s de trilha: ela fica alta até "
              "o fade do fim. Diminua o --respiro ou entre antes.")

    jv = loudness(a.video)
    Lv = float(jv["input_i"])
    if Lv < -60:
        erro("o vídeo está praticamente mudo; não dá pra medir a voz")
    limiar = max(0.000977, min(0.5, 10 ** ((Lv - 30) / 20)))   # a voz passa ~20 dB acima disso quando fala
    alvo = Lv - a.abaixo_da_voz
    print(f"voz do vídeo: {Lv:.1f} LUFS · alvo da trilha: {alvo:.1f} LUFS ({a.abaixo_da_voz:g} dB abaixo)")
    if Dm - a.inicio_trilha < D - a.entrada_em:
        print(f"a música ({Dm:.1f} s) é mais curta que o trecho ({D - a.entrada_em:.1f} s): vai repetir com "
              f"crossfade de {XFADE_LACO:g} s")

    with tempfile.TemporaryDirectory() as td:
        mix_wav, stem_wav = os.path.join(td, "mix.wav"), os.path.join(td, "stem.wav")
        ganho = 0.0
        diff = None
        for passada in range(3):   # o duck depende só da voz, então o ganho é linear: 1 correção basta
            renderizar(a, D, Dm, ganho, limiar, mix_wav, stem_wav)
            Ls = float(loudness(stem_wav)["input_i"])
            diff = Lv - Ls
            if abs(diff - a.abaixo_da_voz) <= 0.3:
                break
            ganho += alvo - Ls
        seco_wav = os.path.join(td, "seco.wav")
        renderizar(a, D, Dm, ganho, limiar, os.path.join(td, "mix_seco.wav"), seco_wav, duck_db=0.0)
        duck = medir_duck(a.video, stem_wav, seco_wav, a.entrada_em + (a.respiro + 0.6 if a.respiro else 0.0))
        jm, tipo, fim = finalizar(a.video, mix_wav, saida)

    so = sonda(saida)
    sinc = desvio_sincronia(a.video, saida)
    print(f"trilha: ganho {ganho:+.1f} dB · ficou {diff:.1f} dB abaixo da voz (LUFS integrado, já com o duck)")
    if duck.get("duck_na_fala_db") is not None:
        print(f"duck: a trilha abaixa {duck['duck_na_fala_db']:.1f} dB quando a voz fala (abaixada em "
              f"{duck['tempo_abaixada_pct']:.0f}% do tempo em que a trilha toca; volta nas pausas longas e no fim) · durante a fala a voz fica "
              f"{duck['voz_menos_trilha_na_fala_db']:.1f} dB acima da trilha")
    print(f"saída: {saida}")
    print(f"  duração vídeo {so['video']:.3f} s, áudio {so['audio']:.3f} s (entrada {D:.3f} s) · volume final "
          f"{fim['input_i']} LUFS, pico {fim['input_tp']} dBTP (ajuste {tipo}) · voz no lugar: desvio {sinc} ms")
    if abs(so["audio"] - D) > 0.05:
        print("  aviso: a duração do áudio não bate com a do vídeo; confira antes de entregar")
    if sinc > 5:
        print("  aviso: a voz escorregou; não entregue sem conferir")
    if tipo != "linear":
        print("  aviso: o ajuste de volume saiu dinâmico; confira se a voz ficou natural")
    rel = dict(video=a.video, trilha=a.trilha, saida=saida, voz_lufs=Lv, abaixo_pedido_db=a.abaixo_da_voz,
               abaixo_medido_db=round(diff, 2), ganho_trilha_db=round(ganho, 2), entrada_em=a.entrada_em,
               respiro_s=a.respiro, respiro_db=a.respiro_db, duck_max_db=a.duck, fade_fim_s=a.fade_fim, **duck,
               duracao_entrada=round(D, 3), duracao_saida_video=round(so["video"], 3),
               duracao_saida_audio=round(so["audio"], 3), lufs_final=float(fim["input_i"]),
               pico_final_dbtp=float(fim["input_tp"]), ajuste=tipo, desvio_voz_ms=sinc)
    Path(saida + ".trilha.json").write_text(json.dumps(rel, ensure_ascii=False, indent=1))


if __name__ == "__main__":
    main()