"""Typing stubs for the `bytetok.bpe` public wrapper module.

This file intentionally exposes Rust-backed symbols via the public path
`bytetok.bpe` so tooling prefers:

`from bytetok.bpe import RustBPETrainer`

over the internal extension path:

`from bytetok._bpe_rs import RustBPETrainer`
"""

from __future__ import annotations

class RustBPETokenizer:
    def __init__(
        self,
        merge_history: list[tuple[tuple[int, int], int]],
        pattern: str,
        special_tokens: dict[str, int] = ...,
    ) -> None: ...
    def encode_text(self, text: str) -> list[int]: ...
    def encode_texts(
        self, texts: list[str], show_progress: bool = True
    ) -> list[list[int]]: ...
    def encode_text_with_special(
        self, text: str, allowed_special: dict[str, int]
    ) -> list[int]: ...
    def encode_texts_with_special(
        self,
        texts: list[str],
        allowed_special: dict[str, int],
        show_progress: bool = True,
    ) -> list[list[int]]: ...
    def encode_bytes(self, text: str) -> list[int]: ...
    def encode_bytes_batch(
        self, texts: list[str], show_progress: bool = True
    ) -> list[list[int]]: ...
    def decode_tokens(self, tokens: list[int], errors: str | None = None) -> str: ...
    def decode_tokens_batch(
        self,
        token_seqs: list[list[int]],
        errors: str | None = None,
        show_progress: bool = True,
    ) -> list[str]: ...
    def vocab_size(self) -> int: ...

class RustBPETrainer:
    def __init__(self, tokens: list[int], next_token_id: int) -> None: ...
    def train(self, num_merges: int, show_progress: bool = True) -> None: ...
    def merge_step(self) -> bool: ...
    def get_tokens(self) -> list[int]: ...
    def get_merge_history(self) -> list[tuple[tuple[int, int], int]]: ...
    def print_state(self) -> None: ...

__all__ = ["RustBPETrainer", "RustBPETokenizer"]
