from typing import Callable, Mapping, Optional, Sequence, Tuple, Union

ByteOrChar = Union[str, int]
TrieNodeID = int
GeneralSAMNodeID = int

NIL_NODE_ID = 0
ROOT_NODE_ID = 1

class TrieNode:
    def is_in_chars(self) -> bool: ...
    def is_in_bytes(self) -> bool: ...
    def get_node_id(self) -> TrieNodeID: ...
    def is_accepting(self) -> bool: ...
    def get_trans(self) -> Mapping[ByteOrChar, TrieNodeID]: ...
    def get_parent(self) -> TrieNodeID: ...

class Trie:
    @staticmethod
    def in_chars() -> "Trie": ...
    @staticmethod
    def in_bytes() -> "Trie": ...
    def is_in_chars(self) -> bool: ...
    def is_in_bytes(self) -> bool: ...
    def num_of_nodes(self) -> int: ...
    def insert_chars(self, s: str) -> TrieNodeID: ...
    def insert_bytes(self, s: bytes) -> TrieNodeID: ...
    def get_bfs_order(self) -> Sequence[TrieNodeID]: ...
    def get_root(self) -> TrieNode: ...
    def get_node(self, node_id: TrieNodeID) -> Optional[TrieNode]: ...
    def dfs_travel(
        self,
        in_stack_callback: Callable[[TrieNodeID, Optional[ByteOrChar]], None],
        out_stack_callback: Callable[[TrieNodeID], None],
        root_node_id: Optional[TrieNodeID] = None,
    ) -> TrieNode: ...
    def bfs_travel(
        self,
        in_queue_callback: Callable[[TrieNodeID, Optional[ByteOrChar]], None],
        out_queue_callback: Callable[[TrieNodeID], None],
        root_node_id: Optional[TrieNodeID] = None,
    ) -> TrieNode: ...

class GeneralSAMState:
    def is_in_chars(self) -> bool: ...
    def is_in_bytes(self) -> bool: ...
    def get_node_id(self) -> GeneralSAMNodeID: ...
    def is_nil(self) -> bool: ...
    def is_root(self) -> bool: ...
    def is_accepting(self) -> bool: ...
    def get_trans(self) -> Mapping[ByteOrChar, GeneralSAMNodeID]: ...
    def get_suffix_parent_id(self) -> GeneralSAMNodeID: ...
    def copy(self) -> "GeneralSAMState": ...
    def goto_suffix_parent(self) -> None: ...
    def goto_char(self, t: str) -> None: ...
    def goto_byte(self, t: int) -> None: ...
    def feed_chars(self, s: str) -> None: ...
    def feed_bytes(self, s: bytes) -> None: ...
    def dfs_along(
        self,
        trie: Trie,
        in_stack_callback: Callable[
            ["GeneralSAMState", TrieNodeID, Optional[ByteOrChar]], None
        ],
        out_stack_callback: Callable[["GeneralSAMState", TrieNodeID], None],
        trie_node_id: Optional[TrieNodeID] = None,
    ) -> TrieNode: ...
    def bfs_along(
        self,
        trie: Trie,
        in_queue_callback: Callable[
            ["GeneralSAMState", TrieNodeID, Optional[ByteOrChar]], None
        ],
        out_queue_callback: Callable[["GeneralSAMState", TrieNodeID], None],
        trie_node_id: Optional[TrieNodeID] = None,
    ) -> TrieNode: ...

class GeneralSAM:
    @staticmethod
    def from_chars(s: str) -> "GeneralSAM": ...
    @staticmethod
    def from_bytes(s: bytes) -> "GeneralSAM": ...
    @staticmethod
    def from_trie(trie: Trie) -> "GeneralSAM": ...
    def is_in_chars(self) -> bool: ...
    def is_in_bytes(self) -> bool: ...
    def num_of_nodes(self) -> int: ...
    def get_root_state(self) -> GeneralSAMState: ...
    def get_state(self, node_id: GeneralSAMNodeID) -> GeneralSAMState: ...
    def get_topo_and_suf_len_sorted_states(self) -> Sequence[GeneralSAMState]: ...

class GreedyTokenizer:
    @staticmethod
    def from_sam_and_trie(sam: GeneralSAM, trie: Trie) -> "GreedyTokenizer": ...
    def get_sam(self) -> GeneralSAM: ...
    def is_in_chars(self) -> bool: ...
    def is_in_bytes(self) -> bool: ...
    def tokenize_str(
        self, s: str, unk_token_id: TrieNodeID = NIL_NODE_ID
    ) -> Sequence[Tuple[TrieNodeID, int]]: ...
    def tokenize_bytes(
        self, s: bytes, unk_token_id: TrieNodeID = NIL_NODE_ID
    ) -> Sequence[Tuple[TrieNodeID, int]]: ...
