align.py

"""Global CTC forced alignment for Kazakh songs.

Pipeline:
  1. Decode the audio to 16 kHz mono, isolate the vocal stem (HDEMUCS).
  2. Load Meta MMS-fl102 with Kazakh adapter — exposes 108-token Cyrillic
     vocab covering ң, ғ, қ, ұ, ү, ә, і, etc. natively.
  3. Normalize the waveform the way MMS was trained (zero mean / unit
     variance, gain-floored) and run the model over the audio in *overlapping*
     chunks to get one continuous emission stream (CTC log-probs per frame,
     50 frames/sec).
  4. Reduce the reference lyric to the character stream the model can actually
     emit — annotations and punctuation dropped, numbers spelled out, Latin
     script transliterated, out-of-vocab characters removed rather than turned
     into <unk> — then run torchaudio.functional.forced_align over the whole
     emission stream. Globally optimal in one pass: no Whisper segment seeds,
     no per-line greedy mistakes.
  5. Group token spans into word units (one per *original* word, so timings and
     display text can't drift apart) and then into lines, and post-process into
     a timeline safe to drive a karaoke UI from.

Output:
    {
      "lines":  [{startMs, endMs, text, score, confidence, words: [...]}],
      "words":  [{lineIdx, wordIdx, startMs, endMs, text, score, confidence}],
      "alignMode": "ctc-mms-fl102-kaz" | "energy-fallback",
      "meanConfidence": 0..1,
      ...
    }

Accuracy claims in this file are measured — see ../eval/README.md for the
ground-truth harness that produces them.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import tempfile
import time
import unicodedata
import urllib.request
from typing import List, Tuple

import torch
import torchaudio

# On a 2-vCPU HF "cpu-basic" Space, torch/OMP inside the container can misread
# the *host* node's core count and spawn far more threads than we have cores,
# thrashing them and slowing every matmul (HDEMUCS + MMS). Pin to the real
# core budget — a free, quality-neutral speedup. Override with
# ANKUI_ALIGN_THREADS if the Space is later upgraded.
_THREADS = int(os.environ.get("ANKUI_ALIGN_THREADS", "2"))
try:
    torch.set_num_threads(_THREADS)
except Exception:  # noqa: BLE001
    pass


def download(url: str, dest: str) -> None:
    urllib.request.urlretrieve(url, dest)


# A structural marker ("Chorus:", "Қайырмасы 2х", "2. Verse") — printed in
# every lyrics site's copy, sung by nobody.
_SECTION_WORDS = (
    r"chorus|verse|bridge|intro|outro|refrain|pre-?chorus|hook|coda|"
    r"қайырма(?:сы)?|шумақ|көпір|қосымша|"
    r"припев|куплет|проигрыш|бридж|вступление|запев"
)
_SECTION_RE = re.compile(
    rf"^\s*(?:\d+\s*[.)]\s*)?(?:{_SECTION_WORDS})\s*[:：]?\s*"
    rf"(?:[(\[]?\s*\d+\s*[xх×]\s*[)\]]?|[(\[]?\s*[xх×]\s*\d+\s*[)\]]?|\d+)?"
    rf"\s*[:：]?\s*$",
    re.IGNORECASE)
_CREDIT_RE = re.compile(
    r"^\s*(әні\s*(мен)?\s*сөзі|әні|сөзі|мәтіні|авторы?|музыка|music|lyrics|words|"
    r"composer|текст|слова|автор)\s*[:：]",
    re.IGNORECASE)
# A bracketed aside that is an annotation rather than sung words: repeat
# counts, bare numbers, or a section name. "(oh oh oh)" is left alone — that
# one really is sung.
_ANNOTATION_RE = re.compile(
    rf"[(\[]\s*(?:\d+\s*[xх×]|[xх×]\s*\d+|\d+|{_SECTION_WORDS})\s*[)\]]",
    re.IGNORECASE)

_APOSTROPHES = "'’‘ʼʻ`´"
_DASHES = "—–‒―−"


# An "Artist — Title" banner at the very top of a scraped lyric page. Sung by
# nobody, but it aligns *somewhere*, so it drags the real first line with it.
# Only ever applied to the first line, and only when it has no sentence
# punctuation, because a lyric line can legitimately contain a dash.
_HEADER_RE = re.compile(r"^[^,.!?;:]{2,44}\s+[—–-]\s+[^,.!?;:]{2,44}$")


def split_lines(raw: str) -> List[str]:
    """Split pasted lyrics into displayable lines, dropping the scaffolding
    (title banner, section markers, author credits, track numbers, blanks)."""
    out: List[str] = []
    for ln in raw.splitlines():
        s = ln.strip()
        if not s:
            continue
        if _SECTION_RE.match(s) or _CREDIT_RE.match(s):
            continue
        if not out and _HEADER_RE.match(s):
            continue
        s = re.sub(r"^\d+\s*[.)]\s*", "", s).strip()
        if s:
            out.append(s)
    return out


_KAZ_UNITS = ["", "бір", "екі", "үш", "төрт", "бес", "алты", "жеті", "сегіз", "тоғыз"]
_KAZ_TENS = ["", "он", "жиырма", "отыз", "қырық", "елу", "алпыс", "жетпіс",
             "сексен", "тоқсан"]


def kazakh_number(n: int) -> str:
    """Spell 0…9999 in Kazakh. Digits are in the MMS vocab as *digit* tokens,
    which the acoustic model will never emit for a sung number — so a written
    "1000" has to become "мың" or it drags the alignment around it."""
    if n == 0:
        return "нөл"
    parts: List[str] = []
    if n >= 1000:
        th = n // 1000
        parts.append(("" if th == 1 else kazakh_number(th) + " ") + "мың")
        n %= 1000
    if n >= 100:
        h = n // 100
        parts.append(("" if h == 1 else _KAZ_UNITS[h] + " ") + "жүз")
        n %= 100
    if n >= 10:
        parts.append(_KAZ_TENS[n // 10])
        n %= 10
    if n:
        parts.append(_KAZ_UNITS[n])
    return " ".join(p for p in parts if p)


# Qazaq latyn → Cyrillic, longest match first. Users paste Latin-script lyrics
# (and the app itself renders Latyn), but the acoustic model only has Cyrillic
# tokens — every Latin letter would otherwise become <unk>. Bare-ASCII Latin
# (no diacritics) still transliterates to a close phonetic match.
_LATYN_TO_CYRL = [
    ("sh", "ш"), ("ch", "ч"), ("ıa", "я"), ("ıý", "ю"), ("ia", "я"), ("iý", "ю"),
    ("á", "ә"), ("ǵ", "ғ"), ("ń", "ң"), ("ó", "ө"), ("ú", "ү"), ("ý", "у"),
    ("ı", "и"), ("i", "і"), ("y", "ы"), ("u", "ұ"),
    ("a", "а"), ("b", "б"), ("v", "в"), ("g", "г"), ("d", "д"), ("e", "е"),
    ("j", "ж"), ("z", "з"), ("k", "к"), ("q", "қ"), ("l", "л"), ("m", "м"),
    ("n", "н"), ("o", "о"), ("p", "п"), ("r", "р"), ("s", "с"), ("t", "т"),
    ("f", "ф"), ("h", "х"), ("c", "с"), ("w", "в"), ("x", "х"),
]
_CYRILLIC_RE = re.compile(r"[Ѐ-ӿ]")
_LATIN_RE = re.compile(r"[a-záǵńóúýı]", re.IGNORECASE)


def latyn_to_cyrillic(s: str) -> str:
    """Transliterate a Latin-script Kazakh line to Cyrillic."""
    out, i = [], 0
    low = s.lower()
    while i < len(low):
        for src, dst in _LATYN_TO_CYRL:
            if low.startswith(src, i):
                out.append(dst)
                i += len(src)
                break
        else:
            out.append(low[i])
            i += 1
    return "".join(out)


def _maybe_transliterate(s: str) -> str:
    """Transliterate the Latin-script *words*, leaving Cyrillic ones alone.

    Word-by-word rather than line-by-line, because both mixes are real: a line
    of Qazaq latyn, and a Cyrillic line with an English word dropped into it
    ("I love you деп айтты"). Either way the Latin letters have to become
    Cyrillic — the Kazakh acoustic head has Latin tokens in its vocab but never
    emits them, so leaving them in is as bad as an <unk>."""
    out = []
    for word in re.split(r"(\s+)", s):
        latin = len(_LATIN_RE.findall(word))
        cyrl = len(_CYRILLIC_RE.findall(word))
        out.append(latyn_to_cyrillic(word) if latin > cyrl else word)
    return "".join(out)


def normalize_for_align(s: str, vocab=None) -> str:
    """Reduce a lyric line to the character stream the acoustic model can
    actually emit.

    Punctuation is dropped (never sung), annotations like "(2x)" and "[chorus]"
    go with it, numbers are spelled out, Latin-script Kazakh is transliterated,
    and — when `vocab` is supplied — anything still outside the model's
    inventory is dropped rather than encoded as <unk>. That last step matters:
    <unk> is a token the model essentially never emits, so leaving one in the
    target forces the CTC path to plant it somewhere anyway, which skews the
    timing of the words either side of it.
    """
    s = unicodedata.normalize("NFC", s)
    s = _ANNOTATION_RE.sub(" ", s)
    s = s.lower()
    for ch in _APOSTROPHES:
        s = s.replace(ch, "'")
    for ch in _DASHES:
        s = s.replace(ch, " ")
    s = _maybe_transliterate(s)
    s = re.sub(r"\d+", lambda m: " " + kazakh_number(int(m.group(0)))
               if len(m.group(0)) <= 4 else " ", s)
    # Everything that isn't a letter, an in-word apostrophe or a hyphen.
    s = re.sub(r"[^\w'\-\s]|_", " ", s, flags=re.UNICODE)
    s = re.sub(r"[\-']+(?=\s|$)|(?<=\s)[\-']+", " ", s)
    if vocab is not None:
        s = "".join(ch for ch in s if ch in vocab or ch.isspace())
    return re.sub(r"\s+", " ", s).strip()


def _is_degenerate(lines_out: List[dict], duration: float) -> bool:
    """True when forced alignment collapsed — the signature of MMS failing
    on a real musical mix (quiet vocals under a loud master, long
    instrumental intro, etc.). We refuse to report that as success.

    Triggers: most lines share one ~0.25 s start bucket (the "everything
    at 0:00" case the user hit), or the aligned span covers a tiny slice
    of the song.
    """
    if not lines_out:
        return True
    from collections import Counter
    starts = [l["startMs"] for l in lines_out]
    buckets = Counter(s // 250 for s in starts)
    if buckets.most_common(1)[0][1] >= max(2, int(0.5 * len(starts))):
        return True
    span = (max(l["endMs"] for l in lines_out) - min(starts)) / 1000.0
    if duration > 0 and span < 0.12 * duration:
        return True
    return False


def _distribute_by_energy(ref_lines: List[str], waveform: torch.Tensor,
                          sample_rate: int, duration: float) -> List[dict]:
    """Fallback timing when forced alignment is unreliable. Find the
    vocal-active region from an RMS energy envelope and spread the lines
    across it, weighted by line length. Not word-accurate, but always
    spread + monotonic so karaoke scrolls and the user can fine-tune.
    Mirrors the original pipeline's "distribute across the detected vocal
    range" behaviour (see lrcalign/README.md)."""
    import numpy as np
    w = waveform[0].detach().cpu().numpy()
    win = max(1, int(0.05 * sample_rate))  # 50 ms frames
    n = len(w) // win
    t0, t1 = 0.0, duration
    if n >= 4:
        rms = np.sqrt(np.array([(w[i * win:(i + 1) * win] ** 2).mean() for i in range(n)]) + 1e-9)
        k = max(1, int(0.5 / 0.05))  # ~0.5 s smoothing
        smooth = np.convolve(rms, np.ones(k) / k, mode="same")
        thr = max(smooth.mean() * 0.5, smooth.max() * 0.12)
        active = np.where(smooth > thr)[0]
        if len(active) >= 2:
            t0 = max(0.0, active[0] * win / sample_rate - 0.3)
            t1 = min(duration, (active[-1] + 1) * win / sample_rate + 0.3)
    span = max(0.5, t1 - t0)
    lengths = [max(1, len(normalize_for_align(l))) for l in ref_lines]
    total = max(1, sum(lengths))
    out: List[dict] = []
    acc = 0
    for li, line in enumerate(ref_lines):
        s = t0 + span * acc / total
        acc += lengths[li]
        e = t0 + span * acc / total
        e = max(s + 0.2, e)
        # Spread the line's own words across it by length. These are estimates,
        # not measurements — but a karaoke line with no word timings can't
        # highlight at all, which reads as broken rather than approximate. The
        # zero score tells the client (and the user) how much to trust them.
        words: List[dict] = []
        orig = line.split()
        wlens = [max(1, len(w)) for w in orig]
        wtotal = max(1, sum(wlens))
        wacc = 0
        for wi, w in enumerate(orig):
            ws = s + (e - s) * wacc / wtotal
            wacc += wlens[wi]
            we = s + (e - s) * wacc / wtotal
            words.append({
                "wordIdx": wi,
                "startMs": int(ws * 1000),
                "endMs": int(max(ws + 0.05, we) * 1000),
                "text": w,
                "score": 0.0,
                # Explicit, because `score` here is a sentinel and not a log
                # probability: exp(0.0) is 1.0, so deriving confidence from it
                # would advertise these estimates as *maximum* certainty.
                "confidence": 0.0,
            })
        out.append({
            "lineIdx": li,
            "startMs": int(s * 1000),
            "endMs": int(e * 1000),
            "text": line,
            "score": 0.0,
            "confidence": 0.0,
            "estimated": True,
            "words": words,
        })
    return out


_MODEL_CACHE = {}


def load_model(device: str = "cpu"):
    """Load MMS-fl102, then swap in the Kazakh adapter weights.

    The right pattern for HF MMS adapter loading is:
        - `from_pretrained` with no language hint (loads base 78-token head)
        - `load_adapter("kaz")` (replaces lm_head + adapter layers)
        - `tokenizer.set_target_lang("kaz")` (swaps tokenizer vocab)

    Setting `target_lang` on `from_pretrained` reinitializes the lm_head to
    random weights — the model would then output garbage. We saw that
    warning in our first run.
    """
    key = ("fl102", "kaz", device)
    if key in _MODEL_CACHE:
        return _MODEL_CACHE[key]
    from transformers import AutoProcessor, Wav2Vec2ForCTC
    proc = AutoProcessor.from_pretrained("facebook/mms-1b-fl102")
    model = Wav2Vec2ForCTC.from_pretrained("facebook/mms-1b-fl102").to(device)
    proc.tokenizer.set_target_lang("kaz")
    model.load_adapter("kaz")
    model.eval()
    _MODEL_CACHE[key] = (proc, model)
    return proc, model


_SEPARATOR_CACHE = {}
_HDEMUCS_SOURCES = ["drums", "bass", "other", "vocals"]


def load_separator(device: str = "cpu"):
    """Load torchaudio's HDEMUCS music source-separation model (bundled,
    no extra pip). Used to isolate the vocal stem before alignment so the
    speech-trained CTC model isn't fighting the instrumentation — this is
    what makes alignment on real musical mixes (vs. clean acapella/TTS)
    actually land on the words instead of collapsing."""
    key = ("hdemucs", device)
    if key in _SEPARATOR_CACHE:
        return _SEPARATOR_CACHE[key]
    import torchaudio
    bundle = torchaudio.pipelines.HDEMUCS_HIGH_MUSDB_PLUS
    model = bundle.get_model().to(device).eval()
    _SEPARATOR_CACHE[key] = (model, bundle.sample_rate)
    return model, bundle.sample_rate


def separate_vocals(waveform: torch.Tensor, sample_rate: int, device: str = "cpu") -> torch.Tensor:
    """Isolate the vocal stem from a (possibly musical) mix. In/out are
    `[1, T]` mono at `sample_rate`. Returns the input unchanged on any
    failure so alignment still proceeds on the raw audio."""
    import torchaudio
    try:
        model, msr = load_separator(device)
        wav = torchaudio.functional.resample(waveform, sample_rate, msr).repeat(2, 1)  # stereo@msr
        ref_mean, ref_std = wav.mean(), wav.std() + 1e-8
        wav = (wav - ref_mean) / ref_std
        ch, total = wav.shape
        step = int(15.0 * msr)
        out = torch.zeros(4, ch, total)
        i = 0
        with torch.inference_mode():
            while i < total:
                chunk = wav[:, i:i + step].unsqueeze(0).to(device)
                out[:, :, i:i + chunk.shape[-1]] = model(chunk)[0].cpu()
                i += step
        vocals = out[_HDEMUCS_SOURCES.index("vocals")] * ref_std + ref_mean  # [2, T]@msr
        mono = vocals.mean(0, keepdim=True)
        return torchaudio.functional.resample(mono, msr, sample_rate)  # [1, T]@sample_rate
    except Exception as exc:  # noqa: BLE001
        sys.stderr.write(f"[align] vocal separation failed, using raw audio: {exc}\n")
        return waveform


def separate_instrumental(waveform: torch.Tensor, sample_rate: int,
                          device: str = "cpu", progress=None) -> torch.Tensor:
    """Isolate the instrumental ("минус") — drums + bass + other, vocals
    removed — from a musical mix. Input is `[1, T]` or `[2, T]` at
    `sample_rate`; output is `[2, T]` stereo at `sample_rate`.

    Same HDEMUCS forward pass as `separate_vocals` — the model emits all
    four sources at once, so producing the instrumental on top of the
    vocal isolation we already do is essentially free. We just keep the
    complement of the vocal stem. Unlike the alignment path (which wants
    16 kHz mono for the CTC model) karaoke wants full-quality stereo, so
    we neither downmix nor downsample here.

    `progress`, if given, is called with a 0..1 fraction after each chunk
    so the job can report honest progress over the multi-minute pass."""
    import torchaudio
    model, msr = load_separator(device)
    wav = waveform if sample_rate == msr else \
        torchaudio.functional.resample(waveform, sample_rate, msr)
    if wav.size(0) == 1:
        wav = wav.repeat(2, 1)          # mono → stereo (HDEMUCS expects 2ch)
    elif wav.size(0) > 2:
        wav = wav[:2]
    ref_mean, ref_std = wav.mean(), wav.std() + 1e-8
    wav = (wav - ref_mean) / ref_std
    ch, total = wav.shape
    step = int(15.0 * msr)
    out = torch.zeros(4, ch, total)
    i = 0
    with torch.inference_mode():
        while i < total:
            chunk = wav[:, i:i + step].unsqueeze(0).to(device)
            out[:, :, i:i + chunk.shape[-1]] = model(chunk)[0].cpu()
            i += step
            if progress is not None and total > 0:
                progress(min(1.0, i / total))
    v = _HDEMUCS_SOURCES.index("vocals")
    # Sum every stem except vocals — drums+bass+other = the backing track.
    instrumental = sum(out[s] for s in range(len(_HDEMUCS_SOURCES)) if s != v)
    # De-normalize with the std ONLY. The four HDEMUCS stems sum to the
    # normalized mix, so the vocals-complement maps back to real units via
    # *ref_std alone; adding ref_mean here would re-inject the whole mix's
    # DC offset onto a signal that already dropped the vocal component.
    instrumental = instrumental * ref_std  # → [2, T]@msr
    if msr == sample_rate:
        return instrumental
    return torchaudio.functional.resample(instrumental, msr, sample_rate)


def load_audio_ffmpeg(path: str, sample_rate: int = 16000) -> torch.Tensor:
    """Decode any audio format to mono float32 at `sample_rate` via ffmpeg.

    torchaudio dropped default mp3 backend in 2.9. ffmpeg → f32le PCM is
    portable and avoids extra Python deps.
    """
    import subprocess
    import numpy as np
    cmd = ["ffmpeg", "-nostdin", "-loglevel", "error",
           "-i", path,
           "-ar", str(sample_rate), "-ac", "1",
           "-f", "f32le", "-"]
    proc = subprocess.run(cmd, check=True, capture_output=True)
    arr = np.frombuffer(proc.stdout, dtype=np.float32).copy()
    return torch.from_numpy(arr).unsqueeze(0)  # [1, T]


def load_audio_stereo_ffmpeg(path: str, sample_rate: int) -> torch.Tensor:
    """Decode any audio format to **stereo** float32 at `sample_rate` via
    ffmpeg. Returns `[2, T]`. The karaoke instrumental path wants full
    stereo (the alignment path uses the mono `load_audio_ffmpeg`)."""
    import subprocess
    import numpy as np
    cmd = ["ffmpeg", "-nostdin", "-loglevel", "error",
           "-i", path,
           "-ar", str(sample_rate), "-ac", "2",
           "-f", "f32le", "-"]
    proc = subprocess.run(cmd, check=True, capture_output=True)
    arr = np.frombuffer(proc.stdout, dtype=np.float32).copy()
    # Interleaved L,R,L,R… → [2, T].
    return torch.from_numpy(arr).reshape(-1, 2).t().contiguous()


def _encode_aac(waveform: torch.Tensor, sample_rate: int, bitrate: str = "192k") -> bytes:
    """Encode a `[1, T]`/`[2, T]` float waveform to AAC/m4a bytes via
    ffmpeg. The summed instrumental can exceed [-1, 1], so clamp first.
    Goes through a temp file because the MP4 muxer needs a seekable output
    for its moov atom (it can't stream to a pipe)."""
    import subprocess
    import numpy as np
    wav = waveform.unsqueeze(0) if waveform.dim() == 1 else waveform
    ch = wav.size(0)
    # [ch, T] → interleaved [T, ch] f32le.
    data = wav.clamp(-1.0, 1.0).t().contiguous().detach().cpu().numpy().astype("float32")
    fd, out_path = tempfile.mkstemp(suffix=".m4a")
    os.close(fd)
    try:
        cmd = ["ffmpeg", "-nostdin", "-loglevel", "error", "-y",
               "-f", "f32le", "-ar", str(sample_rate), "-ac", str(ch), "-i", "pipe:0",
               "-c:a", "aac", "-b:a", bitrate, out_path]
        subprocess.run(cmd, input=data.tobytes(), check=True, capture_output=True)
        with open(out_path, "rb") as fh:
            return fh.read()
    finally:
        try:
            os.unlink(out_path)
        except OSError:
            pass


def render_instrumental(audio_path: str, device: str = "cpu", progress=None) -> dict:
    """Audio → instrumental ("минус") track for real karaoke mode.

    Decode → HDEMUCS → drop the vocal stem (keep drums+bass+other) →
    re-encode to AAC/m4a. Returns ``{"audio": <bytes>, "ext": "m4a",
    "duration": <sec>}``. Stereo at the separator's native rate, so there's
    no resample round-trip and the минус stays karaoke-quality. `progress`
    is forwarded to the separator for per-chunk job progress."""
    _, msr = load_separator(device)
    waveform = load_audio_stereo_ffmpeg(audio_path, msr)        # [2, T] @ msr
    # Reject empty / metadata-only decodes — a zero-length pass would encode
    # a valid-but-silent m4a and report a bogus "done" with a 0s track.
    if waveform.numel() == 0 or waveform.size(1) < int(0.1 * msr):
        raise ValueError("audio is empty or too short to separate")
    duration = waveform.size(1) / msr if msr else 0.0
    instrumental = separate_instrumental(waveform, msr, device, progress=progress)
    return {"audio": _encode_aac(instrumental, msr), "ext": "m4a", "duration": duration}


_MIN_NORM_STD = 0.01          # ≈ -40 dBFS: below this it isn't a vocal stem


def normalize_waveform(waveform: torch.Tensor) -> torch.Tensor:
    """Zero-mean / unit-variance, the way MMS was trained.

    `facebook/mms-1b-fl102`'s feature extractor sets `do_normalize=True`, but we
    never call the feature extractor — we feed the waveform straight to the
    model. Skipping the normalization makes the emissions *level-dependent*,
    which is exactly wrong here: the vocal stem HDEMUCS hands back is de-normal-
    ized to the mix's own scale, so a track with quiet vocals under a loud
    master arrives several dB down and the CTC posteriors turn to mush (measured:
    line-start MAE 2.5 s with a 12 s outlier, and the collapse guard doesn't
    even fire because the timings are wrong rather than identical). Normalizing
    once over the whole utterance — not per chunk, so every chunk keeps the same
    scale and the concatenated emissions stay comparable — makes alignment
    level-invariant.

    The gain is floored, not unbounded: a stem quieter than roughly -40 dBFS is
    not a vocal at all, it's separation residue, and stretching that to unit
    variance manufactures confident nonsense out of music bleed. Normal stems
    (measured std 0.04) are far above the floor and unaffected."""
    w = waveform.to(torch.float32)
    std = torch.sqrt(w.var() + 1e-7)
    return (w - w.mean()) / torch.clamp(std, min=_MIN_NORM_STD)


# wav2vec2's conv stack: 400-sample receptive field, 320-sample stride at
# 16 kHz → 50 frames/sec, frame f covering samples [f*320, f*320+400).
_CONV_STRIDE = 320
_CONV_WINDOW = 400


def get_emissions(model, waveform: torch.Tensor, device: str,
                  chunk_sec: float = 30.0, sample_rate: int = 16000,
                  context_sec: float = 2.0) -> Tuple[torch.Tensor, float]:
    """Run the model over the audio in overlapping chunks and concatenate the
    emissions into one continuous stream.

    Each chunk is widened by `context_sec` on both sides and the extra frames
    are thrown away, so a word straddling a chunk boundary is still scored with
    real audio either side of it. Cutting hard at 30 s (what we used to do) gave
    the transformer a truncated left/right context exactly there, which shows up
    as timing drift around every boundary — for a 3-minute song that's 5 chances
    to derail the global alignment path.

    Returns (emissions [frames, vocab], frames_per_second).
    """
    total = waveform.size(1)
    # Minimum input the conv stack can produce a frame from, with headroom.
    min_samples = max(_CONV_WINDOW * 2, sample_rate // 2)
    chunk_samples = max(min_samples, int(chunk_sec * sample_rate))
    ctx = max(0, int(context_sec * sample_rate))

    waveform = normalize_waveform(waveform)

    # Plan the output spans; fold a too-short tail into its predecessor.
    spans: List[Tuple[int, int]] = []
    offset = 0
    while offset < total:
        spans.append((offset, min(offset + chunk_samples, total)))
        offset = spans[-1][1]
    if len(spans) >= 2 and (spans[-1][1] - spans[-1][0]) < min_samples:
        spans = spans[:-2] + [(spans[-2][0], spans[-1][1])]

    parts = []
    with torch.inference_mode():
        for s, e in spans:
            ws = max(0, s - ctx)
            we = min(total, e + ctx)
            slice_ = waveform[:, ws:we].to(device)
            if slice_.size(1) < min_samples:
                # Song shorter than the conv minimum: pad with silence.
                pad = torch.zeros(slice_.size(0), min_samples - slice_.size(1),
                                  device=slice_.device, dtype=slice_.dtype)
                slice_ = torch.cat([slice_, pad], dim=1)
            logits = model(slice_).logits            # (1, frames, vocab)
            emis = logits.log_softmax(dim=-1).cpu()
            # Keep only the frames whose receptive-field centre lands inside
            # this chunk's own span; the rest were context for the neighbours.
            n = emis.size(1)
            centre0 = ws + _CONV_WINDOW // 2
            lo = 0 if ws == s else max(0, -(-(s - centre0) // _CONV_STRIDE))
            hi = n if we == e else min(n, max(lo, -(-(e - centre0) // _CONV_STRIDE)))
            parts.append(emis[:, lo:hi, :])
    emissions = torch.cat(parts, dim=1).squeeze(0)   # (frames, vocab)
    duration_sec = total / sample_rate
    fps = emissions.size(0) / duration_sec if duration_sec > 0 else 50.0
    return emissions, fps


# Transcription (audio → draft lyric text) uses a Kazakh-fine-tuned
# Whisper — `kk-turbo` = whisper-large-v3-turbo fine-tuned on KSC2
# (issai/abilmansplus), converted to CTranslate2/int8 (faster-whisper).
# Far better on sung Kazakh than the MMS CTC head we align with; reuses
# the same HDEMUCS vocal isolation. Model id overridable for local tests.
def _default_transcribe_model() -> str:
    """Where to load kk-turbo from, most specific first.

    ANKUI_KKTURBO_MODEL wins (the Space sets it to the path baked into the
    image). Failing that, prefer a locally converted copy: the Hub id is a
    *private* repo, so falling straight back to it means every local run — evals,
    dataset builds — dies on a 404 that reads like a bug in the code. Build one
    with convert_kkturbo.py pointed at ../models/kk-turbo-ct2.
    """
    env = os.environ.get("ANKUI_KKTURBO_MODEL")
    if env:
        return env
    local = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                         "models", "kk-turbo-ct2")
    if os.path.isdir(local) and os.path.exists(os.path.join(local, "model.bin")):
        return local
    return "italant7/whisper-turbo-ksc2-ct2"


_TRANSCRIBE_MODEL = _default_transcribe_model()
_TRANSCRIBER_CACHE = {}


def load_transcriber(device: str = "cpu"):
    """Load the kk-turbo faster-whisper model (cached). int8 on CPU,
    float16 on CUDA. CTranslate2 has no MPS path, so non-CUDA → CPU."""
    if device in _TRANSCRIBER_CACHE:
        return _TRANSCRIBER_CACHE[device]
    from faster_whisper import WhisperModel
    ct2_device = "cuda" if device == "cuda" else "cpu"
    compute = "float16" if ct2_device == "cuda" else "int8"
    # Match CTranslate2's CPU threads to our real core budget (see _THREADS).
    model = WhisperModel(_TRANSCRIBE_MODEL, device=ct2_device, compute_type=compute,
                         cpu_threads=(_THREADS if ct2_device == "cpu" else 0))
    _TRANSCRIBER_CACHE[device] = model
    return model


def transcribe_audio(audio_path: str, device: str = "cpu",
                     separate: bool = True) -> dict:
    """Audio → draft Kazakh lyric text with per-phrase line breaks.

    Isolate vocals (HDEMUCS, same as alignment) → kk-turbo Whisper with
    word_timestamps=True. Params are the song-tuned set from KazakhSTT: VAD
    on (skips instrumental intro/outro so it doesn't hallucinate "оооо"),
    beam 10, temperature 0 + condition_on_previous_text False.

    Whisper lumps whole verses into a few giant SEGMENTS, so the old
    "one segment == one line" join returned a handful of monster lines — the
    words are right, only the line breaks are wrong. We request word
    timestamps, flatten every word onto one timeline (so a pause that lands
    on a Whisper segment boundary is treated like any other pause), and
    re-split into singable phrase lines by breath gaps, with char/word/dur
    caps and an orphan-merge so we never emit one-word lines or re-form a
    giant one. It's a *draft* (force-aligned + hand-edited afterward), so
    sensible phrase breaks beat perfect timing. Any failure or missing word
    timestamps falls back to the original per-segment join — worst case is
    byte-identical to before.
    """
    # Phrase-splitting tunables (seconds / chars / counts).
    #
    # The gap threshold is derived from the song's own gap distribution, not
    # fixed. Fixed thresholds were the bug: this used to require 0.40 s to
    # consider a break and 0.65 s to force one, but Whisper's word timestamps
    # come out of a VAD-filtered stream where the silence is already removed —
    # measured on a 16-line song, *every* line boundary showed up as a 0.17–0.21 s
    # gap and every within-line gap as exactly 0.00. Nothing ever reached 0.40,
    # so no break was ever taken on a gap and lines were chopped by the character
    # cap instead: 16 sung lines came back as 7. Half the median non-zero gap
    # sits neatly between the two populations and returns 16 of 16.
    GAP_RATIO = 0.5     # of the median non-zero inter-word gap
    GAP_MIN = 0.08      # s: floor, so a gapless transcript doesn't split on noise
    GAP_MAX = 1.0       # s: ceiling, so one long break can't raise the bar
    GAP_FORCE = 0.45    # s: a pause this long ends a line regardless
    MAX_CHARS = 42      # backstop cap on a singable line's characters
    MIN_CHARS = 12      # below this a line can't be soft-broken
    MAX_WORDS = 9       # backstop cap on runs of tiny particles
    MIN_WORDS = 2       # a soft break may not orphan a single word
    MAX_DUR = 8.0       # s: force a break if a gapless line runs this long
    MERGE_FLOOR = 6     # chars: a fragment this small folds into its neighbour
    PUNCT_HARD = (".", "!", "?", "…")
    PUNCT_SOFT = (",", ";", ":", "—", "–")

    waveform = load_audio_ffmpeg(audio_path, sample_rate=16000)
    duration = waveform.size(1) / 16000
    wave = separate_vocals(waveform, 16000, device) if separate else waveform
    vocals = wave.squeeze(0).detach().cpu().numpy().astype("float32")

    model = load_transcriber(device)
    seg_iter, _info = model.transcribe(
        vocals, language="kk", beam_size=10, vad_filter=True,
        condition_on_previous_text=False,
        # Greedy, no temperature fallback. The hotter retries do recover more
        # words on degraded audio, but measurably the wrong ones: word error on
        # the sung fixture went from 0.76 to 0.93 while character error stayed
        # put. For a draft the user is going to correct by hand, fewer-and-right
        # beats more-and-wrong.
        temperature=0,
        # Mild — song lyrics legitimately repeat words, so this is set to
        # discourage runaway loops without penalizing a real refrain.
        repetition_penalty=1.05,
        # Instrumental stretches are where Whisper invents lyrics; with word
        # timestamps on it can detect and skip them.
        hallucination_silence_threshold=2.0,
        # Default 2 s of silence is longer than the gap between sung lines, and
        # the 400 ms pad smears the boundary we split on.
        vad_parameters=dict(min_silence_duration_ms=500, speech_pad_ms=200),
        word_timestamps=True,
    )
    # Materialize once: seg_iter is a single-use generator and both the
    # phrase-split path and the fallback path need the same segments.
    segments = list(seg_iter)

    def _phrase_lines(segs):
        """Flatten word timestamps onto one timeline and re-split into phrase
        lines. Returns `(groups, words)` — groups are index lists into `words`,
        which is the single source of truth for the flattened timeline — or
        `([], [])` when there are no usable word timestamps."""
        words = []
        for s in segs:
            for w in (getattr(s, "words", None) or []):
                txt = (getattr(w, "word", "") or "").strip()
                if not txt:
                    continue
                st = getattr(w, "start", None)
                en = getattr(w, "end", None)
                if st is None and en is None:
                    continue
                st = float(st) if st is not None else float(en)
                en = float(en) if en is not None else st
                words.append((txt, st, en))
        if not words:
            return [], []

        # Chronological order, always. Words are collected in *segment* order,
        # and a segment whose timestamps run backwards past its predecessor
        # (which happens when a decode is retried, or a VAD region is re-cut)
        # would otherwise emit lyric lines out of sung order — the user-visible
        # "the rows came back swapped". Karaoke can only ever be chronological,
        # so sorting here costs nothing and removes the whole failure class.
        words.sort(key=lambda w: (w[1], w[2]))

        # Collapse a degenerate repeat loop: the same word three or more times
        # in a row is Whisper spinning on silence ("қазақстан қазақстан
        # қазақстан" off the end of a track), not a lyric. Trimmed to two, not
        # one, because doubling is genuinely sung — "Сүйікті етші, сүйікті
        # етші, сүйіктім" must survive intact.
        deduped, run = [], 0
        for w in words:
            same = deduped and w[0].lower() == deduped[-1][0].lower()
            run = run + 1 if same else 0
            if run < 2:
                deduped.append(w)
        words = deduped

        # Break threshold from this song's own gap distribution. Within-phrase
        # gaps cluster at ~0, phrase boundaries an order of magnitude above;
        # half the median non-zero gap lands between the two for both a fast
        # rap and a slow ballad, which a fixed number cannot.
        nonzero = sorted(
            g for g in (words[i][1] - words[i - 1][2] for i in range(1, len(words)))
            if g > 0.02
        )
        gap_bar = (min(GAP_MAX, max(GAP_MIN, GAP_RATIO * nonzero[len(nonzero) // 2]))
                   if nonzero else GAP_MIN)

        groups: List[List[int]] = []
        cur: List[int] = []
        cur_chars = 0
        cur_start = None
        n = len(words)
        for i, (txt, st, en) in enumerate(words):
            if not cur:
                cur_start = st
            cur.append(i)
            cur_chars += len(txt) + (1 if len(cur) > 1 else 0)
            line_dur = en - cur_start
            gap = (words[i + 1][1] - en) if i + 1 < n else float("inf")
            force = (
                gap >= max(gap_bar, GAP_FORCE)
                or i + 1 == n                 # end of stream always closes
                or len(cur) >= MAX_WORDS
                or cur_chars >= MAX_CHARS
                or line_dur >= MAX_DUR
                or txt.endswith(PUNCT_HARD)
            )
            prefer = (
                (gap >= gap_bar or txt.endswith(PUNCT_SOFT))
                and cur_chars >= MIN_CHARS
                and len(cur) >= MIN_WORDS
            )
            if force or prefer:
                groups.append(cur)
                cur = []
                cur_chars = 0
        if cur:
            groups.append(cur)

        # Fold tiny orphan fragments into the previous line.
        merged: List[List[int]] = []
        for g in groups:
            if not g:
                continue
            if merged and sum(len(words[i][0]) for i in g) <= MERGE_FLOOR:
                merged[-1].extend(g)
            else:
                merged.append(g)
        return merged, words

    try:
        groups, flat_words = _phrase_lines(segments)
    except Exception as exc:  # noqa: BLE001 — any failure → safe fallback
        sys.stderr.write(f"[align] phrase-split failed, segment fallback: {exc}\n")
        groups, flat_words = [], []

    timed: List[dict] = []
    for g in groups:
        ws = [flat_words[i] for i in g]
        if not ws:
            continue
        timed.append({
            "startMs": int(ws[0][1] * 1000),
            "endMs": int(ws[-1][2] * 1000),
            "text": " ".join(w[0] for w in ws),
            "words": [{"wordIdx": k, "startMs": int(w[1] * 1000),
                       "endMs": int(w[2] * 1000), "text": w[0]}
                      for k, w in enumerate(ws)],
        })

    lines = [t["text"] for t in timed]
    if not lines:
        # Original behaviour: one Whisper segment per line, no timings.
        lines = [s.text.strip() for s in segments if s.text.strip()]
        timed = []

    try:
        from truecase import truecase as _truecase
        lines = [_truecase(ln) for ln in lines]
        for t, ln in zip(timed, lines):
            t["text"] = ln
            # Push the same casing down to the words, or a client rendering
            # `lines[].words[]` shows lowercase text under a truecased line.
            parts = ln.split()
            if len(parts) == len(t.get("words") or []):
                for w, p in zip(t["words"], parts):
                    w["text"] = p
    except Exception as exc:  # noqa: BLE001
        sys.stderr.write(f"[align] truecase skipped: {exc}\n")

    text = "\n".join(lines).strip()
    # `lines` carries Whisper's own word timings for the draft. They're rough —
    # good enough to show the user where each line sits, and to give /align a
    # sanity reference — but the karaoke timings still come from forced
    # alignment, which is an order of magnitude tighter.
    return {"duration": duration, "text": text, "lines": timed,
            "model": "kk-turbo-ksc2"}


def tokenize(processor, text: str) -> torch.Tensor:
    """Encode normalized Kazakh text to token IDs that the model knows."""
    ids = processor.tokenizer(text, return_tensors="pt").input_ids[0]
    # The HF tokenizer wraps with pad/sos sometimes — strip pad.
    pad = processor.tokenizer.pad_token_id
    ids = ids[ids != pad]
    return ids


def merge_token_spans(alignments: torch.Tensor, scores: torch.Tensor,
                      blank: int = 0) -> List[dict]:
    """Compress per-frame token output into [token, start_frame, end_frame, score].

    forced_align emits one token id per frame (with blanks). Consecutive
    frames sharing the same non-blank token form a span.

    `blank` must be the same id passed to `forced_align`. It used to be
    hard-coded to 0, which happens to be right for MMS but silently produces
    garbage for any CTC head whose pad token sits elsewhere in the vocab.
    """
    out: List[dict] = []
    cur_id = None
    cur_start = 0
    cur_scores: List[float] = []
    for i, tok in enumerate(alignments.tolist()):
        if tok == blank:  # blank — close current span if any
            if cur_id is not None:
                out.append({"id": cur_id, "start": cur_start, "end": i,
                             "score": float(sum(cur_scores) / max(1, len(cur_scores)))})
                cur_id = None
                cur_scores = []
            continue
        if tok != cur_id:
            if cur_id is not None:
                out.append({"id": cur_id, "start": cur_start, "end": i,
                             "score": float(sum(cur_scores) / max(1, len(cur_scores)))})
            cur_id = tok
            cur_start = i
            cur_scores = [float(scores[i])]
        else:
            cur_scores.append(float(scores[i]))
    if cur_id is not None:
        out.append({"id": cur_id, "start": cur_start, "end": len(alignments),
                     "score": float(sum(cur_scores) / max(1, len(cur_scores)))})
    return out


def build_targets(ref_lines: List[str], line_idxs: List[int], vocab: dict,
                  space_id: int) -> Tuple[List[int], List[tuple], dict]:
    """Encode the reference as a CTC target sequence.

    Returns `(target_ids, provenance, unit_text)` where `provenance[i]` is
    `(line_idx, unit_idx, kind, char)` for target token `i`, and
    `unit_text[(line_idx, unit_idx)]` is the *original* word to display.

    Normalization runs per original word rather than per line. The old code
    normalized the whole line, aligned the normalized words, then paired them
    back to `line.split()` by position — which silently skews as soon as
    normalization changes the word count, and it always does: "асты — кең"
    splits into three whitespace words but two spoken ones, so every word after
    the dash carried its neighbour's timing. Normalizing per word keeps the
    mapping exact by construction, and a word that normalizes away entirely
    (a lone dash, "(2x)") just contributes no tokens.
    """
    target_ids: List[int] = []
    prov: List[tuple] = []
    unit_text: dict = {}
    for li in line_idxs:
        line = ref_lines[li]
        unit = 0
        for orig_word in line.split():
            norm = normalize_for_align(orig_word, vocab)
            if not norm:
                continue                      # punctuation-only "word"
            if target_ids:
                target_ids.append(space_id)
                prov.append((li, unit, "sep", " "))
            for ch in norm:
                if ch == " ":
                    continue                  # one word, spelled straight through
                target_ids.append(vocab[ch])
                prov.append((li, unit, "char", ch))
            unit_text[(li, unit)] = orig_word
            unit += 1
    return target_ids, prov, unit_text


def _lines_from_alignment(ref_lines: List[str], prov: List[tuple], unit_text: dict,
                          spans: List[dict], sec_per_frame: float) -> List[dict]:
    """Group token spans into word units and then lines, in reference order."""
    units: dict = {}
    for tok_idx, p in enumerate(prov):
        if p[2] != "char" or tok_idx >= len(spans):
            continue
        sp = spans[tok_idx]
        key = (p[0], p[1])
        u = units.get(key)
        if u is None:
            units[key] = {"start": sp["start"], "end": sp["end"], "scores": [sp["score"]]}
        else:
            u["start"] = min(u["start"], sp["start"])
            u["end"] = max(u["end"], sp["end"])
            u["scores"].append(sp["score"])

    by_line: dict = {}
    for (li, ui), u in sorted(units.items()):
        by_line.setdefault(li, []).append((ui, u))

    out: List[dict] = []
    for li in sorted(by_line):
        words = []
        for ui, u in by_line[li]:
            s = u["start"] * sec_per_frame
            e = max(s + 0.05, u["end"] * sec_per_frame)
            words.append({
                "wordIdx": len(words),
                "startMs": int(s * 1000),
                "endMs": int(e * 1000),
                "text": unit_text.get((li, ui), ""),
                "score": sum(u["scores"]) / max(1, len(u["scores"])),
            })
        if not words:
            continue
        out.append({
            "lineIdx": li,
            "startMs": words[0]["startMs"],
            "endMs": words[-1]["endMs"],
            "text": ref_lines[li],
            "score": sum(w["score"] for w in words) / len(words),
            "words": words,
        })
    return out


# Rejecting a low-confidence alignment and substituting the energy spread is OFF
# by default, because measurement says it makes things worse: on a vocal 10 dB
# under the backing (mean confidence 0.03, i.e. as bad as it gets) the real CTC
# alignment still scored 4.0 s mean line error while the energy spread scored
# 13.2 s. A hard-won alignment beats an evenly-distributed guess even when the
# model isn't sure — so we keep it and report `meanConfidence` instead, letting
# the client tell the user which lines to check. Raise ANKUI_ALIGN_MIN_CONF above
# 0 only to deliberately trade accuracy for "never look confident".
_MIN_MEAN_CONFIDENCE = float(os.environ.get("ANKUI_ALIGN_MIN_CONF", "0"))


def _confidence(score: float) -> float:
    """forced_align hands back per-frame log-probabilities of the chosen token;
    exponentiating gives a plain 0…1 'how sure was the model' number that is
    safe to show a user and to threshold on."""
    import math
    return max(0.0, min(1.0, math.exp(score)))


def _postprocess(lines: List[dict], duration: float, hold_sec: float = 0.6) -> List[dict]:
    """Make the timeline safe to drive a karaoke UI from.

    Forced alignment is monotonic in *token* order but the derived line times
    can still be degenerate at the edges — a zero-length line, an end past the
    end of the file, or a line whose start equals the previous line's start.
    Also extends each line's end toward the next line's start (bounded), because
    a line that closes on its final consonant flickers off mid-note.
    """
    dur_ms = int(duration * 1000)
    out = [dict(l) for l in lines]
    prev_end = 0
    for i, l in enumerate(out):
        start = max(0, min(int(l["startMs"]), dur_ms))
        start = max(start, prev_end - 200 if i else 0)   # allow a little overlap
        end = max(start + 200, min(int(l["endMs"]), dur_ms))
        nxt = int(out[i + 1]["startMs"]) if i + 1 < len(out) else dur_ms
        if nxt > end:
            end = min(max(end, min(nxt, end + int(hold_sec * 1000))), dur_ms)
        l["startMs"], l["endMs"] = start, end
        l.setdefault("confidence", round(_confidence(l.get("score", 0.0)), 3))
        prev_end = end
        words = l.get("words") or []
        wprev = start
        for w in words:
            ws = max(start, min(int(w["startMs"]), end))
            ws = max(ws, wprev)
            we = max(ws + 50, min(int(w["endMs"]), end))
            w["startMs"], w["endMs"] = ws, we
            # setdefault, mirroring the line above: a caller that already knows
            # the confidence (the energy fallback) has the final say.
            w.setdefault("confidence", round(_confidence(w.get("score", 0.0)), 3))
            # Floor the next word at this word's END, not its start, so adjacent
            # word spans cannot overlap and light up two "current" words at once.
            wprev = we
    return out


def prepare_alignment(audio_path: str, device: str = "cpu",
                      separate: bool = True) -> dict:
    """Do everything that depends only on the *audio*: decode, isolate vocals,
    run the acoustic model.

    Split out from `run_alignment` so one audio pass can serve many references.
    That is what makes "fetch several candidate lyric sheets and keep whichever
    actually matches the recording" affordable: separation plus the MMS forward
    pass is essentially the entire cost, while `forced_align` over an existing
    emission matrix is milliseconds. Verifying five candidates this way costs
    barely more than aligning one.
    """
    proc, model = load_model(device)
    sample_rate = proc.feature_extractor.sampling_rate

    # Audio: decode straight to 16 kHz mono float via ffmpeg (no torchaudio
    # backend dance, works with mp3/m4a/flac/anything ffmpeg knows).
    waveform = load_audio_ffmpeg(audio_path, sample_rate=sample_rate)
    duration = waveform.size(1) / sample_rate

    # Isolate vocals before alignment so a real musical mix aligns to the
    # words instead of collapsing (clean acapella/TTS passes through fine).
    align_wave = separate_vocals(waveform, sample_rate, device) if separate else waveform

    # Run model — chunked (with context overlap) to avoid OOM on long songs.
    emissions, fps = get_emissions(model, align_wave, device, chunk_sec=30.0,
                                   sample_rate=sample_rate)
    return {"proc": proc, "sample_rate": sample_rate, "duration": duration,
            "align_wave": align_wave, "emissions": emissions, "fps": fps}


def align_reference(ctx: dict, ref_lines: List[str]) -> dict:
    """Align `ref_lines` against a context from `prepare_alignment`."""
    proc = ctx["proc"]
    sample_rate = ctx["sample_rate"]
    duration = ctx["duration"]
    align_wave = ctx["align_wave"]
    emissions = ctx["emissions"]
    fps = ctx["fps"]

    vocab = proc.tokenizer.get_vocab()
    space_id = vocab.get("|")            # MMS uses '|' as the word separator
    pad_id = proc.tokenizer.pad_token_id  # == the CTC blank
    # Only single characters can be targets; the specials (<s>, <unk>, …) must
    # never end up in the reference — see normalize_for_align's vocab filter.
    char_vocab = {t: i for t, i in vocab.items() if len(t) == 1}

    sec_per_frame = 1.0 / fps
    emissions_b = emissions.unsqueeze(0)

    def align_subset(line_idxs: List[int]):
        target_ids, prov, unit_text = build_targets(ref_lines, line_idxs,
                                                    char_vocab, space_id)
        if not target_ids:
            return None
        targets = torch.tensor([target_ids], dtype=torch.int32)
        alignments, scores = torchaudio.functional.forced_align(
            emissions_b, targets, blank=pad_id)
        spans = merge_token_spans(alignments[0], scores[0], blank=pad_id)
        return _lines_from_alignment(ref_lines, prov, unit_text, spans, sec_per_frame)

    lines_out = align_subset(list(range(len(ref_lines))))
    if lines_out is None:
        return {"duration": duration, "lines": [], "words": [],
                "fps": fps, "alignMode": "no-reference"}

    align_mode = "ctc-mms-fl102-kaz"
    # Reported to the client so it can tell the user which lines to check. A
    # clean alignment measures ~0.7, a hard sung mix ~0.3, and a vocal buried
    # under the backing ~0.03.
    mean_conf = (sum(_confidence(l["score"]) for l in lines_out) / len(lines_out)
                 if lines_out else 0.0)

    # Guard against the classic collapse (every line at 0:00). Note this is a
    # *degeneracy* test, not a quality one: rejecting merely low-confidence
    # alignments and substituting the energy spread measured worse than keeping
    # them (see _MIN_MEAN_CONFIDENCE), so quality is reported, not acted on.
    if _is_degenerate(lines_out, duration) or mean_conf < _MIN_MEAN_CONFIDENCE:
        lines_out = _distribute_by_energy(ref_lines, align_wave, sample_rate, duration)
        align_mode = "energy-fallback"

    if align_mode == "energy-fallback":
        # The CTC mean describes an alignment that is no longer in the payload.
        mean_conf = 0.0
    lines_out = _postprocess(lines_out, duration)
    words_out = [dict(w, lineIdx=l["lineIdx"]) for l in lines_out
                 for w in (l.get("words") or [])]

    return {
        "duration": duration,
        "lines":    lines_out,
        "words":    words_out,
        "fps":      fps,
        "alignMode": align_mode,
        "meanConfidence": round(mean_conf, 3),
    }


def run_alignment(audio_path: str, ref_lines: List[str], device: str = "cpu",
                  separate: bool = True) -> dict:
    """Full pipeline: audio + reference lyric -> timed lines. Unchanged API."""
    ctx = prepare_alignment(audio_path, device=device, separate=separate)
    return align_reference(ctx, ref_lines)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--url", required=True, help="MP3 URL")
    p.add_argument("--text-file", required=True, help="reference lyric text path")
    p.add_argument("--out", default="-", help="output JSON path or '-' for stdout")
    p.add_argument("--device", default="cpu", choices=["cpu", "mps", "cuda"])
    p.add_argument("--keep-audio", action="store_true")
    # Compatibility no-ops (so the Go driver's --model and --align-mode flags don't break).
    p.add_argument("--model", default="ignored")
    p.add_argument("--align-mode", default="ignored")
    args = p.parse_args()

    with open(args.text_file, encoding="utf-8") as f:
        raw = f.read().replace("\\n", "\n").replace("\\r", "")
    ref_lines = split_lines(raw)
    if not ref_lines:
        print(json.dumps({"error": "no usable reference lines"}))
        sys.exit(1)

    fd, audio_path = tempfile.mkstemp(suffix=".mp3")
    os.close(fd)
    try:
        t0 = time.time()
        download(args.url, audio_path)
        dl_s = time.time() - t0

        t0 = time.time()
        result = run_alignment(audio_path, ref_lines, device=args.device)
        align_s = time.time() - t0
    finally:
        if not args.keep_audio:
            try: os.unlink(audio_path)
            except OSError: pass

    payload = {
        "duration":         result["duration"],
        # Report what actually ran — this can come back "energy-fallback".
        "alignMode":        result.get("alignMode", "ctc-mms-fl102-kaz"),
        "meanConfidence":   result.get("meanConfidence", 0.0),
        "downloadSeconds":  round(dl_s, 2),
        "alignSeconds":     round(align_s, 2),
        "transcribeSeconds": 0.0,  # Whisper no longer in the loop
        "fps":              round(result.get("fps", 50.0), 2),
        "refLines":         len(ref_lines),
        "lines":            result["lines"],
    }
    out = json.dumps(payload, ensure_ascii=False, indent=2)
    if args.out == "-":
        print(out)
    else:
        with open(args.out, "w", encoding="utf-8") as f:
            f.write(out)


if __name__ == "__main__":
    main()