from khmer_keyboard.model_io import load_model

from .chargram_embedder import CharGramEmbedder
from .language_model import LanguageModel
from .personalization import Personalization


class KhmerKeyboardEngine:
    def __init__(self, model_path):
        model = load_model(model_path=model_path)

        self.lm = LanguageModel(
            model["unigram"],
            model["bigram_next"],
            model["trigram_topk"],
        )
        self.chargram = CharGramEmbedder(model["chargram"])
        self.personal = Personalization()
        self.model_format_version = model["model_format_version"]

        self.sorted_unigram = sorted(
            model["unigram"].items(),
            key=lambda item: item[1],
            reverse=True,
        )
        self.unigram_top5 = [word for word, _ in self.sorted_unigram[:5]]
        self.vocab_ranked = [word for word, _ in self.sorted_unigram]
        self.vocab = set(self.vocab_ranked)
        self.char_candidates = self.vocab_ranked[:2000]

    def predict(self, text):
        words = text.strip().split()

        if not words:
            return self.unigram_top5

        if not text.endswith(" "):
            prefix = words[-1]

            if prefix not in self.vocab:
                candidates = []
                for word in self.char_candidates:
                    sim = self.chargram.similarity(prefix, word)
                    if sim > 0.3:
                        candidates.append((word, sim))
                candidates.sort(key=lambda item: item[1], reverse=True)
                return [word for word, _ in candidates[:5]]

            return self.unigram_top5

        if len(words) >= 2:
            w1, w2 = words[-2], words[-1]
            candidates = []
            for word in self.vocab_ranked:
                score = self.lm.score(w1, w2, word) + self.personal.boost(word)
                candidates.append((word, score))

            candidates.sort(key=lambda item: item[1], reverse=True)
            return [word for word, _ in candidates[:5]]

        return self.unigram_top5

