"""
Spelling correction module using SymSpell (Symmetric Delete Spelling Correction).
Optimized for O(1) delete lookups to match column names and operation keywords.
"""
from __future__ import annotations

import re

def _get_deletes(word: str, max_edit_distance: int = 2) -> set[str]:
    """Generate deletes for a word up to a maximum edit distance."""
    deletes = set()
    queue = {word}
    for _ in range(max_edit_distance):
        next_queue = set()
        for w in queue:
            if len(w) > 1:
                for i in range(len(w)):
                    del_w = w[:i] + w[i+1:]
                    deletes.add(del_w)
                    next_queue.add(del_w)
        queue = next_queue
    return deletes

def _lev_dist(s1: str, s2: str) -> int:
    """Calculate the Levenshtein distance between two strings."""
    if len(s1) < len(s2):
        return _lev_dist(s2, s1)
    if len(s2) == 0:
        return len(s1)
    previous_row = list(range(len(s2) + 1))
    for i, c1 in enumerate(s1):
        current_row = [i + 1]
        for j, c2 in enumerate(s2):
            insertions = previous_row[j + 1] + 1
            deletions = current_row[j] + 1
            substitutions = previous_row[j] + (c1 != c2)
            current_row.append(min(insertions, deletions, substitutions))
        previous_row = current_row
    return previous_row[-1]

class SymSpell:
    """A lightweight symmetric delete spelling corrector."""
    def __init__(self, max_edit_distance: int = 2):
        self.max_edit_distance = max_edit_distance
        # Maps delete_item -> set of original_words
        self.deletes: dict[str, set[str]] = {}
        self.words: set[str] = set()

    def add_word(self, word: str) -> None:
        """Index a word and its deletes for O(1) spelling lookup."""
        word = word.lower().strip()
        if not word or word in self.words:
            return
        self.words.add(word)
        
        # Index word itself
        if word not in self.deletes:
            self.deletes[word] = set()
        self.deletes[word].add(word)
        
        # Index deletions
        for delete in _get_deletes(word, self.max_edit_distance):
            if delete not in self.deletes:
                self.deletes[delete] = set()
            self.deletes[delete].add(word)

    def lookup(self, word: str) -> list[str]:
        """Find candidate words matching the spelling of input word."""
        word = word.lower().strip()
        if not word:
            return []
        if word in self.words:
            return [word]
        
        candidates: set[str] = set()
        # 1. Direct delete match
        if word in self.deletes:
            candidates.update(self.deletes[word])
            
        # 2. Deletes of word match
        for delete in _get_deletes(word, self.max_edit_distance):
            if delete in self.deletes:
                candidates.update(self.deletes[delete])
            if delete in self.words:
                candidates.add(delete)
                
        # Score and rank candidates by Levenshtein distance
        scored = []
        for cand in candidates:
            dist = _lev_dist(word, cand)
            if dist <= self.max_edit_distance:
                scored.append((cand, dist))
                
        # Sort by distance first, then length (longer words first for ties)
        scored.sort(key=lambda x: (x[1], -len(x[0])))
        return [c for c, _ in scored]

    def correct_query(self, query: str) -> str:
        """Correct typos in words within the query string."""
        # Find all alphabet-only words
        words = re.findall(r'[a-zA-Z]+', query)
        corrected = query
        for w in words:
            if len(w) > 2:  # Only correct words longer than 2 characters
                suggestions = self.lookup(w)
                if suggestions:
                    # Match exact word boundary to prevent partial replacements
                    corrected = re.sub(rf'\b{re.escape(w)}\b', suggestions[0], corrected, flags=re.IGNORECASE)
        return corrected