#!/usr/bin/env python3
"""Optional voices for on-screen dialogue (a visitor and an assistant), with word timings for captions.

  python3 voice.py tts lines.json --out vo/ [--model eleven_v3] [--take t1] [--only 0,2]
  python3 voice.py check vo/ [--take t1]               speech-to-text each line; flags any word that came out wrong
  python3 voice.py tighten vo/ --pick 0=t1,2=t2 --out vo-final/ [--tempo assistant=1.1] [--max-gap 0.22]
  python3 voice.py sync vo-final/                      how far captions run ahead of the audio (shift the audio by it)

lines.json:
  {"voices": {"user": "<ElevenLabs voice id>", "assistant": "<voice id>"},
   "say": {"$129,": "a hundred and twenty-nine dollars,", "12th,": "twelfth,"},
   "lines": {"0": ["user", "How do I export my data?"], "1": ["assistant", "It's under Settings. I'll take you there."]}}
Captions keep the written words; `say` swaps in what the voice should pronounce (numbers, prices,
dates, codes), and each caption word is mapped to the start of its spoken words, so captions still
appear exactly as they are said. Each line becomes <take>-m<i>.mp3 + .json ({"words": [[word, start_s]...], "end": s}).
eleven_v3 sounds the most human (no speed setting; use `tighten --tempo`); eleven_multilingual_v2
accepts voice_settings.speed but sometimes adds filler words, so always run `check`.
Needs ELEVENLABS_API_KEY with Text to Speech and Speech to Text permissions; tighten needs numpy.
"""
import argparse, base64, glob, json, os, re, subprocess, uuid, urllib.request
from common import key, http_json, ffmpeg, decode

ap = argparse.ArgumentParser()
ap.add_argument("cmd", choices=["tts", "check", "tighten", "sync"]); ap.add_argument("path")
ap.add_argument("--out", default="vo"); ap.add_argument("--model", default="eleven_v3"); ap.add_argument("--take", default="t1")
ap.add_argument("--only", default=""); ap.add_argument("--pick", default=""); ap.add_argument("--tempo", default="")
ap.add_argument("--max-gap", type=float, default=0.22)
a = ap.parse_args()
H = lambda: {"xi-api-key": key("ELEVENLABS_API_KEY")}


def transcribe(path, words=False):
    b = uuid.uuid4().hex; ext = os.path.splitext(path)[1][1:]
    fields = f"--{b}\r\nContent-Disposition: form-data; name=\"model_id\"\r\n\r\nscribe_v1\r\n"
    if words: fields += f"--{b}\r\nContent-Disposition: form-data; name=\"timestamps_granularity\"\r\n\r\nword\r\n"
    body = (fields + f"--{b}\r\nContent-Disposition: form-data; name=\"file\"; filename=\"a.{ext}\"\r\nContent-Type: audio/{'wav' if ext == 'wav' else 'mpeg'}\r\n\r\n").encode() \
        + open(path, "rb").read() + f"\r\n--{b}--\r\n".encode()
    req = urllib.request.Request("https://api.elevenlabs.io/v1/speech-to-text", data=body, headers={**H(), "Content-Type": f"multipart/form-data; boundary={b}"})
    with urllib.request.urlopen(req, timeout=120) as r:
        return json.load(r)


norm = lambda s: re.sub(r"[^a-z0-9 ]", "", s.lower().replace("’", "'").replace("-", " ")).split()

if a.cmd == "tts":
    L = json.load(open(a.path)); os.makedirs(a.out, exist_ok=True)
    only = {s for s in a.only.split(",") if s}
    for i, (role, caption) in L["lines"].items():
        if only and i not in only: continue
        cap = caption.split(" "); spoken, starts, pos = [], [], 0
        for w in cap:
            s = L.get("say", {}).get(w, w); starts.append(pos); spoken.append(s); pos += len(s) + 1
        vs = {"stability": 0.5, "similarity_boost": 0.8, "style": 0.25, "use_speaker_boost": True}
        if a.model != "eleven_v3": vs["speed"] = 1.1
        d = http_json(f"https://api.elevenlabs.io/v1/text-to-speech/{L['voices'][role]}/with-timestamps?output_format=mp3_44100_192",
                      {"text": " ".join(spoken), "model_id": a.model, "voice_settings": vs}, H(), timeout=180)
        open(f"{a.out}/{a.take}-m{i}.mp3", "wb").write(base64.b64decode(d["audio_base64"]))
        al = d.get("alignment") or d["normalized_alignment"]; st = al["character_start_times_seconds"]
        meta = {"role": role, "text": caption, "spoken": " ".join(spoken), "end": al["character_end_times_seconds"][-1],
                "words": [[w, st[min(p, len(st) - 1)]] for w, p in zip(cap, starts)]}
        json.dump(meta, open(f"{a.out}/{a.take}-m{i}.json", "w"), ensure_ascii=False)
        print(f"m{i} {role:9s} {meta['end']:.2f}s  {len(spoken) / max(meta['end'], 0.1):.1f} words/s")

elif a.cmd == "check":
    for js in sorted(glob.glob(f"{a.path}/{a.take}-m*.json") or glob.glob(f"{a.path}/m*.json")):
        meta = json.load(open(js)); audio = next(p for p in (js[:-5] + ".wav", js[:-5] + ".mp3") if os.path.exists(p))
        heard = transcribe(audio)["text"]
        want = norm(meta.get("spoken", meta["text"]))
        ok = norm(heard) == want or norm(heard) == norm(meta["text"])
        print(f"{'OK  ' if ok else 'DIFF'} {os.path.basename(js)[:-5]:8s} heard: {heard!r}" + ("" if ok else f"  wanted: {meta.get('spoken', meta['text'])!r}"))

elif a.cmd == "tighten":
    import numpy as np
    SR = 44100; os.makedirs(a.out, exist_ok=True)
    tempo = {k: float(v) for k, v in (p.split("=") for p in a.tempo.split(",") if p)}
    for pair in a.pick.split(","):
        i, tk = pair.split("=")
        meta = json.load(open(f"{a.path}/{tk}-m{i}.json"))
        x = decode(f"{a.path}/{tk}-m{i}.mp3", SR).copy()
        w = int(0.01 * SR); db = 20 * np.log10(np.array([np.sqrt(np.mean(x[k:k + w] ** 2)) for k in range(0, len(x) - w, w)]) + 1e-9)
        loud = db > max(-45, db.max() - 38); first, last = int(np.argmax(loud)), len(loud) - 1 - int(np.argmax(loud[::-1]))
        cuts, k = [], first
        while k <= last:  # silent runs inside the line longer than max-gap become 0.14 s
            if not loud[k]:
                j = k
                while j <= last and not loud[j]: j += 1
                if (j - k) * 0.01 > a.max_gap: cuts.append((k * 0.01 + 0.07, j * 0.01 - 0.07))
                k = j
            else: k += 1
        a0 = max(0, first * 0.01 - 0.03); a1 = min(len(x) / SR, max((last + 1) * 0.01 + 0.08, meta["end"] + 0.12))
        keep, s = [], a0
        for c0, c1 in cuts: keep.append((s, c0)); s = c1
        keep.append((s, a1))
        def remap(t):
            o = 0.0
            for k0, k1 in keep:
                if t <= k0: return o
                if t < k1: return o + t - k0
                o += k1 - k0
            return o
        f = int(0.008 * SR); parts = []
        for k0, k1 in keep:
            seg = x[int(k0 * SR):int(k1 * SR)].copy(); m = min(f, len(seg) // 2)
            seg[:m] *= np.linspace(0, 1, m); seg[len(seg) - m:] *= np.linspace(1, 0, m); parts.append(seg)
        tp = tempo.get(meta["role"], 1.0)
        subprocess.run([ffmpeg(), "-v", "error", "-y", "-f", "f32le", "-ar", str(SR), "-ac", "1", "-i", "-", "-af", f"atempo={tp}", "-ac", "2", f"{a.out}/m{i}.wav"],
                       input=np.concatenate(parts).astype(np.float32).tobytes(), check=True)
        out = {"role": meta["role"], "text": meta["text"], "spoken": meta.get("spoken", meta["text"]),
               "words": [[wd, round(remap(t) / tp, 3)] for wd, t in meta["words"]], "end": round(remap(meta["end"]) / tp, 3)}
        json.dump(out, open(f"{a.out}/m{i}.json", "w"), ensure_ascii=False)
        print(f"m{i:<3} {meta['end']:.2f}s -> {out['end']:.2f}s ({len(cuts)} pauses shortened, tempo {tp})")

else:  # sync
    offs = []
    for js in sorted(glob.glob(f"{a.path}/m*.json")):
        meta = json.load(open(js)); heard = [w for w in transcribe(js[:-5] + ".wav", True)["words"] if w.get("type", "word") == "word"]
        d = [h["start"] - o[1] for o, h in zip(meta["words"], heard)]; offs += d
        print(f"{os.path.basename(js)[:-5]:4s} captions lead the audio by {sorted(d)[len(d) // 2] * 1000:.0f} ms")
    print(f"median lead {sorted(offs)[len(offs) // 2] * 1000:.0f} ms: place the audio that much earlier in the mix (vo_shift)")