"""
Finite State Transducer (FST) for term dictionaries.

Provides O(|key|) exact lookup with extreme prefix/suffix compression, modeled
after the approach used by Lucene and Tantivy.  This simplified implementation
uses a trie with shared prefixes and suffix deduplication at build time.

Two classes are exposed:

    - ``FSTBuilder``  Incrementally constructs an FST from keys added in sorted
                      order.  Call :meth:`finish` to produce an immutable FST.
    - ``FST``         Immutable, lookup-only transducer supporting exact match,
                      prefix search, and fuzzy (edit distance) search.

Zero external dependencies.
"""

from __future__ import annotations

from typing import Optional


# ---------------------------------------------------------------------------
# Internal node representation
# ---------------------------------------------------------------------------

class _Node:
    """A single state in the transducer graph.

    Attributes:
        children: Mapping from character label to child node.
        output:   Associated integer value if this node is a final state.
        is_final: Whether this node marks the end of a stored key.
    """

    __slots__ = ("children", "output", "is_final")

    def __init__(self) -> None:
        self.children: dict[str, _Node] = {}
        self.output: Optional[int] = None
        self.is_final: bool = False

    # _signature is computed externally by the builder to avoid deep recursion.


# ---------------------------------------------------------------------------
# FST (immutable, read-only)
# ---------------------------------------------------------------------------

class FST:
    """Immutable Finite State Transducer.

    Supports O(|key|) exact lookup, prefix search, and fuzzy search with edit
    distance.  Instances are produced by :meth:`FSTBuilder.finish`.
    """

    __slots__ = ("_root", "_size")

    def __init__(self, root: _Node, size: int) -> None:
        self._root = root
        self._size = size

    # -- exact lookup ------------------------------------------------------

    def get(self, key: str) -> Optional[int]:
        """Return the value associated with *key*, or ``None`` if absent.

        Time complexity: O(|key|).
        """
        node = self._root
        for ch in key:
            child = node.children.get(ch)
            if child is None:
                return None
            node = child
        if node.is_final:
            return node.output
        return None

    def has_key(self, key: str) -> bool:
        """Return ``True`` if *key* exists in the transducer."""
        return self.get(key) is not None

    # -- prefix search -----------------------------------------------------

    def prefix_search(self, prefix: str) -> list[tuple[str, int]]:
        """Return all ``(key, value)`` pairs whose key starts with *prefix*.

        Results are returned in lexicographic order.
        """
        node = self._root
        for ch in prefix:
            child = node.children.get(ch)
            if child is None:
                return []
            node = child

        results: list[tuple[str, int]] = []
        self._collect(node, prefix, results)
        return results

    def _collect(
        self, node: _Node, path: str, results: list[tuple[str, int]]
    ) -> None:
        """DFS collect all final states reachable from *node*."""
        if node.is_final:
            results.append((path, node.output))  # type: ignore[arg-type]
        for ch in sorted(node.children):
            self._collect(node.children[ch], path + ch, results)

    # -- fuzzy search (Levenshtein) ----------------------------------------

    def fuzzy_search(
        self, key: str, max_edit_distance: int = 1
    ) -> list[tuple[str, int, int]]:
        """Find all keys within *max_edit_distance* edits of *key*.

        Returns a list of ``(matched_key, value, edit_distance)`` tuples,
        sorted by (edit_distance, matched_key).

        Uses a standard Levenshtein automaton traversal over the trie.
        """
        results: list[tuple[str, int, int]] = []
        # Initial row of the Levenshtein matrix: [0, 1, 2, ..., len(key)]
        initial_row = list(range(len(key) + 1))
        self._fuzzy_recurse(
            self._root, "", key, initial_row, max_edit_distance, results
        )
        results.sort(key=lambda t: (t[2], t[0]))
        return results

    def _fuzzy_recurse(
        self,
        node: _Node,
        path: str,
        key: str,
        prev_row: list[int],
        max_dist: int,
        results: list[tuple[str, int, int]],
    ) -> None:
        """Recursive Levenshtein traversal of the trie."""
        columns = len(key) + 1

        for ch, child in node.children.items():
            current_row = [prev_row[0] + 1]

            for col in range(1, columns):
                insert_cost = current_row[col - 1] + 1
                delete_cost = prev_row[col] + 1
                replace_cost = prev_row[col - 1]
                if key[col - 1] != ch:
                    replace_cost += 1
                current_row.append(min(insert_cost, delete_cost, replace_cost))

            # If this child is a final state and distance is acceptable, record
            if child.is_final and current_row[columns - 1] <= max_dist:
                results.append(
                    (path + ch, child.output, current_row[columns - 1])  # type: ignore[arg-type]
                )

            # Prune: only recurse if any entry in current_row <= max_dist
            if min(current_row) <= max_dist:
                self._fuzzy_recurse(
                    child, path + ch, key, current_row, max_dist, results
                )

    # -- dunder protocols --------------------------------------------------

    def __len__(self) -> int:
        """Return the number of keys stored in the transducer."""
        return self._size

    def __contains__(self, key: object) -> bool:
        """Membership test: ``key in fst``."""
        if not isinstance(key, str):
            return False
        return self.has_key(key)


# ---------------------------------------------------------------------------
# FSTBuilder (write-once, sorted-order)
# ---------------------------------------------------------------------------

class FSTBuilder:
    """Incrementally construct an :class:`FST` from sorted keys.

    Keys **must** be added in strict lexicographic (sorted) order.  Adding an
    out-of-order key raises :class:`ValueError`.

    Example::

        builder = FSTBuilder()
        for i, word in enumerate(sorted(vocabulary)):
            builder.add(word, i)
        fst = builder.finish()
        assert fst.get("hello") is not None
    """

    __slots__ = ("_root", "_last_key", "_size", "_finished")

    def __init__(self) -> None:
        self._root = _Node()
        self._last_key: Optional[str] = None
        self._size = 0
        self._finished = False

    def add(self, key: str, value: int) -> None:
        """Insert a key-value pair into the builder.

        Parameters
        ----------
        key : str
            The term to index.  Must be greater than the previously added key
            (lexicographic order).
        value : int
            Non-negative integer value associated with *key*.

        Raises
        ------
        ValueError
            If *key* is out of order, a duplicate, or the builder is already
            finished.
        """
        if self._finished:
            raise ValueError("FSTBuilder has already been finished")
        if self._last_key is not None and key <= self._last_key:
            raise ValueError(
                f"Keys must be added in strict sorted order: "
                f"{key!r} <= {self._last_key!r}"
            )

        node = self._root
        for ch in key:
            child = node.children.get(ch)
            if child is None:
                child = _Node()
                node.children[ch] = child
            node = child

        node.is_final = True
        node.output = value
        self._last_key = key
        self._size += 1

    def finish(self) -> FST:
        """Freeze the builder and return an immutable :class:`FST`.

        After calling this method, no further keys may be added.  The builder
        performs suffix deduplication to compress the trie.
        """
        if self._finished:
            raise ValueError("FSTBuilder has already been finished")
        self._finished = True

        # Suffix deduplication: merge nodes with identical sub-automata.
        # Uses iterative post-order traversal to avoid stack overflow on
        # deep tries (e.g. a 1000-char key produces depth 1000).
        self._dedup_iterative()

        return FST(self._root, self._size)

    def _dedup_iterative(self) -> None:
        """Iterative post-order deduplication of suffix sub-trees."""
        cache: dict[tuple, _Node] = {}
        sig_cache: dict[int, tuple] = {}  # node id -> signature

        # Phase 1: collect nodes in post-order (children before parents).
        post_order: list[_Node] = []
        stack: list[_Node] = [self._root]
        visited: set[int] = set()

        while stack:
            node = stack[-1]
            nid = id(node)

            if nid in visited:
                stack.pop()
                post_order.append(node)
                continue

            visited.add(nid)
            # Push children (they will be processed before this node).
            for ch in sorted(node.children, reverse=True):
                child = node.children[ch]
                if id(child) not in visited:
                    stack.append(child)

        # Phase 2: compute signatures bottom-up and deduplicate.
        for node in post_order:
            # First, replace children with canonical versions.
            for ch in list(node.children):
                child = node.children[ch]
                child_sig = sig_cache.get(id(child))
                if child_sig is not None:
                    canonical = cache.get(child_sig)
                    if canonical is not None and canonical is not child:
                        node.children[ch] = canonical

            # Compute this node's signature using already-computed child sigs.
            child_sigs = tuple(
                (ch, sig_cache[id(node.children[ch])])
                for ch in sorted(node.children)
            )
            sig = (node.is_final, node.output, child_sigs)
            sig_cache[id(node)] = sig

            if sig not in cache:
                cache[sig] = node
