backend/app/query_normalizer.py

from __future__ import annotations

import re
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple

try:
    from rapidfuzz import fuzz, process  # type: ignore
except Exception:  # pragma: no cover
    fuzz = None  # type: ignore
    process = None  # type: ignore

try:
    from symspellpy import SymSpell, Verbosity  # type: ignore
except Exception:  # pragma: no cover
    SymSpell = None  # type: ignore
    Verbosity = None  # type: ignore

try:
    from importlib.resources import files as resource_files
except Exception:  # pragma: no cover
    resource_files = None  # type: ignore


_WORD_RE = re.compile(r"[A-Za-z][A-Za-z0-9']*[A-Za-z0-9]|[A-Za-z0-9]")
_TOKEN_OR_SEP_RE = re.compile(r"[A-Za-z][A-Za-z0-9']*[A-Za-z0-9]|[A-Za-z0-9]|\s+|[^A-Za-z0-9\s]+")
_MODELISH_RE = re.compile(r"^(?:[A-Za-z]{1,10}[-_]?\d+[A-Za-z0-9\-_.]*|\d+[A-Za-z]+[A-Za-z0-9\-_.]*)$")
_EMAIL_RE = re.compile(r"^[^@\s]+@[^@\s]+\.[^@\s]+$")
_URL_RE = re.compile(r"^https?://", flags=re.IGNORECASE)

_COMMON_TYPO_MAP = {
    "repalcement": "replacement",
    "repalce": "replace",
    "thier": "their",
    "teh": "the",
    "woudl": "would",
    "woud": "would",
    "reccomend": "recommend",
    "reccomendation": "recommendation",
    "comparision": "comparison",
    "compair": "compare",
    "diffrence": "difference",
    "differnce": "difference",
    "incontrol": "incontrol2",
}

_CORRECTABLE_STOPWORDS = {
    "a",
    "an",
    "and",
    "are",
    "as",
    "at",
    "be",
    "by",
    "for",
    "from",
    "has",
    "have",
    "how",
    "i",
    "in",
    "is",
    "it",
    "of",
    "on",
    "or",
    "the",
    "to",
    "what",
    "which",
    "with",
}


def _norm_text(value: Any) -> str:
    return re.sub(r"\s+", " ", str(value or "").strip())


def _preserve_case(src: str, replacement: str) -> str:
    if not src:
        return replacement
    if src.isupper():
        return replacement.upper()
    if src[0].isupper():
        return replacement.capitalize()
    return replacement


class QueryNormalizer:
    def __init__(
        self,
        *,
        protected_terms: Optional[Iterable[str]] = None,
        domain_terms: Optional[Iterable[str]] = None,
        enabled: bool = True,
        max_edit_distance: int = 2,
        min_ratio: int = 90,
    ) -> None:
        self.enabled = bool(enabled)
        self.max_edit_distance = max(1, min(2, int(max_edit_distance)))
        self.min_ratio = max(80, min(99, int(min_ratio)))
        self._protected: Set[str] = {
            re.sub(r"[^a-z0-9]", "", str(t or "").lower())
            for t in (protected_terms or [])
            if str(t or "").strip()
        }
        self._domain_terms: Set[str] = {str(t or "").strip().lower() for t in (domain_terms or []) if str(t or "").strip()}
        self._symspell = self._build_symspell()

    def _build_symspell(self) -> Optional[Any]:
        if (not self.enabled) or (SymSpell is None):
            return None
        try:
            sym = SymSpell(max_dictionary_edit_distance=self.max_edit_distance, prefix_length=7)
            if resource_files is not None:
                try:
                    pkg_root = resource_files("symspellpy")
                    dict_path = pkg_root.joinpath("frequency_dictionary_en_82_765.txt")
                    if dict_path and Path(str(dict_path)).exists():
                        sym.load_dictionary(str(dict_path), 0, 1)
                    bigram_path = pkg_root.joinpath("frequency_bigramdictionary_en_243_342.txt")
                    if bigram_path and Path(str(bigram_path)).exists():
                        sym.load_bigram_dictionary(str(bigram_path), 0, 2)
                except Exception:
                    pass
            for term in self._domain_terms:
                if len(term) >= 3:
                    sym.create_dictionary_entry(term, 10_000_000)
            return sym
        except Exception:
            return None

    def _is_protected(self, token: str) -> bool:
        t = str(token or "")
        if not t:
            return True
        if _EMAIL_RE.match(t) or _URL_RE.match(t):
            return True
        if any(ch.isdigit() for ch in t):
            if _MODELISH_RE.match(t):
                return True
            # Tokens with digits are often model/SKU-like in this app.
            return True
        normalized = re.sub(r"[^a-z0-9]", "", t.lower())
        if normalized in self._protected:
            return True
        up = t.upper()
        if ("-" in t) and up == t and len(t) >= 4:
            return True
        return False

    def _suggest_with_symspell(self, token: str) -> Optional[str]:
        if (self._symspell is None) or (Verbosity is None):
            return None
        try:
            suggestions = self._symspell.lookup(
                token,
                Verbosity.TOP,
                max_edit_distance=self.max_edit_distance,
                include_unknown=True,
                transfer_casing=False,
            )
        except Exception:
            return None
        if not suggestions:
            return None
        candidate = str(suggestions[0].term or "").strip().lower()
        if not candidate:
            return None
        return candidate

    def _suggest_with_rapidfuzz(self, token: str) -> Optional[str]:
        if (process is None) or (fuzz is None) or (not self._domain_terms):
            return None
        try:
            match = process.extractOne(token, list(self._domain_terms), scorer=fuzz.ratio)
        except Exception:
            return None
        if not match:
            return None
        candidate, score, _ = match
        if int(score) < self.min_ratio:
            return None
        return str(candidate).strip().lower() or None

    def normalize(self, text: str) -> Tuple[str, List[Dict[str, str]]]:
        raw = str(text or "")
        if (not self.enabled) or (not raw.strip()) or raw.strip().startswith("/"):
            return raw, []

        corrections: List[Dict[str, str]] = []
        out: List[str] = []
        for piece in _TOKEN_OR_SEP_RE.findall(raw):
            if not piece:
                continue
            if piece.isspace():
                out.append(piece)
                continue
            if not _WORD_RE.fullmatch(piece):
                out.append(piece)
                continue

            source = piece
            low = source.lower()
            if self._is_protected(source):
                out.append(source)
                continue
            if low in _CORRECTABLE_STOPWORDS:
                out.append(source)
                continue
            if len(low) < 4:
                out.append(source)
                continue

            fixed = _COMMON_TYPO_MAP.get(low)
            if not fixed:
                fixed = self._suggest_with_symspell(low)
            if not fixed:
                fixed = self._suggest_with_rapidfuzz(low)
            if (not fixed) or (fixed == low):
                out.append(source)
                continue
            if abs(len(fixed) - len(low)) > 3:
                out.append(source)
                continue
            if (fuzz is not None) and (int(fuzz.ratio(low, fixed)) < self.min_ratio):
                out.append(source)
                continue

            replaced = _preserve_case(source, fixed)
            out.append(replaced)
            corrections.append({"from": source, "to": replaced})

        normalized = "".join(out)
        normalized = _norm_text(normalized)
        if not normalized:
            return raw, corrections
        return normalized, corrections