"""
ComfyUI-ForcedAlignTimings — word-level forced alignment for TTS markup.

Aligns generated speech audio against a KNOWN text using torchaudio's
MMS_FA (wav2vec2 multilingual forced aligner) and returns per-word timings +
character offsets into the text we aligned against.

Pipeline role (Synthy/iRead "boris_tts_*" provider PoC):
    CosyVoice (text -> audio)
        -> ForcedAlignTimings (audio + known text -> raw word timings)
            -> n8n Code node (raw timings + ORIGINAL text -> X-Voice-Markup)

This node deliberately does NOT compute UTF-16 offsets and does NOT map text
normalization back to the original — that is pure string logic and lives in the
n8n Code node (see 1/tts-poc/map_timings.mjs). Here we only emit, per word:
    value, char_start, char_end (Python codepoint offsets), time(ms), end(ms), score

Author: iMind Studio
License: MIT
"""

import json
import re

import torch
import torchaudio


# ── word tokenizer ───────────────────────────────────
# One token per run of letters (with internal apostrophes) OR run of digits.
# Surrounding punctuation/space is excluded, so char_start/char_end point at the
# word core only (matches SY-A-39: commas / "!" are not emitted as marks).
_WORD_RE = re.compile(r"[0-9]+|[^\W\d_]+(?:['’][^\W\d_]+)*", re.UNICODE)

# Compact Cyrillic -> Latin fallback, used only if the `uroman` package is absent.
_RU_MAP = {
    "а": "a", "б": "b", "в": "v", "г": "g", "д": "d", "е": "e", "ё": "yo",
    "ж": "zh", "з": "z", "и": "i", "й": "y", "к": "k", "л": "l", "м": "m",
    "н": "n", "о": "o", "п": "p", "р": "r", "с": "s", "т": "t", "у": "u",
    "ф": "f", "х": "kh", "ц": "ts", "ч": "ch", "ш": "sh", "щ": "shch",
    "ъ": "", "ы": "y", "ь": "", "э": "e", "ю": "yu", "я": "ya",
}

_uroman = None
_uroman_tried = False


def _romanize(word):
    """Romanize a single word to lowercase [a-z'] for the MMS_FA dictionary."""
    global _uroman, _uroman_tried
    if not _uroman_tried:
        _uroman_tried = True
        try:
            import uroman as _ur
            _uroman = _ur.Uroman()
        except Exception:
            _uroman = None

    if _uroman is not None:
        try:
            text = _uroman.romanize_string(word)
        except Exception:
            text = word
    else:
        out = []
        for ch in word:
            low = ch.lower()
            out.append(_RU_MAP[low] if low in _RU_MAP else ch)
        text = "".join(out)

    text = text.lower()
    text = re.sub(r"[^a-z']", "", text)
    return text


class ForcedAlignTimings:
    """Forced-align an AUDIO clip against a KNOWN text; return raw word timings."""

    TARGET_SR = 16000  # MMS_FA expects 16 kHz mono

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "audio": ("AUDIO",),
                "text": ("STRING", {"default": "", "multiline": True}),
            },
            "optional": {
                "device": (["auto", "cuda", "cpu"],),
            },
        }

    RETURN_TYPES = ("STRING", "INT", "INT")
    RETURN_NAMES = ("timings_json", "audio_ms", "word_count")
    FUNCTION = "align"
    CATEGORY = "audio/tts"
    DESCRIPTION = ("Forced alignment (MMS_FA) of audio against known text. "
                   "Outputs per-word timings + char offsets as JSON for TTS markup.")
    OUTPUT_NODE = True  # surface timings_json into /history outputs so n8n can read it via API

    def _wrap(self, timings_json, audio_ms, count):
        # Downstream sockets receive `result`; n8n reads `ui` from /history[pid].outputs[nodeId].
        return {
            "ui": {"timings_json": [timings_json], "audio_ms": [audio_ms], "word_count": [count]},
            "result": (timings_json, audio_ms, count),
        }

    # ── helpers ──────────────────────────────────────
    def _prep_audio(self, audio):
        """ComfyUI AUDIO dict -> mono 16 kHz [1, T] tensor."""
        wav = audio["waveform"]            # [B, C, T]
        sr = int(audio["sample_rate"])
        if wav.dim() == 3:
            wav = wav[0]                   # [C, T]
        if wav.dim() == 2 and wav.shape[0] > 1:
            wav = wav.mean(dim=0, keepdim=True)
        elif wav.dim() == 1:
            wav = wav.unsqueeze(0)
        if sr != self.TARGET_SR:
            wav = torchaudio.functional.resample(wav, sr, self.TARGET_SR)
        return wav.contiguous(), self.TARGET_SR

    @staticmethod
    def _get_dict(bundle):
        try:
            return bundle.get_dict(star="*")
        except TypeError:
            return bundle.get_dict()

    @staticmethod
    def _get_model(bundle, with_star):
        try:
            return bundle.get_model(with_star=with_star)
        except TypeError:
            return bundle.get_model()

    # ── main ─────────────────────────────────────────
    def align(self, audio, text, device="auto"):
        from torchaudio.pipelines import MMS_FA as bundle

        dev = ("cuda" if torch.cuda.is_available() else "cpu") if device == "auto" else device

        waveform, sr = self._prep_audio(audio)
        waveform = waveform.to(dev)
        audio_ms = int(round(waveform.shape[-1] / sr * 1000))

        # 1) tokenize text into words with codepoint char offsets
        words = [{"value": m.group(0), "char_start": m.start(), "char_end": m.end()}
                 for m in _WORD_RE.finditer(text)]

        if not words:
            payload = {"sample_rate": sr, "audio_ms": audio_ms, "align_text": text, "words": []}
            return self._wrap(json.dumps(payload, ensure_ascii=False), audio_ms, 0)

        # 2) romanize words -> MMS_FA dictionary token ids (digits/unknown -> '*')
        dictionary = self._get_dict(bundle)
        star = dictionary.get("*")

        tokens_flat, counts = [], []
        for w in words:
            roman = _romanize(w["value"])
            ids = [dictionary[c] for c in roman if c in dictionary]
            if not ids:
                ids = [star if star is not None else 0]
            tokens_flat.extend(ids)
            counts.append(len(ids))

        # 3) emission + forced alignment
        model = self._get_model(bundle, with_star=(star is not None)).to(dev)
        with torch.inference_mode():
            emission, _ = model(waveform)
        targets = torch.tensor([tokens_flat], dtype=torch.int32, device=dev)
        aligned, scores = torchaudio.functional.forced_align(emission, targets, blank=0)
        token_spans = torchaudio.functional.merge_tokens(aligned[0], scores[0].exp())

        num_frames = emission.shape[1]
        ratio = waveform.shape[-1] / num_frames / sr   # seconds per emission frame

        # 4) regroup token spans into words by token counts
        out_words = []
        idx = 0
        for w, n in zip(words, counts):
            span = token_spans[idx: idx + n]
            idx += n
            if span:
                t0 = int(round(span[0].start * ratio * 1000))
                t1 = int(round(span[-1].end * ratio * 1000))
                sc = sum(s.score for s in span) / len(span)
            else:
                t0 = t1 = 0
                sc = 0.0
            out_words.append({
                "value": w["value"],
                "char_start": w["char_start"],
                "char_end": w["char_end"],
                "time": t0,
                "end": t1,
                "score": round(float(sc), 4),
            })

        payload = {"sample_rate": sr, "audio_ms": audio_ms, "align_text": text, "words": out_words}
        return self._wrap(json.dumps(payload, ensure_ascii=False), audio_ms, len(out_words))


NODE_CLASS_MAPPINGS = {"ForcedAlignTimings": ForcedAlignTimings}
NODE_DISPLAY_NAME_MAPPINGS = {"ForcedAlignTimings": "Forced Align Timings 🎯⏱️"}


# ── GPU smoke test (run on the 5090 box BEFORE docker commit) ──
# Usage:  python nodes.py path/to/audio.wav "Привет, мир"
if __name__ == "__main__":
    import sys

    wav_path, in_text = sys.argv[1], sys.argv[2]
    wav, sr = torchaudio.load(wav_path)
    audio_in = {"waveform": wav.unsqueeze(0), "sample_rate": sr}
    out = ForcedAlignTimings().align(audio_in, in_text)
    print(out["result"][0])