# -*- coding: utf-8 -*-
"""Generates one WAV per segment with qwen-tts (cloned Rioplatense voice),
checks each one by transcribing it with whisper and retries the bad ones."""
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__)))
# qwentts.cpp checkout and the cloned reference voice. The voice is not part of
# the repository: VOICE_REF is the path prefix of its .spk/.rvq/.txt files
# (produced by `qwen-codec --talker` from a recording and its transcript).
QWENTTS = os.path.expanduser(os.environ.get("QWENTTS_DIR", "~/GIT-MIRRO/qwentts.cpp"))
TTS    = os.path.join(QWENTTS, "build", "qwen-tts")
MODEL  = os.path.join(QWENTTS, "models", "qwen-talker-1.7b-base-Q8_0.gguf")
CODEC  = os.path.join(QWENTTS, "models", "qwen-tokenizer-12hz-Q8_0.gguf")
VOICE_REF = os.path.expanduser(os.environ.get("VOICE_REF", "/tmp/Firefox/output"))
REF    = dict(spk=VOICE_REF + ".spk", rvq=VOICE_REF + ".rvq", txt=VOICE_REF + ".txt")

MIN_RATIO  = 0.90   # minimum accepted similarity against the transcription
TIMEOUT    = 150    # seconds per segment; more than that is runaway generation
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_safe(text, out, seed, label=""):
    """Sometimes the model never emits an end and generates forever: cut and change the seed."""
    for k in range(3):
        try:
            synth(text, out, seed + 7919 * k)
            return seed + 7919 * k
        except (subprocess.TimeoutExpired, subprocess.CalledProcessError):
            print(f"  {label} hung with seed {seed + 7919 * k}, retrying", flush=True)
    raise RuntimeError(f"{label}: three attempts hung")

def transcribe(files, workdir):
    """A single whisper invocation for all files (the model loads once)."""
    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, path of the best attempt)
    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_safe(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():        # keep the best attempt, not the last one
        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}] done. audio={data['total_audio']}s  unverified={bad}", flush=True)

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