"""Slovak lemmatization for lexical retrieval and concept keys (Stanza).

"chatbotom", "chatbota" and "chatbot" share the lemma "chatbot"; "reči" is
"reč". Lemmas extend the full-text index and the literal keys of the activity
graph; they never replace the source text shown in answers.

The Stanza Slovak model (tokenize, pos, lemma) is prepared into the image with
the other runtime assets; no download happens at run time. Lemmas of source
units are cached on disk by text hash, so only new units are processed after
a sync. RETRIEVAL_LEMMATIZER=off (or a missing model) falls back to a rough
suffix stem; the switch exists for the documented A/B comparison.
"""
from __future__ import annotations

import hashlib
import json
import os
import re
import sqlite3
import threading
from contextlib import closing
from functools import lru_cache
from pathlib import Path

from tools.answers.context_policy import normalized

LEMMA_VERSION = "stanza-sk-1"
_lock = threading.Lock()
_SK_ENDINGS = ("ami", "ach", "och", "ovi", "ych", "ymi", "eho", "emu", "ej", "om", "ov", "ou",
               "ia", "ie", "iu", "a", "u", "e", "y", "i", "o")


def enabled() -> bool:
    return os.getenv("RETRIEVAL_LEMMATIZER", "stanza").lower() != "off"


def word_stem(word: str) -> str:
    """Fallback: rough Slovak case stem ("chatbotom" -> "chatbot")."""
    if len(word) < 5 or not word.isalpha():
        return word
    for ending in _SK_ENDINGS:
        if word.endswith(ending) and len(word) - len(ending) >= 4:
            return word[: -len(ending)]
    return word


def stanza_dir() -> Path | None:
    try:
        from stanza.resources.common import DEFAULT_MODEL_DIR
    except Exception:
        DEFAULT_MODEL_DIR = None
    for candidate in (os.getenv("STANZA_RESOURCES_DIR"),
                      os.getenv("ZP_RUNTIME_ASSETS") and str(Path(os.environ["ZP_RUNTIME_ASSETS"]) / "stanza"),
                      DEFAULT_MODEL_DIR):
        if candidate and (Path(candidate) / "sk").is_dir():
            return Path(candidate)
    return None


@lru_cache(maxsize=1)
def _pipeline():
    import logging
    import stanza
    logging.getLogger("stanza").setLevel(logging.ERROR)
    directory = stanza_dir()
    if directory is None:
        raise RuntimeError("Stanza Slovak model is not available")
    return stanza.Pipeline("sk", dir=str(directory), processors="tokenize,pos,lemma", verbose=False,
                           use_gpu=False, download_method=None)


def available() -> bool:
    if not enabled():
        return False
    try:
        _pipeline()
        return True
    except Exception:
        return False


def _fallback(text: str) -> list[str]:
    return [word_stem(w) for w in re.findall(r"[a-z0-9+#]+", normalized(text))]


def _from_doc(doc) -> list[str]:
    out = []
    for sentence in doc.sentences:
        for word in sentence.words:
            lemma = normalized(word.lemma or word.text)
            out.extend(re.findall(r"[a-z0-9+#]+", lemma))
    return out


def lemmas(text: str) -> list[str]:
    """Normalized lemma tokens (diacritics removed, lower case)."""
    if not text or not text.strip():
        return []
    if not available():
        return _fallback(text)
    with _lock:
        return _from_doc(_pipeline()(text))


def lemma_map(text: str) -> dict:
    """Normalized surface word -> normalized lemma for the words of a text."""
    if not text or not text.strip() or not available():
        return {}
    with _lock:
        doc = _pipeline()(text)
    out = {}
    for sentence in doc.sentences:
        for word in sentence.words:
            surface, lemma = normalized(word.text), normalized(word.lemma or word.text)
            if re.fullmatch(r"[a-z0-9+#]+", surface) and re.fullmatch(r"[a-z0-9+#]+", lemma):
                out[surface] = lemma
    return out


def lemma_text(text: str) -> str:
    return " ".join(lemmas(text))


def _key(text: str) -> str:
    return hashlib.sha256((LEMMA_VERSION + "\0" + text).encode()).hexdigest()


class LemmaCache:
    """Persistent text -> lemma string cache (SQLite, one row per text hash)."""

    def __init__(self, path: Path):
        self.path = Path(path)

    def _connect(self):
        self.path.parent.mkdir(parents=True, exist_ok=True)
        connection = sqlite3.connect(self.path, timeout=30)
        connection.execute("CREATE TABLE IF NOT EXISTS lemmas (key TEXT PRIMARY KEY, lemmas TEXT NOT NULL)")
        return connection

    def many(self, texts: list[str]) -> list[str]:
        """Lemma strings for texts; missing ones are computed in bulk and stored."""
        if not available():
            return [" ".join(_fallback(t)) for t in texts]
        keys = [_key(t) for t in texts]
        found = {}
        try:
            with closing(self._connect()) as connection:
                for start in range(0, len(keys), 500):
                    chunk = keys[start:start + 500]
                    rows = connection.execute(
                        f"SELECT key, lemmas FROM lemmas WHERE key IN ({','.join('?' * len(chunk))})", chunk)
                    found.update(rows.fetchall())
        except sqlite3.Error:
            pass  # a cache is an optimization, never the source of truth
        missing = [(k, t) for k, t in zip(keys, texts) if k not in found]
        if missing:
            import stanza
            with _lock:
                pipeline = _pipeline()
                docs = pipeline.bulk_process([stanza.Document([], text=t) for _, t in missing])
            fresh = {k: " ".join(_from_doc(doc)) for (k, _), doc in zip(missing, docs)}
            found.update(fresh)
            try:
                with closing(self._connect()) as connection, connection:
                    connection.executemany("INSERT OR REPLACE INTO lemmas VALUES (?,?)", list(fresh.items()))
            except sqlite3.Error:
                pass
        return [found[k] for k in keys]


def cache_path(data_dir: Path | None = None) -> Path:
    from tools.rag.common import DATA_DIR
    return Path(data_dir or DATA_DIR) / "knowledge" / "lemma_cache.sqlite"


def warm(data_dir: Path | None = None) -> dict:
    """Precompute lemmas of every published content unit (after a build or sync)."""
    from tools.knowledge.store import load_index
    from tools.sync.index_state import read_version
    from tools.rag.common import DATA_DIR
    data_dir = Path(data_dir or DATA_DIR)
    index = load_index(data_dir, read_version(data_dir))
    texts = [u["text"] for u in index["units"]]
    LemmaCache(cache_path(data_dir)).many(texts)
    return {"units": len(texts), "lemmatizer": LEMMA_VERSION if available() else "fallback-stem"}


if __name__ == "__main__":
    print(json.dumps(warm()))