import math
from collections import defaultdict

import Levenshtein

from .model_io import load_model


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

        self.unigram = model["unigram"]
        self.bigram_next = model["bigram_next"]
        self.trigram_topk = model["trigram_topk"]
        self.model_format_version = model["model_format_version"]

        self.sorted_unigram = sorted(
            self.unigram.items(),
            key=lambda x: x[1],
            reverse=True,
        )

        self.unigram_top5 = [w for w, _ in self.sorted_unigram[:5]]
        self.vocab = set(self.unigram.keys())
        self.total_unigram = sum(self.unigram.values()) or 1
        self.max_word_len = max((len(w) for w in self.vocab), default=1)

        self.prefix_index = defaultdict(list)
        for word, _ in self.sorted_unigram:
            for i in range(1, min(len(word), 6) + 1):
                prefix = word[:i]
                self.prefix_index[prefix].append(word)

        self.bigram_sorted = {}
        for w2, next_words in self.bigram_next.items():
            if isinstance(next_words, dict):
                sorted_items = sorted(
                    next_words.items(),
                    key=lambda x: x[1],
                    reverse=True,
                )
                self.bigram_sorted[w2] = [w for w, _ in sorted_items]
            else:
                self.bigram_sorted[w2] = [w for w, _ in next_words]

        self.first_char_index = defaultdict(list)
        for word in self.vocab:
            if word:
                self.first_char_index[word[0]].append(word)

    def correct_word(self, typo, top_n=3, include_exact=False):
        if not typo:
            return []

        candidates = []
        for word in self.first_char_index.get(typo[0], []):
            if not include_exact and word == typo:
                continue

            if abs(len(word) - len(typo)) > 2:
                continue

            dist = Levenshtein.distance(typo, word)
            if dist > 2:
                continue

            freq = self.unigram.get(word, 1)
            score = freq / (1 + dist)
            candidates.append((word, score))

        candidates.sort(key=lambda x: x[1], reverse=True)
        return [w for w, _ in candidates[:top_n]]

    def predict_next(self, w1, w2, top_n=5):
        key = f"{w1}\t{w2}"

        tri = self.trigram_topk.get(key)
        if tri:
            return [w for (w, _) in tri[:top_n]]

        bi = self.bigram_sorted.get(w2)
        if bi:
            return bi[:top_n]

        return self.unigram_top5[:top_n]

    def suggest_prefix(self, prefix, top_n=5):
        return self.prefix_index.get(prefix, [])[:top_n]

    def _segment_compound_token(self, token, max_words=4):
        # Split merged token into likely known words via unigram cost.
        if not token:
            return []

        n = len(token)
        inf = float("inf")
        best_cost = [inf] * (n + 1)
        best_path = [None] * (n + 1)
        best_cost[0] = 0.0
        best_path[0] = []

        for i in range(n):
            if best_path[i] is None:
                continue

            limit = min(n, i + self.max_word_len)
            for j in range(i + 1, limit + 1):
                piece = token[i:j]
                if piece not in self.vocab:
                    continue

                freq = self.unigram.get(piece, 0)
                prob = (freq + 1) / (self.total_unigram + len(self.vocab))
                cost = best_cost[i] + (-math.log(prob))

                path = best_path[i] + [piece]
                if len(path) > max_words:
                    continue

                if cost < best_cost[j]:
                    best_cost[j] = cost
                    best_path[j] = path

        return best_path[n] or []

    def _recover_from_suffix(self, token, top_n=5):
        # Try likely trailing sub-token for partial merged input.
        if not token:
            return []

        max_len = min(len(token), self.max_word_len)
        for size in range(max_len, 0, -1):
            suffix = token[-size:]
            suggestions = self.suggest_prefix(suffix, top_n=top_n)
            if suggestions:
                return suggestions

            corrected = self.correct_word(suffix, top_n=top_n)
            if corrected:
                return corrected

        return []

    def smart(
        self,
        text=None,
        context_words=None,
        current_prefix=None,
        top_n=5,
    ):
        # Optional text input for keyboard-like usage:
        # - trailing space => next-word prediction
        # - no trailing space => prefix suggestion/correction
        if text is not None:
            words = text.strip().split()
            if not words:
                return self.unigram_top5[:top_n]

            if text.endswith(" "):
                if len(words) >= 2:
                    return self.predict_next(words[-2], words[-1], top_n=top_n)
                if len(words) == 1:
                    bi = self.bigram_sorted.get(words[-1])
                    if bi:
                        return bi[:top_n]
                return self.unigram_top5[:top_n]

            current_prefix = words[-1]
            context_words = words[:-1]

        context_words = context_words or []

        if current_prefix:
            # Try splitting merged no-space phrase into words first.
            if (
                current_prefix not in self.vocab
                and not context_words
            ):
                split_words = self._segment_compound_token(current_prefix)
                if len(split_words) >= 2:
                    w1, w2 = split_words[-2], split_words[-1]
                    return self.predict_next(w1, w2, top_n=top_n)

            # If token is not in vocab and correction exists:
            # 1) try next-word prediction from corrected token
            # 2) fallback to corrected candidates
            if current_prefix not in self.vocab:
                corrected = self.correct_word(current_prefix, top_n=top_n)
                if corrected:
                    best = corrected[0]

                    if len(context_words) >= 1:
                        next_words = self.predict_next(
                            context_words[-1],
                            best,
                            top_n=top_n,
                        )
                        if next_words:
                            return next_words
                    else:
                        bi = self.bigram_sorted.get(best)
                        if bi:
                            return bi[:top_n]

                    return corrected

            suggestions = self.suggest_prefix(current_prefix, top_n=top_n)
            if suggestions:
                return suggestions

            corrected = self.correct_word(current_prefix, top_n=top_n)
            if corrected:
                return corrected

            recovered = self._recover_from_suffix(current_prefix, top_n=top_n)
            if recovered:
                return recovered

            return self.unigram_top5[:top_n]

        if len(context_words) >= 2:
            return self.predict_next(
                context_words[-2],
                context_words[-1],
                top_n=top_n,
            )

        if len(context_words) == 1:
            bi = self.bigram_sorted.get(context_words[-1])
            if bi:
                return bi[:top_n]

        return self.unigram_top5[:top_n]

    def predict(self, context_words=None, current_prefix=None):
        return self.smart(
            context_words=context_words,
            current_prefix=current_prefix,
            top_n=5,
        )
