"""Russian Snowball stemmer (https://snowballstem.org/algorithms/russian/stemmer.html).

Self-contained because `pyproject.toml` is a frozen contract (#34): no new dependencies for the RAG.
Latin words and digits pass through lowercased; "ё" is folded to "е".
"""

import re
from functools import lru_cache

_VOWELS = set("аеиоуыэюя")

# (suffixes, needs_a_ya): group-1 suffixes count only when preceded by "а"/"я" (which stays in the stem).
_PERFECTIVE_GERUND = ((("вшись", "вши", "в"), True), (("ившись", "ывшись", "ивши", "ывши", "ив", "ыв"), False))
_REFLEXIVE = ((("ся", "сь"), False),)
_ADJECTIVE = (
    (
        (
            "ими", "ыми", "его", "ого", "ему", "ому", "ее", "ие", "ые", "ое", "ей", "ий", "ый", "ой", "ем",
            "им", "ым", "ом", "их", "ых", "ую", "юю", "ая", "яя", "ою", "ею",
        ),
        False,
    ),
)  # fmt: skip
_PARTICIPLE = ((("ем", "нн", "вш", "ющ", "щ"), True), (("ивш", "ывш", "ующ"), False))
_VERB = (
    (("ла", "на", "ете", "йте", "ли", "й", "л", "ем", "н", "ло", "но", "ет", "ют", "ны", "ть", "ешь", "нно"), True),
    (
        (
            "ила", "ыла", "ена", "ейте", "уйте", "ите", "или", "ыли", "ей", "уй", "ил", "ыл", "им", "ым", "ен",
            "ило", "ыло", "ено", "ят", "ует", "уют", "ит", "ыт", "ены", "ить", "ыть", "ишь", "ую", "ю",
        ),
        False,
    ),
)  # fmt: skip
_NOUN = (
    (
        (
            "иями", "ями", "ами", "ией", "иям", "ием", "иях", "ев", "ов", "ие", "ье", "еи", "ии", "ей", "ой",
            "ий", "ям", "ем", "ам", "ом", "ах", "ях", "ию", "ью", "ия", "ья", "а", "е", "и", "й", "о", "у",
            "ы", "ь", "ю", "я",
        ),
        False,
    ),
)  # fmt: skip
_DERIVATIONAL = ((("ость", "ост"), False),)
_SUPERLATIVE = ((("ейше", "ейш"), False),)

_CYRILLIC_WORD = re.compile(r"^[а-я]+$")


def _strip(word: str, start: int, groups: tuple[tuple[tuple[str, ...], bool], ...]) -> str | None:
    """Remove the longest matching suffix lying in word[start:], or `None` if none matches."""
    best: tuple[int, str] | None = None
    for suffixes, needs_a_ya in groups:
        for suffix in suffixes:
            if not word.endswith(suffix) or len(word) - len(suffix) < start:
                continue
            if needs_a_ya and (len(word) - len(suffix) - 1 < start or word[-len(suffix) - 1] not in "ая"):
                continue
            if best is None or len(suffix) > best[0]:
                best = (len(suffix), word[: -len(suffix)])
    return None if best is None else best[1]


def _regions(word: str) -> tuple[int, int]:
    """RV (after the first vowel) and R2 start indices."""
    rv = next((i + 1 for i, ch in enumerate(word) if ch in _VOWELS), len(word))

    def after_vc(start: int) -> int:
        for i in range(start + 1, len(word)):
            if word[i] not in _VOWELS and word[i - 1] in _VOWELS:
                return i + 1
        return len(word)

    return rv, after_vc(after_vc(0))


@lru_cache(maxsize=65536)
def stem(word: str) -> str:
    word = word.lower().replace("ё", "е")
    if not _CYRILLIC_WORD.match(word):
        return word
    rv, r2 = _regions(word)

    # Step 1
    out = _strip(word, rv, _PERFECTIVE_GERUND)
    if out is None:
        word = _strip(word, rv, _REFLEXIVE) or word
        out = _strip(word, rv, _ADJECTIVE)
        if out is not None:
            out = _strip(out, rv, _PARTICIPLE) or out
        else:
            out = _strip(word, rv, _VERB)
            if out is None:
                out = _strip(word, rv, _NOUN)
    word = out if out is not None else word

    # Step 2
    if word.endswith("и") and len(word) - 1 >= rv:
        word = word[:-1]
    # Step 3
    word = _strip(word, r2, _DERIVATIONAL) or word
    # Step 4
    if word.endswith("нн") and len(word) - 1 >= rv:
        return word[:-1]
    sup = _strip(word, rv, _SUPERLATIVE)
    if sup is not None:
        word = sup
        return word[:-1] if word.endswith("нн") else word
    if word.endswith("ь") and len(word) - 1 >= rv:
        word = word[:-1]
    return word