"""
RBSP Spell Checker — "Did you mean?" correction via Symmetric Delete.

The Symmetric Delete algorithm (as used by SymSpell) pre-computes all single-
character deletion variants within a configurable edit distance at index time.
At lookup time, deletion variants of the query term are generated and matched
against the pre-computed dictionary, achieving O(1) average-case lookup speed
compared to computing full Levenshtein distances at query time.

Thread-safe.  Zero external dependencies.
"""

from __future__ import annotations

import threading
from typing import Dict, List, Optional, Set, Tuple


class SpellChecker:
    """Symmetric Delete spell checker with corpus-frequency awareness.

    Pre-computes all deletion variants within *max_edit_distance* for every
    dictionary word.  Lookups generate deletion variants of the query term
    and intersect them with the pre-computed set to find candidates, then
    verify each candidate with a true Levenshtein distance check.

    Args:
        max_edit_distance: Maximum edit distance for both indexing and lookup
                           (default 2).  Higher values increase memory usage
                           exponentially.
    """

    def __init__(self, max_edit_distance: int = 2) -> None:
        if max_edit_distance < 0:
            raise ValueError("max_edit_distance must be >= 0")
        self._max_edit_distance = max_edit_distance
        # {word: frequency}
        self._words: Dict[str, int] = {}
        # {delete_variant: set_of_original_words}
        self._deletes: Dict[str, Set[str]] = {}
        self._lock = threading.Lock()

    # ------------------------------------------------------------------
    # Public API — dictionary building
    # ------------------------------------------------------------------

    def add_word(self, word: str, frequency: int = 1) -> None:
        """Add a word to the dictionary with its corpus frequency.

        If the word already exists, *frequency* is added to the existing
        count.

        Args:
            word:      The dictionary word (will be lowercased).
            frequency: Corpus frequency count (default 1).
        """
        word = word.lower()
        with self._lock:
            self._words[word] = self._words.get(word, 0) + frequency
            deletes = _generate_deletes(word, self._max_edit_distance)
            for variant in deletes:
                bucket = self._deletes.get(variant)
                if bucket is None:
                    bucket = set()
                    self._deletes[variant] = bucket
                bucket.add(word)

    def add_words(self, words: Dict[str, int]) -> None:
        """Bulk-add words with their corpus frequencies.

        Args:
            words: Mapping of {word: frequency}.
        """
        for word, freq in words.items():
            self.add_word(word, freq)

    # ------------------------------------------------------------------
    # Public API — lookup & correction
    # ------------------------------------------------------------------

    def lookup(
        self,
        query: str,
        max_distance: Optional[int] = None,
        top_k: int = 5,
    ) -> List[Tuple[str, int, int]]:
        """Find spelling suggestions for *query*.

        Args:
            query:        The (possibly misspelled) term to look up.
            max_distance: Maximum edit distance for this lookup.  Defaults to
                          the instance ``max_edit_distance``.
            top_k:        Maximum number of suggestions to return (default 5).

        Returns:
            A list of ``(suggestion, distance, frequency)`` tuples sorted by
            ascending distance, then descending frequency.
        """
        query = query.lower()
        if max_distance is None:
            max_distance = self._max_edit_distance
        max_distance = min(max_distance, self._max_edit_distance)

        candidates: Dict[str, int] = {}  # {word: best_distance}

        with self._lock:
            words_snapshot = dict(self._words)
            deletes_snapshot = dict(self._deletes)

        # Exact match check
        if query in words_snapshot:
            candidates[query] = 0

        # Generate deletion variants of the query and look up each one
        query_deletes = _generate_deletes(query, max_distance)
        query_deletes.add(query)

        for variant in query_deletes:
            bucket = deletes_snapshot.get(variant)
            if bucket is None:
                continue
            for suggestion in bucket:
                if suggestion in candidates:
                    continue
                dist = _levenshtein(query, suggestion)
                if dist <= max_distance:
                    candidates[suggestion] = dist

        # Build result list
        results: List[Tuple[str, int, int]] = []
        for word, dist in candidates.items():
            freq = words_snapshot.get(word, 0)
            results.append((word, dist, freq))

        # Sort: ascending distance, then descending frequency
        results.sort(key=lambda r: (r[1], -r[2]))
        return results[:top_k]

    def correct_query(self, query_tokens: List[str]) -> List[str]:
        """Correct each token in a tokenized query.

        Tokens that exist in the dictionary are kept as-is.  Tokens that do
        not exist are replaced by the best suggestion (lowest distance, then
        highest frequency).  If no suggestion is found, the token is kept
        unchanged.

        Args:
            query_tokens: List of query tokens.

        Returns:
            A new list of tokens with corrections applied.
        """
        corrected: List[str] = []
        for token in query_tokens:
            lower = token.lower()
            with self._lock:
                known = lower in self._words
            if known:
                corrected.append(lower)
                continue
            suggestions = self.lookup(lower, top_k=1)
            if suggestions:
                corrected.append(suggestions[0][0])
            else:
                corrected.append(lower)
        return corrected

    def did_you_mean(self, query: str) -> Optional[str]:
        """Return a corrected query string, or ``None`` if no corrections.

        Splits *query* on whitespace, corrects each token via
        ``correct_query``, and returns the joined result only if at least
        one token was changed.

        Args:
            query: The raw query string.

        Returns:
            The corrected query string, or ``None`` if no corrections were
            made.
        """
        tokens = query.lower().split()
        if not tokens:
            return None
        corrected = self.correct_query(tokens)
        if corrected == tokens:
            return None
        return " ".join(corrected)


# ----------------------------------------------------------------------
# Internal helpers
# ----------------------------------------------------------------------

def _generate_deletes(word: str, max_distance: int) -> Set[str]:
    """Generate all deletion variants of *word* within *max_distance*.

    For each character position, produces the word with that character
    removed.  Recurses to produce multi-character deletions up to
    *max_distance*.

    Args:
        word:         The source word.
        max_distance: Maximum number of character deletions.

    Returns:
        A set of all unique deletion variants (does not include the
        original word itself).
    """
    deletes: Set[str] = set()
    if max_distance < 1 or not word:
        return deletes
    queue: List[Tuple[str, int]] = [(word, 0)]
    seen: Set[str] = {word}
    while queue:
        current, depth = queue.pop()
        next_depth = depth + 1
        for i in range(len(current)):
            delete = current[:i] + current[i + 1:]
            if delete not in seen:
                seen.add(delete)
                deletes.add(delete)
                if next_depth < max_distance:
                    queue.append((delete, next_depth))
    return deletes


def _levenshtein(s: str, t: str) -> int:
    """Compute the Levenshtein (edit) distance between two strings.

    Uses the classic two-row dynamic programming approach with O(min(m,n))
    space.

    Args:
        s: First string.
        t: Second string.

    Returns:
        The minimum number of single-character edits (insertions,
        deletions, substitutions) to transform *s* into *t*.
    """
    if s == t:
        return 0
    len_s = len(s)
    len_t = len(t)
    if len_s == 0:
        return len_t
    if len_t == 0:
        return len_s

    # Ensure s is the shorter string for space efficiency
    if len_s > len_t:
        s, t = t, s
        len_s, len_t = len_t, len_s

    prev_row: List[int] = list(range(len_s + 1))
    for j in range(1, len_t + 1):
        curr_row: List[int] = [j] + [0] * len_s
        for i in range(1, len_s + 1):
            cost = 0 if s[i - 1] == t[j - 1] else 1
            curr_row[i] = min(
                curr_row[i - 1] + 1,      # insertion
                prev_row[i] + 1,           # deletion
                prev_row[i - 1] + cost,    # substitution
            )
        prev_row = curr_row
    return prev_row[len_s]
