"""Distribui um bloco de texto colado (a letra inteira, copiada do Google/Spotify) nos 31 campos do letra.json.

Ninguém precisa saber qual frase vai onde: o script transcreve o áudio do trecho (faster-whisper, tempo por palavra),
casa as palavras transcritas com as do texto colado, dá a cada palavra colada o instante em que é cantada e
põe em cada campo as palavras cantadas na janela de tempo dele. A lista do Spotify (D1–D9) usa as quebras de linha
do texto colado: D1 = a linha antes da 1ª cantada no trecho, D2–D9 = as linhas seguintes.
O script não imprime o texto: só quantas palavras caíram em cada campo e quantas o original tinha.

Uso: cole o texto em ae/LETRA-COMPLETA.txt e rode  python -P ae/letra_automatica.py   (áudio: audio/musica.wav ou audio/musica-original.wav)
     --teste usa a própria transcrição como texto colado (para conferir o encaixe sem texto nenhum).
"""
import json, os, re, sys, difflib
sys.stdout.reconfigure(encoding="utf-8")
AQUI = os.path.dirname(os.path.abspath(__file__)); RAIZ = os.path.dirname(AQUI)
TESTE = "--teste" in sys.argv

# Janelas de cada campo (s): palavras cantadas em [ini, fim) vão para o campo. Caixa alta quando o original usa.
JANELAS = [("A1", 0.0, 1.52), ("A2", 1.52, 2.36), ("A3", 2.36, 2.80), ("B1", 2.80, 4.16), ("B2", 4.16, 4.94),
           ("E2", 8.34, 9.54), ("E3", 9.54, 10.52), ("F1", 10.52, 12.56), ("F2", 12.56, 14.04), ("F3", 14.04, 15.18),
           ("F4", 15.18, 16.54), ("G1", 16.54, 17.54), ("G2", 17.54, 18.78), ("H1", 18.78, 20.64), ("H2", 20.64, 22.38),
           ("I1", 22.38, 23.58), ("J1", 24.95, 28.06), ("K1", 28.06, 28.84), ("K2", 28.84, 31.64), ("K4", 31.64, 99)]
CAIXA_ALTA = {"A3", "B2", "E2", "E3", "F1", "F2", "F3", "F4", "G1", "H1", "H2", "J1", "K1", "K3"}
INICIAL_MAIUSCULA = {"G2", "I1", "I2", "K2", "K4"}   # no original esses campos começam com maiúscula mesmo no meio do verso

def norm(w): return re.sub(r"[^a-z0-9]", "", w.lower().replace("’", "'").replace("'", ""))

def transcrever():
    from faster_whisper import WhisperModel
    m = WhisperModel("large-v3", device="cuda", compute_type="float16")
    audio = next(p for p in (os.path.join(RAIZ, "audio", n) for n in ("musica.wav", "musica-original.wav")) if os.path.exists(p))
    segs, _ = m.transcribe(audio, language=None, word_timestamps=True, condition_on_previous_text=False)   # idioma detectado
    return [(w.word.strip(), w.start) for s in segs for w in (s.words or [])]

ditas = transcrever()
if TESTE:
    linhas = [" ".join(w for w, _ in ditas[i:i + 6]) for i in range(0, len(ditas), 6)]
else:
    bruto = open(os.path.join(AQUI, "LETRA-COMPLETA.txt"), encoding="utf-8").read()
    linhas = [l.strip() for l in bruto.splitlines() if l.strip() and not l.strip().startswith("#") and not re.match(r"^\[.*\]$", l.strip())]
coladas = []  # (palavra como está, índice da linha)
for li, l in enumerate(linhas):
    for w in l.split(): coladas.append((w, li))

a = [norm(w) for w, _ in coladas]; b = [norm(w) for w, _ in ditas]
sm = difflib.SequenceMatcher(None, a, b, autojunk=False)
tempo = [None] * len(coladas)
for blk in sm.get_matching_blocks():
    for k in range(blk.size): tempo[blk.a + k] = ditas[blk.b + k][1]
casadas = [i for i, t in enumerate(tempo) if t is not None]
if len(casadas) < 0.5 * len(ditas):
    print(f"AVISO: só {len(casadas)} de {len(ditas)} palavras cantadas casaram. O texto colado é mesmo desta música?")
i0, i1 = casadas[0], casadas[-1]
for i in range(i0, i1 + 1):            # palavras do meio sem par ganham tempo interpolado
    if tempo[i] is None:
        p = max(j for j in casadas if j < i); q = min(j for j in casadas if j > i)
        tempo[i] = tempo[p] + (tempo[q] - tempo[p]) * (i - p) / (q - p)

d = json.load(open(os.path.join(AQUI, "letra.json"), encoding="utf-8")); slots = {s["id"]: s for s in d["slots"]}
def por_janela(ini, fim): return [i for i in range(i0, i1 + 1) if ini <= tempo[i] < fim]
res = {}
for sid, ini, fim in JANELAS: res[sid] = por_janela(ini, fim)
res["I2"] = res["I1"] + por_janela(23.58, 24.95)              # título do vídeo = busca + as palavras grifadas
res["K3"] = res["K2"][-2:]                                     # rótulo do pino 2 = 2 últimas palavras da busca
# Spotify: linha antes da 1ª cantada + 8 linhas seguintes; D5/D6 = linhas cantadas em 4,94–7,36 e 7,36–8,34
l0 = coladas[i0][1]
for k in range(9):
    li = l0 - 1 + k; res[f"D{k + 1}"] = [i for i, (_, l) in enumerate(coladas) if l == li]
for sid, idx in res.items():
    s = slots[sid]; txt = " ".join(coladas[i][0] for i in idx)
    if sid in INICIAL_MAIUSCULA and txt: txt = txt[0].upper() + txt[1:]
    s["texto"] = txt.upper() if sid in CAIXA_ALTA else txt
    s["tempos_texto"] = [round(tempo[i], 3) for i in idx if tempo[i] is not None] if not sid.startswith("D") or sid in ("D5", "D6") else []
    if sid == "J1" and s["tempos_texto"] and s["tempos_texto"][0] < s["canta_em"][0] - 0.1: s["tempos_texto"][0] = s["canta_em"][0]
    if sid in ("D5", "D6"):  # karaokê usa o instante de cada palavra da linha
        s["tempos_texto"] = [round(tempo[i], 3) if tempo[i] is not None else None for i in idx]
if TESTE:
    for s in d["slots"]: s["texto"] = ""; s.pop("tempos_texto", None)
else:
    json.dump(d, open(os.path.join(AQUI, "letra.json"), "w", encoding="utf-8"), ensure_ascii=False, indent=1)
print(f"{len(casadas)} de {len(ditas)} palavras cantadas casaram com o texto colado")
print("campo  palavras (colado / original)")
for s in d["slots"]:
    n = len(res.get(s["id"], [])); print(f"  {s['id']:3s}   {n:2d} / {s['palavras']}" + ("" if abs(n - s["palavras"]) <= 1 else "   <- conferir"))