import re
from operator import attrgetter
from typing import Dict, Literal, Tuple

from smb.common.multiruntime.lib import is_arcadia_python

if is_arcadia_python:
    from yandex_lemmer import AnalyzeWord as analyze_word
else:

    def analyze_word(*args, **kwargs):
        return []


__all__ = ["extract_lemmas"]


T = Dict[Literal["input", "lemmas"], Tuple[str, ...]]


def extract_lemmas(input_: str) -> Tuple[T, ...]:
    words = re.sub(r"[^\w-]", " ", input_).split()

    return tuple(_extract_lemmas_from_word(word) for word in words)


def _extract_lemmas_from_word(word: str) -> T:
    got = analyze_word(word, langs=("ru",), split=False)

    max_weight = max(map(attrgetter("Weight"), got))

    return dict(
        input=word,
        lemmas=tuple(
            map(
                attrgetter("Lemma"),
                filter(lambda el: el.Weight == max_weight, got),
            )
        ),
    )
