"""
Stage 1: Tokenization (character-level).

This maps directly onto the "Tokenization summary" slide you saw:
we're starting at the CHARACTER level because it's the simplest possible
scheme — near-zero risk of out-of-vocabulary tokens, trivially
interpretable, and the vocab is tiny (~65 symbols for Shakespeare).
The cost is longer sequences and less meaningful individual tokens.

Later (Stage 4) you'll swap this out for a subword tokenizer (BPE) and
watch how vocab size, sequence length, and sample quality change.
That contrast is one of the best lessons in this whole exercise.

Key idea: a tokenizer is just two lookup tables.
    encode: string  -> list of integers
    decode: list of integers -> string
Everything downstream (embeddings, attention) operates on the integers.
"""

from __future__ import annotations

import json
from pathlib import Path


class CharTokenizer:
    """A character-level tokenizer built from a training corpus."""

    def __init__(self, text: str) -> None:
        # The vocabulary is simply every unique character, sorted so the
        # mapping is deterministic across runs.
        chars = sorted(set(text))
        self.vocab_size = len(chars)

        # stoi = "string to integer", itos = "integer to string".
        # These two dicts ARE the tokenizer. That's it.
        self.stoi: dict[str, int] = {ch: i for i, ch in enumerate(chars)}
        self.itos: dict[int, str] = {i: ch for i, ch in enumerate(chars)}

    def encode(self, text: str) -> list[int]:
        """Convert a string into a list of token ids."""
        return [self.stoi[ch] for ch in text]

    def decode(self, ids: list[int]) -> str:
        """Convert a list of token ids back into a string."""
        return "".join(self.itos[i] for i in ids)

    # --- persistence, so a trained model ships with its tokenizer ---

    def save(self, path: str | Path) -> None:
        Path(path).write_text(
            json.dumps({"stoi": self.stoi}, indent=2), encoding="utf-8"
        )

    @classmethod
    def load(cls, path: str | Path) -> "CharTokenizer":
        stoi = json.loads(Path(path).read_text(encoding="utf-8"))["stoi"]
        tok = cls.__new__(cls)
        tok.stoi = {k: int(v) for k, v in stoi.items()}
        tok.itos = {int(v): k for k, v in stoi.items()}
        tok.vocab_size = len(tok.stoi)
        return tok


if __name__ == "__main__":
    # Quick self-test / demo
    data_path = Path(__file__).resolve().parent.parent / "data" / "input.txt"
    text = data_path.read_text(encoding="utf-8")

    tok = CharTokenizer(text)
    print(f"Vocab size: {tok.vocab_size}")
    print(f"Vocabulary: {''.join(tok.itos[i] for i in range(tok.vocab_size))!r}")

    sample = "To be, or not to be"
    ids = tok.encode(sample)
    print(f"\nencode({sample!r})\n  -> {ids}")
    print(f"decode(...)\n  -> {tok.decode(ids)!r}")
    assert tok.decode(ids) == sample, "Round-trip failed!"
    print("\nRound-trip OK ✔")