"""Paprastas raktažodžių pagrindu veikiantis kategorizatorius.

Jei šaltinis pateikė kategoriją – ja pasitikim. Jei ne – atspėjam pagal
raktažodžių žodynus keliomis kalbomis. Tikslas – ne idealiai klasifikuoti,
o pateikti naudingą greitą filtrą.
"""

from __future__ import annotations

import re
from typing import Iterable, Optional

CATEGORIES = ("sports", "politics", "entertainment", "events", "news", "other")

KEYWORDS = {
    "sports": [
        # EN
        "match", "goal", "score", "league", "football", "soccer", "basketball",
        "nba", "nfl", "uefa", "fifa", "tennis", "olympic", "tournament",
        "champion", "playoffs", "boxing", "mma", "hockey", "cricket", "rugby",
        "vs", "fc", "afc", "euroleague", "world cup",
        # LT
        "rungtynės", "krepšinis", "futbolas", "čempion", "varžybos", "olimp",
        "treneris", "lyga", "žalgiris", "rytas", "lietkabelis", "neptūnas",
        # ES/PT/IT/DE
        "fútbol", "futebol", "calcio", "fußball", "bundesliga", "liga",
        # gerai žinomi klubai / lygos
        "cska", "real madrid", "barcelona", "bayern", "liverpool",
        "manchester", "psg", "juventus",
    ],
    "politics": [
        "election", "parliament", "president", "prime minister", "minister",
        "government", "senate", "congress", "vote", "ballot", "policy",
        "diplomat", "sanction", "summit", "coalition", "law", "bill",
        # LT
        "rinkim", "seimas", "premjeras", "ministras", "vyriausybė",
        "prezident", "balsavim",
        # other
        "putin", "biden", "trump", "macron", "zelensky", "xi jinping",
        "nato", "eu summit", "kremlin", "white house",
    ],
    "entertainment": [
        "movie", "film", "actor", "actress", "celebrity", "singer", "concert",
        "album", "song", "music", "netflix", "hbo", "oscar", "grammy",
        "billboard", "tiktok", "youtube", "drama", "series", "show",
        # LT
        "filmas", "muzika", "daina", "albumas", "aktor", "koncert", "žvaigždė",
        # generic
        "trailer", "premiere",
    ],
    "events": [
        "festival", "concert", "expo", "summit", "conference", "meeting",
        "march", "protest", "rally", "ceremony", "parade",
        # LT
        "festivalis", "konferencija", "renginys", "eitynės", "mitingas",
    ],
}

# Leidžiame ne tik tikslų žodžio matchą, bet ir kamienus su priesagomis
# (svarbu lietuvių/latvių/rusų ir t.t. linksniavimui): "prezident" matchina
# "prezidentas", "prezidento", "prezidentinis"...
# Žodžio pradžia – \b, gale – \w* (priesagai).
_COMPILED = {
    cat: re.compile(
        r"\b(?:" + "|".join(re.escape(k) for k in kws) + r")\w*",
        flags=re.IGNORECASE,
    )
    for cat, kws in KEYWORDS.items()
}


def categorize(text: str, hint: Optional[str] = None) -> str:
    """Grąžina vieną iš CATEGORIES.

    `hint` – iš šaltinio gauta kategorija (jei pateikta) – pasitikima, jei
    atitinka mūsų aibę. Kitu atveju – paprastas keyword match.
    """
    if hint:
        h = hint.strip().lower()
        if h in CATEGORIES:
            return h
        # mapinam dažnus aliasus
        alias = {
            "sport": "sports", "policy": "politics", "election": "politics",
            "movies": "entertainment", "music": "entertainment",
            "celebrity": "entertainment", "event": "events", "news": "news",
        }
        if h in alias:
            return alias[h]

    if not text:
        return "other"

    scores = {cat: 0 for cat in CATEGORIES if cat != "other"}
    for cat, rgx in _COMPILED.items():
        scores[cat] = len(rgx.findall(text))

    best = max(scores, key=scores.get)
    if scores[best] == 0:
        return "news"   # jei nieko neaptikta – plačiausia kategorija "news"
    return best