"""Local ONNX Russian stress annotator adapted from RUAccent.

This is a deliberately small, self-contained adaptation of RUAccent's MIT
licensed runtime.  Model weights and dictionaries are stored below
``release/ruaccent`` so using TeraTTS never downloads a second model repo.
"""

from __future__ import annotations

import gzip
import json
import re
from pathlib import Path

import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer


def _fix_capital(source: str, target: str) -> str:
    if len(source) != len(target):
        return target
    return "".join(b.upper() if a.isupper() else b.lower() for a, b in zip(source, target))


def _softmax(values: np.ndarray, axis: int = -1) -> np.ndarray:
    values = values - np.max(values, axis=axis, keepdims=True)
    exponent = np.exp(values)
    return exponent / exponent.sum(axis=axis, keepdims=True)


def _session(path: Path, device: str) -> ort.InferenceSession:
    provider = "CUDAExecutionProvider" if device.upper() == "CUDA" else "CPUExecutionProvider"
    return ort.InferenceSession(str(path / "model.onnx"), providers=[provider])


def _session_inputs(session: ort.InferenceSession, values: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
    """Pass only inputs that a particular exported graph declares."""
    required = {input_.name for input_ in session.get_inputs()}
    return {name: value for name, value in values.items() if name in required}


class _CharAccentModel:
    def __init__(self, path: Path, device: str) -> None:
        self.session = _session(path, device)
        self.id2label = json.loads((path / "config.json").read_text())["id2label"]
        self.vocab = {
            token.rstrip("\n"): index
            for index, token in enumerate((path / "vocab.txt").read_text().splitlines())
        }
        self.pad = self.vocab["[pad]"]
        self.unknown = self.vocab["[unk]"]
        self.bos = self.vocab["[bos]"]
        self.eos = self.vocab["[eos]"]

    def put_accent(self, word: str) -> str:
        ids = [self.bos] + [self.vocab.get(char, self.unknown) for char in word.lower()] + [self.eos]
        ids_array = np.asarray([ids], dtype=np.int64)
        inputs = _session_inputs(
            self.session,
            {
                "input_ids": ids_array,
                "attention_mask": np.ones_like(ids_array),
                "token_type_ids": np.zeros_like(ids_array),
            },
        )
        logits = self.session.run(None, inputs)[0]
        scores = _softmax(logits)[0]
        rendered = list(word)
        # The character model has [bos] at position zero, so character i is
        # prediction i+1.  Do not add a second marker to an explicit stress.
        for position, (label, score) in enumerate(zip(np.argmax(scores, axis=-1), np.max(scores, axis=-1))):
            character = position - 1
            if not 0 <= character < len(rendered):
                continue
            name = self.id2label[str(int(label))]
            if name not in {"NO", "STRESS_SECONDARY"} and float(score) >= 0.55:
                rendered[character] = "+" + rendered[character]
        return "".join(rendered)


class _TokenClassifier:
    def __init__(self, path: Path, device: str) -> None:
        self.session = _session(path, device)
        self.id2label = json.loads((path / "config.json").read_text())["id2label"]
        self.tokenizer = AutoTokenizer.from_pretrained(path, local_files_only=True)

    def classify_words(self, text: str) -> list[dict[str, object]]:
        encoded = self.tokenizer(
            text,
            return_offsets_mapping=True,
            return_special_tokens_mask=True,
            return_tensors="np",
        )
        offsets = encoded.pop("offset_mapping")[0]
        special = encoded.pop("special_tokens_mask")[0]
        token_ids = encoded["input_ids"][0]
        inputs = _session_inputs(
            self.session, {name: value.astype(np.int64) for name, value in encoded.items()}
        )
        scores = _softmax(self.session.run(None, inputs)[0])[0]
        pieces: list[dict[str, object]] = []
        for index, token_scores in enumerate(scores):
            if special[index]:
                continue
            start, end = (int(value) for value in offsets[index])
            token = self.tokenizer.convert_ids_to_tokens(int(token_ids[index]))
            reference = text[start:end]
            # Fast tokenizers use a prefix for continuation pieces.  The
            # offset fallback also works for WordPiece tokenizers.
            prefix = getattr(self.tokenizer._tokenizer.model, "continuing_subword_prefix", None)
            subword = len(token) != len(reference) if prefix else (start > 0 and text[start - 1 : start] != " ")
            if int(token_ids[index]) == self.tokenizer.unk_token_id:
                token, subword = reference, False
            pieces.append(
                {
                    "token": token,
                    "scores": token_scores,
                    "start": start,
                    "end": end,
                    "subword": subword,
                }
            )

        groups: list[list[dict[str, object]]] = []
        for piece in pieces:
            if groups and bool(piece["subword"]):
                groups[-1].append(piece)
            else:
                groups.append([piece])
        output = []
        for group in groups:
            averaged = np.mean(np.stack([item["scores"] for item in group]), axis=0)
            label = int(np.argmax(averaged))
            output.append(
                {
                    "entity": self.id2label[str(label)],
                    "score": float(averaged[label]),
                    "word": self.tokenizer.convert_tokens_to_string([str(item["token"]) for item in group]),
                    "start": group[0]["start"],
                    "end": group[-1]["end"],
                }
            )
        return output


class _OmographModel:
    def __init__(self, path: Path, device: str) -> None:
        self.session = _session(path, device)
        self.tokenizer = AutoTokenizer.from_pretrained(path, local_files_only=True)

    def choose(self, sentence: str, variants: list[str]) -> str:
        prepared = re.sub(r"\s+(?=(?:[,.?!:;…]))", "", sentence)
        encoded = self.tokenizer(
            [prepared] * len(variants), variants, max_length=512, truncation=True, padding=True, return_tensors="np"
        )
        inputs = _session_inputs(
            self.session, {name: value.astype(np.int64) for name, value in encoded.items()}
        )
        probabilities = _softmax(self.session.run(None, inputs)[0], axis=-1)
        return variants[int(np.argmax(probabilities[:, 1]))]


class RUAccent:
    """Local, ONNX-only RUAccent-compatible ``process_all`` implementation."""

    _normalize = re.compile(r"[^a-zA-Z0-9\sа-яА-ЯёЁ—.,!?:;'(){}\[\]«»„“”\-]")
    _tokens = re.compile(r"\w*(?:\+\w+)*|[^\w\s]+")
    _sentence = re.compile(r"[^.!?…]+[.!?…]*[\"»“]*")

    def __init__(
        self,
        root: Path,
        *,
        model_size: str = "turbo3.1",
        device: str = "CPU",
        mode: str = "full",
    ) -> None:
        root = root.resolve()
        if model_size != "turbo3.1":
            raise ValueError("the bundled RUAccent model is turbo3.1")
        if mode not in {"full", "dictionary"}:
            raise ValueError("RUAccent mode must be 'full' or 'dictionary'")
        self.mode = mode
        dictionary = root / "dictionary"
        self.accents = json.load(gzip.open(dictionary / "accents.json.gz"))
        self.omographs = json.load(gzip.open(dictionary / "omographs.json.gz"))
        self.omographs["коса"] = ["к+оса", "кос+а"]
        self.yo_words = json.load(gzip.open(dictionary / "yo_words.json.gz"))
        self.yo_homographs = json.load(gzip.open(dictionary / "yo_homographs.json.gz"))
        self.accents.update({"о": "+о", "О": "+О"})
        if mode == "full":
            self.accent_model = _CharAccentModel(root / "nn" / "nn_accent", device)
            self.omograph_model = _OmographModel(root / "nn" / "nn_omograph" / model_size, device)
            self.stress_usage = _TokenClassifier(root / "nn" / "nn_stress_usage_predictor", device)
            self.yo_classifier = _TokenClassifier(root / "nn" / "nn_yo_homograph_resolver", device)

    @staticmethod
    def _remaining(sentence: str, matches: list[re.Match[str]]) -> tuple[list[str], list[str]]:
        words = [match.group(0) for match in matches if match.group(0)]
        valid = [match for match in matches if match.group(0)]
        if not valid:
            return [], [sentence]
        gaps = [sentence[: valid[0].start()]]
        gaps.extend(sentence[left.end() : right.start()] for left, right in zip(valid, valid[1:]))
        gaps.append(sentence[valid[-1].end() :])
        return words, gaps

    @staticmethod
    def _delete_spaces_before_punctuation(text: str) -> str:
        for punctuation in '!"#$%&\'()*,./:;<=>?@[\\]^_`{|}~-':
            text = text.replace(" " + punctuation, punctuation)
            if punctuation == "-":
                text = text.replace(punctuation + " ", punctuation)
        return text.replace("~", "-")

    @staticmethod
    def _vowels(word: str) -> int:
        return sum(letter in "аеёиоуыэюяАЕЁИОУЫЭЮЯ" for letter in word)

    @staticmethod
    def _has_punctuation(word: str) -> bool:
        return any(letter in '!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~' for letter in word)

    def _process_yo(self, words: list[str], sentence: str) -> list[str]:
        predictions: list[str] = []
        if "е" in sentence.lower():
            predictions = [str(item["entity"]) for item in self.yo_classifier.classify_words(sentence.lower())]
        output = []
        for index, word in enumerate(words):
            lowered = word.lower()
            converted = _fix_capital(word, self.yo_words.get(lowered, word))
            if index < len(predictions) and predictions[index] == "YO":
                converted = _fix_capital(word, self.yo_homographs.get(lowered, word))
            output.append(converted)
        return output

    def _process_omographs(self, words: list[str]) -> list[str]:
        for index, word in enumerate(words):
            variants = self.omographs.get(word.lower())
            if not variants:
                continue
            context = words.copy()
            context[index] = f" <w>{word}</w> "
            words[index] = self.omograph_model.choose(" ".join(context), variants)
        return words

    def _process_accents(self, words: list[str], usages: list[str]) -> list[str]:
        for index, word in enumerate(words):
            if "+" in word or index >= len(usages) or usages[index] != "STRESS":
                continue
            lowered = word.lower()
            accented = self.accents.get(lowered, lowered)
            if accented == lowered and not self._has_punctuation(lowered) and self._vowels(lowered) > 1:
                words[index] = self.accent_model.put_accent(word)
                continue
            # Transfer dictionary marker positions to the source casing.
            target = list(word)
            inserted = 0
            for marker in re.finditer(r"\+", accented):
                position = marker.start() + inserted
                target.insert(position, "+")
                inserted += 1
            words[index] = "".join(target)
        return words

    def _dictionary_word(self, match: re.Match[str]) -> str:
        """Apply deterministic ``ё`` and accent dictionary entries only."""
        word = match.group(0)
        normalized = _fix_capital(word, self.yo_words.get(word.lower(), word))
        accented = self.accents.get(normalized.lower(), normalized.lower())
        if accented == normalized.lower():
            return normalized
        target = list(normalized)
        inserted = 0
        for marker in re.finditer(r"\+", accented):
            target.insert(marker.start() + inserted, "+")
            inserted += 1
        return "".join(target)

    def _process_dictionary(self, text: str) -> str:
        # Dictionary mode deliberately avoids loading or calling every neural
        # RUAccent ONNX graph. Unknown words and unresolved homographs remain
        # untouched rather than receiving a neural prediction.
        return re.sub(r"[A-Za-zА-Яа-яЁё]+", self._dictionary_word, text)

    def process_all(self, text: str) -> str:
        text = self._normalize.sub("", text)
        if self.mode == "dictionary":
            return self._process_dictionary(text)
        output: list[str] = []
        # Keep sentence delimiters attached to the sentence, matching the
        # original RUAccent intent without its optional razdel dependency.
        cursor = 0
        for match in self._sentence.finditer(text):
            output.append(text[cursor : match.start()])
            sentence = match.group(0)
            cursor = match.end()
            matches = list(self._tokens.finditer(sentence.replace(" - ", " ~ ")))
            words, gaps = self._remaining(sentence.replace(" - ", " ~ "), matches)
            if not words:
                output.append(sentence)
                continue
            usages = [str(item["entity"]) for item in self.stress_usage.classify_words(sentence)]
            words = self._process_yo(words, sentence)
            words = self._process_omographs(words)
            words = self._process_accents(words, usages)
            output.append(self._delete_spaces_before_punctuation("".join(gap + word for gap, word in zip(gaps, words)) + gaps[-1]))
        output.append(text[cursor:])
        return "".join(output)