"""
RBSP Auto-complete / Type-ahead Suggestion Engine.

Provides prefix-based and fuzzy suggestion lookup with O(1) prefix retrieval
using an internal prefix-to-entries mapping.  All operations are thread-safe.

Classes:
    - SuggestionEntry:   A single suggestion with text, weight, and payload.
    - Suggester:         Core prefix + fuzzy suggestion engine.
    - ContextSuggester:  Context-filtered suggestions (file type, directory, etc.).
"""

from __future__ import annotations

import heapq
import threading
from dataclasses import dataclass, field
from typing import Dict, FrozenSet, List, Optional, Set


_MAX_PREFIX_LEN = 20


@dataclass(frozen=True, slots=True)
class SuggestionEntry:
    """A single auto-complete suggestion.

    Attributes:
        text:    The full suggestion text.
        weight:  Numeric weight for ranking (higher = better).
        payload: Optional opaque payload attached to this entry.
    """

    text: str
    weight: int = 1
    payload: Optional[str] = None


class Suggester:
    """Prefix-based and fuzzy auto-complete suggestion engine.

    Internally maps every prefix substring (up to ``_MAX_PREFIX_LEN`` chars)
    to the set of entries whose text starts with that prefix, giving O(1)
    prefix lookup.  Top-K selection uses ``heapq.nlargest`` by weight.

    Thread-safe: all mutations and reads are protected by a lock.
    """

    def __init__(self) -> None:
        self._lock = threading.Lock()
        # prefix string -> set of SuggestionEntry
        self._prefix_map: Dict[str, Set[SuggestionEntry]] = {}
        # text -> SuggestionEntry (canonical)
        self._entries: Dict[str, SuggestionEntry] = {}

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def add(self, text: str, weight: int = 1, payload: Optional[str] = None) -> None:
        """Add or update a suggestion entry.

        If *text* already exists, the entry is replaced (old prefixes are
        cleaned up first).
        """
        entry = SuggestionEntry(text=text, weight=weight, payload=payload)
        normalised = text.lower()
        with self._lock:
            # Remove old entry if updating
            if text in self._entries:
                self._remove_unlocked(text)
            self._entries[text] = entry
            for length in range(1, min(len(normalised), _MAX_PREFIX_LEN) + 1):
                prefix = normalised[:length]
                bucket = self._prefix_map.get(prefix)
                if bucket is None:
                    bucket = set()
                    self._prefix_map[prefix] = bucket
                bucket.add(entry)

    def suggest(self, prefix: str, top_k: int = 10) -> List[SuggestionEntry]:
        """Return up to *top_k* suggestions whose text starts with *prefix*.

        Results are sorted by descending weight, then alphabetically by text
        for deterministic ordering among equal weights.
        """
        normalised = prefix.lower()
        with self._lock:
            bucket = self._prefix_map.get(normalised)
            if not bucket:
                return []
            return heapq.nlargest(
                top_k,
                bucket,
                key=lambda e: (e.weight, e.text),
            )

    def suggest_fuzzy(
        self,
        query: str,
        top_k: int = 10,
        max_edits: int = 1,
    ) -> List[SuggestionEntry]:
        """Return fuzzy suggestions within *max_edits* edit distance of *query*.

        Uses Levenshtein distance computed against each stored entry's text.
        """
        normalised = query.lower()
        with self._lock:
            candidates: List[SuggestionEntry] = []
            for entry in self._entries.values():
                if _edit_distance(normalised, entry.text.lower(), max_edits) <= max_edits:
                    candidates.append(entry)
            return heapq.nlargest(
                top_k,
                candidates,
                key=lambda e: (e.weight, e.text),
            )

    def remove(self, text: str) -> bool:
        """Remove a suggestion by its exact text.  Returns ``True`` if found."""
        with self._lock:
            return self._remove_unlocked(text)

    def __len__(self) -> int:
        with self._lock:
            return len(self._entries)

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

    def _remove_unlocked(self, text: str) -> bool:
        """Remove *text* from entries and all prefix buckets.  Caller holds lock."""
        entry = self._entries.pop(text, None)
        if entry is None:
            return False
        normalised = text.lower()
        for length in range(1, min(len(normalised), _MAX_PREFIX_LEN) + 1):
            prefix = normalised[:length]
            bucket = self._prefix_map.get(prefix)
            if bucket is not None:
                bucket.discard(entry)
                if not bucket:
                    del self._prefix_map[prefix]
        return True


class ContextSuggester(Suggester):
    """Suggestions filtered by context (e.g. file type, directory).

    Each entry may be associated with one or more context strings.  At query
    time callers can restrict results to entries matching any of the supplied
    contexts.  If no contexts are provided the full set is searched.
    """

    def __init__(self) -> None:
        super().__init__()
        # text -> frozenset of context strings
        self._contexts: Dict[str, FrozenSet[str]] = {}

    def add(  # type: ignore[override]
        self,
        text: str,
        weight: int = 1,
        payload: Optional[str] = None,
        contexts: Optional[List[str]] = None,
    ) -> None:
        """Add a suggestion with optional context tags."""
        with self._lock:
            # Store contexts before super().add() which also acquires the lock,
            # so we inline the parent logic to avoid double-locking.
            if text in self._entries:
                self._remove_unlocked(text)
                self._contexts.pop(text, None)

            entry = SuggestionEntry(text=text, weight=weight, payload=payload)
            normalised = text.lower()
            self._entries[text] = entry
            for length in range(1, min(len(normalised), _MAX_PREFIX_LEN) + 1):
                prefix = normalised[:length]
                bucket = self._prefix_map.get(prefix)
                if bucket is None:
                    bucket = set()
                    self._prefix_map[prefix] = bucket
                bucket.add(entry)

            if contexts:
                self._contexts[text] = frozenset(contexts)

    def suggest(  # type: ignore[override]
        self,
        prefix: str,
        top_k: int = 10,
        contexts: Optional[List[str]] = None,
    ) -> List[SuggestionEntry]:
        """Prefix suggestions optionally filtered by context.

        If *contexts* is provided, only entries tagged with at least one of
        the given context strings are returned.
        """
        normalised = prefix.lower()
        with self._lock:
            bucket = self._prefix_map.get(normalised)
            if not bucket:
                return []
            if contexts is None:
                candidates = bucket
            else:
                ctx_set = frozenset(contexts)
                candidates = {
                    e
                    for e in bucket
                    if self._contexts.get(e.text, frozenset()) & ctx_set
                }
            return heapq.nlargest(
                top_k,
                candidates,
                key=lambda e: (e.weight, e.text),
            )

    def remove(self, text: str) -> bool:
        """Remove a suggestion and its context mapping."""
        with self._lock:
            removed = self._remove_unlocked(text)
            if removed:
                self._contexts.pop(text, None)
            return removed


# ------------------------------------------------------------------
# Edit distance helper
# ------------------------------------------------------------------


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

    Returns early with ``threshold + 1`` when the distance is guaranteed
    to exceed *threshold*, keeping the cost bounded for short max_edits.
    """
    la, lb = len(a), len(b)
    if abs(la - lb) > threshold:
        return threshold + 1

    # Use the shorter string for the column to save memory
    if la > lb:
        a, b = b, a
        la, lb = lb, la

    prev = list(range(la + 1))
    for j in range(1, lb + 1):
        curr = [j] + [0] * la
        row_min = j
        for i in range(1, la + 1):
            cost = 0 if a[i - 1] == b[j - 1] else 1
            curr[i] = min(
                prev[i] + 1,       # deletion
                curr[i - 1] + 1,   # insertion
                prev[i - 1] + cost, # substitution
            )
            if curr[i] < row_min:
                row_min = curr[i]
        if row_min > threshold:
            return threshold + 1
        prev = curr

    return prev[la]
