#!/usr/bin/env python3
"""voz_kokoro.py — a voz local, sem cota e sem servidor (Kokoro-82M, na CPU).

O Gemini é a voz de gente, mas a cota grátis é pequena e às vezes o áudio
volta cortado (medido em 25/09). O Kokoro roda nesta máquina: 11 s de áudio
em 3,5 s, mais 4,3 s para carregar o modelo, sem corte no fim (05/10).

Roda pelo `uv` (lib/estudio/kokoro.js monta o comando), com as versões
fixadas: sem `transformers>=4.45` o uv escolhe a 4.12 e o tokenizers não
compila; sem ESPEAK_DATA_PATH o espeak embutido aponta para uma pasta que
não existe.

Não se chama kokoro.py: a pasta do script entra no sys.path e o nome
esconderia o pacote `kokoro` (medido em 05/10).

Uso (a mesma entrada e a mesma saída do `narrar.py roteiro`):
    voz_kokoro.py roteiro --arquivo partes.json --saida narracao.mp3 [--voz pf_dora]
"""

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

TAXA = 24000
RESPIRO = 0.12  # o mesmo respiro curto entre as partes do narrar.py


def _palavras(texto, de, ate):
    """O Kokoro em português não devolve o tempo de cada palavra: ele é
    repartido pelo tamanho da palavra dentro da parte (melhor que nada para
    a tipografia seguir a voz)."""
    ws = texto.split()
    total = sum(len(w) for w in ws) or 1
    out, t = [], de
    for w in ws:
        d = (ate - de) * len(w) / total
        out.append({"de": round(t, 3), "ate": round(t + d, 3), "texto": w})
        t += d
    return out


def comando_roteiro(args):
    import numpy as np
    import soundfile as sf
    from kokoro import KPipeline

    partes = json.loads(Path(args.arquivo).read_text(encoding="utf-8"))
    if not partes:
        return {"ok": False, "error": "o roteiro veio vazio"}
    pipe = KPipeline(lang_code="p")  # português do Brasil
    pausa = np.zeros(int(TAXA * RESPIRO), dtype=np.float32)
    blocos, trechos, palavras, t = [], [], [], 0.0
    for parte in partes:
        texto = re.sub(r"\s+", " ", str(parte.get("texto") or "")).strip()
        if not texto:
            trechos.append({"de": round(t, 3), "ate": round(t, 3), "palavras": []})
            continue
        audio = np.concatenate([np.asarray(a, dtype=np.float32) for _, _, a in pipe(texto, voice=args.voz)])
        dur = len(audio) / TAXA
        ws = _palavras(texto, t, t + dur)
        palavras += ws
        trechos.append({"de": round(t, 3), "ate": round(t + dur, 3), "palavras": ws})
        blocos += [audio, pausa]
        t += dur + RESPIRO
    with tempfile.TemporaryDirectory() as tmp:
        wav = Path(tmp) / "voz.wav"
        sf.write(str(wav), np.concatenate(blocos), TAXA)
        r = subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", str(wav),
                            "-af", "loudnorm=I=-14:TP=-1.5:LRA=9", "-c:a", "libmp3lame", "-q:a", "3",
                            str(args.saida)], capture_output=True, timeout=180)
        if r.returncode != 0:
            return {"ok": False, "error": (r.stderr or b"").decode("utf-8", "ignore")[-300:]}
    return {"ok": True, "arquivo": str(args.saida), "motor": "kokoro", "voz": f"kokoro {args.voz}",
            "segundos": round(t, 3), "partes": trechos, "palavras": palavras}


def main():
    p = argparse.ArgumentParser(description="Narração local do PrimoCode Studio (Kokoro)")
    sub = p.add_subparsers(dest="comando", required=True)
    r = sub.add_parser("roteiro")
    r.add_argument("--arquivo", required=True)
    r.add_argument("--saida", required=True)
    r.add_argument("--voz", default="pm_alex")
    args = p.parse_args()
    try:
        saida = comando_roteiro(args)
    except Exception as e:  # nunca derruba a corrente: o JS passa para a próxima voz
        saida = {"ok": False, "error": f"o kokoro falhou: {str(e)[:300]}"}
    print(json.dumps(saida, ensure_ascii=False))
    sys.exit(0 if saida.get("ok") else 1)


if __name__ == "__main__":
    main()