# -*- coding: utf-8 -*-
"""Genera un WAV por segmento con qwen-tts (voz clonada rioplatense),
verifica cada uno transcribiendo con whisper y reintenta los que salen mal."""
import json, os, re, shutil, subprocess, sys, unicodedata, wave
from difflib import SequenceMatcher

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from narration import SCRIPTS
from numspell import expand

ROOT   = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
TTS    = os.path.expanduser("~/GIT-MIRRO/qwentts.cpp/build/qwen-tts")
MODEL  = os.path.expanduser("~/GIT-MIRRO/qwentts.cpp/models/qwen-talker-1.7b-base-Q8_0.gguf")
CODEC  = os.path.expanduser("~/GIT-MIRRO/qwentts.cpp/models/qwen-tokenizer-12hz-Q8_0.gguf")
REF    = dict(spk="/tmp/Firefox/output.spk", rvq="/tmp/Firefox/output.rvq", txt="/tmp/Firefox/output.txt")

MIN_RATIO  = 0.90   # similitud minima aceptada contra la transcripcion
TIMEOUT    = 150    # segundos por segmento; mas que eso es generacion desbocada
MAX_ROUNDS = 4

def norm(s):
    s = expand(s)
    s = unicodedata.normalize("NFD", s.lower())
    s = "".join(c for c in s if unicodedata.category(c) != "Mn")
    s = re.sub(r"[^a-z0-9ñ ]+", " ", s)
    return " ".join(s.split())

def dur(path):
    with wave.open(path) as w:
        return w.getnframes() / w.getframerate()

def synth(text, out, seed, timeout=TIMEOUT):
    cmd = [TTS, "--model", MODEL, "--codec", CODEC,
           "--ref-spk", REF["spk"], "--ref-rvq", REF["rvq"], "--ref-text", REF["txt"],
           "--lang", "spanish", "--seed", str(seed), "--temp", "0.85", "--top-p", "0.95",
           "-o", out]
    subprocess.run(cmd, input=text.encode(), stdout=subprocess.DEVNULL,
                   stderr=subprocess.DEVNULL, check=True, timeout=timeout)


def synth_seguro(text, out, seed, etiqueta=""):
    """A veces el modelo no emite fin y genera hasta el infinito: corto y cambio semilla."""
    for k in range(3):
        try:
            synth(text, out, seed + 7919 * k)
            return seed + 7919 * k
        except (subprocess.TimeoutExpired, subprocess.CalledProcessError):
            print(f"  {etiqueta} se colgo con semilla {seed + 7919 * k}, reintento", flush=True)
    raise RuntimeError(f"{etiqueta}: tres intentos colgados")

def transcribe(files, workdir):
    """Una sola invocacion de whisper para todos los archivos (carga el modelo una vez)."""
    if not files:
        return {}
    subprocess.run(["conda", "run", "-n", "stt", "--no-capture-output",
                    "whisper-ctranslate2", "--language", "es", "--model", "small",
                    "--output_dir", workdir, "--output_format", "txt", *files],
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
    out = {}
    for f in files:
        t = os.path.join(workdir, os.path.splitext(os.path.basename(f))[0] + ".txt")
        out[f] = open(t, encoding="utf-8").read().strip() if os.path.exists(t) else ""
    return out

def build(name):
    spec = SCRIPTS[name]
    outdir = os.path.join(ROOT, "audio", name)
    work   = os.path.join(ROOT, "tmp", name)
    os.makedirs(outdir, exist_ok=True); os.makedirs(work, exist_ok=True)

    segs  = spec["segments"]
    paths = [os.path.join(outdir, f"{i:02d}.wav") for i in range(len(segs))]
    seeds = [1000 + i for i in range(len(segs))]
    pend  = list(range(len(segs)))
    score = {}
    best  = {}                      # i -> (ratio, ruta del mejor intento)
    keep  = os.path.join(work, "best")
    os.makedirs(keep, exist_ok=True)

    for rnd in range(MAX_ROUNDS):
        if not pend:
            break
        for i in pend:
            seeds[i] = synth_seguro(segs[i][0], paths[i], seeds[i], f"[{name}] {i:02d}")
        tr = transcribe([paths[i] for i in pend], work)
        nxt = []
        for i in pend:
            r = SequenceMatcher(None, norm(segs[i][0]), norm(tr[paths[i]])).ratio()
            score[i] = r
            if i not in best or r > best[i][0]:
                b = os.path.join(keep, f"{i:02d}.wav")
                shutil.copy(paths[i], b)
                best[i] = (r, b)
            flag = "ok " if r >= MIN_RATIO else "RE "
            print(f"  [{name}] {i:02d} r={r:.2f} {flag}{dur(paths[i]):5.2f}s | {tr[paths[i]][:60]}", flush=True)
            if r < MIN_RATIO:
                seeds[i] += 7919
                nxt.append(i)
        pend = nxt
        if pend:
            print(f"  [{name}] ronda {rnd+2}: reintentando {pend}", flush=True)

    for i, (r, b) in best.items():        # me quedo con el mejor intento, no el ultimo
        shutil.copy(b, paths[i])
        score[i] = r

    data = {"name": name, "title": spec["title"], "segments": [
        {"i": i, "tts": segs[i][0], "sub": segs[i][1], "wav": paths[i],
         "dur": round(dur(paths[i]), 3), "score": round(score.get(i, 0), 3)}
        for i in range(len(segs))]}
    data["total_audio"] = round(sum(s["dur"] for s in data["segments"]), 2)
    with open(os.path.join(outdir, "segments.json"), "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    bad = [s["i"] for s in data["segments"] if s["score"] < MIN_RATIO]
    print(f"[{name}] listo. audio={data['total_audio']}s  sin_verificar={bad}", flush=True)

if __name__ == "__main__":
    for n in (sys.argv[1:] or list(SCRIPTS)):
        build(n)