"""
RBSP Levenshtein Distance & Fuzzy Matching.

Provides edit-distance computation and fuzzy term matching used for
approximate string search (analogous to Lucene/Elasticsearch ``~`` queries).

Three distance functions are included:

    - ``edit_distance``            Classic Levenshtein DP.
    - ``edit_distance_threshold``  Early-exit variant that aborts when
                                   the distance exceeds a caller-supplied
                                   maximum, avoiding unnecessary work.
    - ``damerau_levenshtein``      Optimal string alignment variant that
                                   counts transpositions as a single edit.

``FuzzyMatcher`` wraps these into a convenient object that finds all
vocabulary terms within *max_edits* of a query term, optionally
filtering by a shared prefix for speed.
"""

from __future__ import annotations

from typing import Iterable, List, Tuple


# ---------------------------------------------------------------------------
# Core distance functions
# ---------------------------------------------------------------------------

def edit_distance(a: str, b: str) -> int:
    """Compute the Levenshtein distance between *a* and *b*.

    Classic O(mn) dynamic-programming algorithm using a single-row
    optimisation for O(min(m,n)) space.
    """
    if a == b:
        return 0

    la, lb = len(a), len(b)
    if la == 0:
        return lb
    if lb == 0:
        return la

    # Ensure *a* is the shorter string so the row array is smaller.
    if la > lb:
        a, b = b, a
        la, lb = lb, la

    # Previous row of distances — only need one row at a time.
    prev: List[int] = list(range(la + 1))

    for j in range(1, lb + 1):
        curr: List[int] = [j] + [0] * la
        for i in range(1, la + 1):
            cost = 0 if a[i - 1] == b[j - 1] else 1
            curr[i] = min(
                curr[i - 1] + 1,       # insertion
                prev[i] + 1,           # deletion
                prev[i - 1] + cost,    # substitution
            )
        prev = curr

    return prev[la]


def edit_distance_threshold(a: str, b: str, max_dist: int) -> int:
    """Compute Levenshtein distance with early exit.

    Returns the true distance when it is <= *max_dist*, otherwise returns
    ``max_dist + 1``.  The DP only evaluates a diagonal band of width
    ``2 * max_dist + 1``, giving O(n * max_dist) time instead of O(mn).
    """
    if a == b:
        return 0

    la, lb = len(a), len(b)

    # If the length difference alone exceeds the threshold, short-circuit.
    if abs(la - lb) > max_dist:
        return max_dist + 1

    if la == 0:
        return lb
    if lb == 0:
        return la

    # Keep *a* as the shorter string.
    if la > lb:
        a, b = b, a
        la, lb = lb, la

    sentinel = max_dist + 1

    prev: List[int] = list(range(la + 1))

    for j in range(1, lb + 1):
        curr: List[int] = [j] + [0] * la

        # Only evaluate columns within the diagonal band.
        lo = max(1, j - max_dist)
        hi = min(la, j + max_dist)

        # Fill entries outside the band with the sentinel.
        if lo > 1:
            curr[lo - 1] = sentinel
        for i in range(lo, hi + 1):
            cost = 0 if a[i - 1] == b[j - 1] else 1
            curr[i] = min(
                curr[i - 1] + 1,
                prev[i] + 1,
                prev[i - 1] + cost,
            )

        # If every value in the active band exceeds max_dist, bail out.
        if min(curr[lo:hi + 1]) > max_dist:
            return sentinel

        prev = curr

    return prev[la] if prev[la] <= max_dist else sentinel


# ---------------------------------------------------------------------------
# Damerau-Levenshtein (optimal string alignment)
# ---------------------------------------------------------------------------

def damerau_levenshtein(a: str, b: str) -> int:
    """Compute optimal-string-alignment distance (restricted Damerau-Levenshtein).

    Counts insertions, deletions, substitutions, **and adjacent transpositions**
    each as a single edit.  Uses the full O(mn) matrix to track transpositions.
    """
    if a == b:
        return 0

    la, lb = len(a), len(b)
    if la == 0:
        return lb
    if lb == 0:
        return la

    # Full matrix needed for the transposition look-back.
    d: List[List[int]] = [[0] * (lb + 1) for _ in range(la + 1)]

    for i in range(la + 1):
        d[i][0] = i
    for j in range(lb + 1):
        d[0][j] = j

    for i in range(1, la + 1):
        for j in range(1, lb + 1):
            cost = 0 if a[i - 1] == b[j - 1] else 1
            d[i][j] = min(
                d[i - 1][j] + 1,       # deletion
                d[i][j - 1] + 1,       # insertion
                d[i - 1][j - 1] + cost,  # substitution
            )
            # Transposition
            if (
                i > 1
                and j > 1
                and a[i - 1] == b[j - 2]
                and a[i - 2] == b[j - 1]
            ):
                d[i][j] = min(d[i][j], d[i - 2][j - 2] + 1)

    return d[la][lb]


# ---------------------------------------------------------------------------
# DamerauLevenshtein convenience class
# ---------------------------------------------------------------------------

class DamerauLevenshtein:
    """Convenience wrapper exposing Damerau-Levenshtein as a callable class.

    >>> dl = DamerauLevenshtein()
    >>> dl.distance("ca", "abc")
    3
    >>> dl.distance("abcd", "abdc")
    1
    """

    @staticmethod
    def distance(a: str, b: str) -> int:
        """Return the Damerau-Levenshtein distance between *a* and *b*."""
        return damerau_levenshtein(a, b)

    @staticmethod
    def normalized(a: str, b: str) -> float:
        """Return the normalised distance in [0.0, 1.0].

        0.0 means the strings are identical; 1.0 means maximally different.
        """
        if not a and not b:
            return 0.0
        d = damerau_levenshtein(a, b)
        return d / max(len(a), len(b))


# ---------------------------------------------------------------------------
# FuzzyMatcher
# ---------------------------------------------------------------------------

class FuzzyMatcher:
    """Find vocabulary terms within a given edit distance of a query term.

    Args:
        term:      The query term to match against.
        max_edits: Maximum allowed edit distance (default 2).
    """

    __slots__ = ("_term", "_max_edits")

    def __init__(self, term: str, max_edits: int = 2) -> None:
        if max_edits < 0:
            raise ValueError("max_edits must be non-negative")
        self._term = term
        self._max_edits = max_edits

    @property
    def term(self) -> str:
        """The query term."""
        return self._term

    @property
    def max_edits(self) -> int:
        """Maximum edit distance threshold."""
        return self._max_edits

    def distance(self, candidate: str) -> int:
        """Return the exact Levenshtein distance to *candidate*."""
        return edit_distance(self._term, candidate)

    def matches(self, candidate: str) -> bool:
        """Return ``True`` if *candidate* is within *max_edits* of the query term."""
        return edit_distance_threshold(self._term, candidate, self._max_edits) <= self._max_edits

    def find_matches(
        self, vocabulary: Iterable[str]
    ) -> List[Tuple[str, int]]:
        """Return all terms from *vocabulary* within *max_edits*, sorted by distance.

        Returns:
            List of ``(term, distance)`` pairs sorted by distance ascending,
            then lexicographically by term.
        """
        results: List[Tuple[str, int]] = []
        for word in vocabulary:
            d = edit_distance_threshold(self._term, word, self._max_edits)
            if d <= self._max_edits:
                results.append((word, d))
        results.sort(key=lambda t: (t[1], t[0]))
        return results

    def find_matches_prefix(
        self, vocabulary: Iterable[str], prefix_length: int = 3
    ) -> List[Tuple[str, int]]:
        """Find matches restricted to candidates sharing a common prefix.

        Only candidates whose first *prefix_length* characters match the query
        term's prefix are evaluated, dramatically reducing the search space for
        large vocabularies.

        If the query term is shorter than *prefix_length*, the full term is used
        as the prefix.

        Args:
            vocabulary:    Iterable of candidate strings.
            prefix_length: Number of leading characters that must match (default 3).

        Returns:
            List of ``(term, distance)`` pairs sorted by distance ascending,
            then lexicographically by term.
        """
        effective = min(prefix_length, len(self._term))
        if effective == 0:
            # Empty term — fall back to full search.
            return self.find_matches(vocabulary)

        prefix = self._term[:effective]
        results: List[Tuple[str, int]] = []
        for word in vocabulary:
            if not word[:effective] == prefix:
                continue
            d = edit_distance_threshold(self._term, word, self._max_edits)
            if d <= self._max_edits:
                results.append((word, d))
        results.sort(key=lambda t: (t[1], t[0]))
        return results
