"""From-scratch tokenizers for the playground.

Two schemes share one interface so the rest of the app (training + serving) can
treat them identically:

    encode(text)            -> list[int]                 ids for the model
    decode(ids)             -> str                        inverse of encode
    tokens_with_spans(text) -> list[{text,id,start,end}]  for the colored chips
    vocab_size              -> int
    save(dir) / load(dir)

`CharTokenizer` is one-token-per-character. `BPETokenizer` is the classic
Sennrich byte-pair-encoding learned *from scratch* over characters, and it
records every merge it makes (so the UI can let you "watch it merge").

Nothing here uses the HuggingFace `tokenizers` library on purpose — the point is
to see the algorithm.
"""

from __future__ import annotations

import json
import os
import re
from collections import Counter
from typing import Dict, List, Tuple

# Pre-tokenization: split text into "pieces" before BPE so merges never cross
# word boundaries. This is a *reversible* partition — "".join(pieces) == text —
# using the GPT-2 leading-space convention (a single ASCII space attaches to the
# following word, which is why real tokenizers show tokens like " the"). \w and
# \s are Unicode-aware in Python 3, so accented letters are handled too.
_PRETOK = re.compile(r" ?\w+| ?[^\w\s]+|\s+", re.UNICODE)

UNK = "<unk>"  # id 0 in every tokenizer; only produced for chars unseen in training


def pretokenize(text: str) -> List[str]:
    """Reversible split into word / punctuation / whitespace pieces."""
    return _PRETOK.findall(text)


def pretokenize_with_offsets(text: str) -> List[Tuple[str, int]]:
    """Like pretokenize but also returns each piece's start offset in `text`."""
    out = []
    for m in _PRETOK.finditer(text):
        out.append((m.group(0), m.start()))
    return out


# --------------------------------------------------------------------------- #
# Character-level
# --------------------------------------------------------------------------- #
class CharTokenizer:
    kind = "char"

    def __init__(self, itos: List[str]):
        # itos[0] is always UNK; the rest are single characters.
        self.itos = itos
        self.stoi = {s: i for i, s in enumerate(itos)}

    @classmethod
    def train(cls, text: str) -> "CharTokenizer":
        chars = sorted(set(text))
        return cls([UNK] + chars)

    @property
    def vocab_size(self) -> int:
        return len(self.itos)

    def encode(self, text: str) -> List[int]:
        get = self.stoi.get
        return [get(ch, 0) for ch in text]

    def decode(self, ids: List[int]) -> str:
        itos = self.itos
        return "".join("" if i == 0 else itos[i] for i in ids if 0 <= i < len(itos))

    def tokens_with_spans(self, text: str) -> List[dict]:
        out = []
        for i, ch in enumerate(text):
            out.append({"text": ch, "id": self.stoi.get(ch, 0), "start": i, "end": i + 1})
        return out

    def token_str(self, idx: int) -> str:
        return UNK if idx == 0 else self.itos[idx]

    # symmetry with BPETokenizer (char has no merges)
    @property
    def merges_history(self) -> List[dict]:
        return []

    def to_meta(self) -> dict:
        return {"kind": self.kind, "vocab_size": self.vocab_size, "num_merges": 0}

    def save(self, path: str) -> None:
        os.makedirs(path, exist_ok=True)
        with open(os.path.join(path, "tokenizer.json"), "w", encoding="utf-8") as f:
            json.dump({"kind": self.kind, "itos": self.itos}, f, ensure_ascii=False)
        # keep a (empty) merge history so the loader is uniform
        with open(os.path.join(path, "merges_history.json"), "w", encoding="utf-8") as f:
            json.dump([], f)

    @classmethod
    def load(cls, path: str) -> "CharTokenizer":
        with open(os.path.join(path, "tokenizer.json"), encoding="utf-8") as f:
            data = json.load(f)
        return cls(data["itos"])


# --------------------------------------------------------------------------- #
# Byte-pair encoding (from scratch)
# --------------------------------------------------------------------------- #
class BPETokenizer:
    kind = "bpe"

    def __init__(self, itos: List[str], merges: List[Tuple[str, str]],
                 merges_history: List[dict] | None = None):
        self.itos = itos
        self.stoi = {s: i for i, s in enumerate(itos)}
        self.merges = [tuple(m) for m in merges]          # learned order
        self.merge_rank = {tuple(m): r for r, m in enumerate(self.merges)}
        self._history = merges_history or []
        self._piece_cache: Dict[str, List[str]] = {}

    # ----- training ----------------------------------------------------- #
    @classmethod
    def train(cls, text: str, vocab_size: int, verbose: bool = False) -> "BPETokenizer":
        """Learn `vocab_size`-ish tokens of BPE from `text`, recording each merge."""
        # 1. pre-tokenize and count unique pieces (BPE works on word *types*,
        #    weighted by frequency — far fewer than the raw token stream).
        piece_freq = Counter(pretokenize(text))

        # 2. each piece starts as a tuple of its characters
        words: Dict[Tuple[str, ...], int] = {}
        base_chars = set()
        for piece, freq in piece_freq.items():
            syms = tuple(piece)
            words[syms] = words.get(syms, 0) + freq
            base_chars.update(syms)

        itos = [UNK] + sorted(base_chars)
        stoi = {s: i for i, s in enumerate(itos)}
        num_merges = max(0, vocab_size - len(itos))

        merges: List[Tuple[str, str]] = []
        history: List[dict] = []

        for step in range(num_merges):
            # count every adjacent symbol pair, weighted by piece frequency
            pairs: Counter = Counter()
            for syms, freq in words.items():
                for a, b in zip(syms, syms[1:]):
                    pairs[(a, b)] += freq
            if not pairs:
                break
            # most frequent pair; deterministic tie-break on the pair itself
            best, freq = max(pairs.items(), key=lambda kv: (kv[1], kv[0]))
            new_tok = best[0] + best[1]

            # apply the merge everywhere
            words = _merge_words(words, best, new_tok)

            stoi[new_tok] = len(itos)
            itos.append(new_tok)
            merges.append(best)
            history.append({
                "rank": step,
                "pair": [best[0], best[1]],
                "token": new_tok,
                "freq": freq,
                "vocab_size": len(itos),
            })
            if verbose and (step % 200 == 0 or step == num_merges - 1):
                print(f"  merge {step:4d}: {best[0]!r}+{best[1]!r} -> {new_tok!r} "
                      f"(freq {freq}, vocab {len(itos)})")

        return cls(itos, merges, history)

    # ----- encoding ----------------------------------------------------- #
    def _bpe_piece(self, piece: str) -> List[str]:
        """Apply learned merges, in learned order, to one pre-tokenized piece."""
        cached = self._piece_cache.get(piece)
        if cached is not None:
            return cached
        syms = list(piece)
        if len(syms) >= 2:
            for a, b in self.merges:
                if len(syms) < 2:
                    break
                if a not in syms:           # cheap skip
                    continue
                syms = _merge_seq(syms, a, b)
        self._piece_cache[piece] = syms
        return syms

    def encode(self, text: str) -> List[int]:
        get = self.stoi.get
        out: List[int] = []
        for piece in pretokenize(text):
            for sym in self._bpe_piece(piece):
                out.append(get(sym, 0))
        return out

    def decode(self, ids: List[int]) -> str:
        itos = self.itos
        return "".join("" if i == 0 else itos[i] for i in ids if 0 <= i < len(itos))

    def tokens_with_spans(self, text: str) -> List[dict]:
        out = []
        for piece, start in pretokenize_with_offsets(text):
            off = start
            for sym in self._bpe_piece(piece):
                out.append({"text": sym, "id": self.stoi.get(sym, 0),
                            "start": off, "end": off + len(sym)})
                off += len(sym)
        return out

    def token_str(self, idx: int) -> str:
        return UNK if idx == 0 else self.itos[idx]

    # ----- the "watch it merge" trace ----------------------------------- #
    def bpe_trace(self, word: str) -> List[dict]:
        """Replay training on a single word, capturing every state it passes
        through as merges are applied in the order they were learned. Returns
        only the change-points, so each entry is a real step in the collapse."""
        syms = list(word)
        steps = [{"step": 0, "applied": None, "tokens": list(syms)}]
        for rank, (a, b) in enumerate(self.merges):
            if len(syms) < 2:
                break
            if a not in syms:
                continue
            merged = _merge_seq(syms, a, b)
            if merged != syms:
                syms = merged
                steps.append({"step": rank + 1, "applied": [a, b], "tokens": list(syms)})
        return steps

    @property
    def merges_history(self) -> List[dict]:
        return self._history

    def to_meta(self) -> dict:
        return {"kind": self.kind, "vocab_size": self.vocab_size,
                "num_merges": len(self.merges)}

    @property
    def vocab_size(self) -> int:
        return len(self.itos)

    # ----- persistence -------------------------------------------------- #
    def save(self, path: str) -> None:
        os.makedirs(path, exist_ok=True)
        with open(os.path.join(path, "tokenizer.json"), "w", encoding="utf-8") as f:
            json.dump({"kind": self.kind, "itos": self.itos,
                       "merges": [list(m) for m in self.merges]},
                      f, ensure_ascii=False)
        with open(os.path.join(path, "merges_history.json"), "w", encoding="utf-8") as f:
            json.dump(self._history, f, ensure_ascii=False)

    @classmethod
    def load(cls, path: str) -> "BPETokenizer":
        with open(os.path.join(path, "tokenizer.json"), encoding="utf-8") as f:
            data = json.load(f)
        history = []
        hp = os.path.join(path, "merges_history.json")
        if os.path.exists(hp):
            with open(hp, encoding="utf-8") as f:
                history = json.load(f)
        return cls(data["itos"], data["merges"], history)


# --------------------------------------------------------------------------- #
# helpers
# --------------------------------------------------------------------------- #
def _merge_seq(syms: List[str], a: str, b: str) -> List[str]:
    """Merge every adjacent (a, b) in a symbol list into a+b."""
    merged = a + b
    out: List[str] = []
    i, n = 0, len(syms)
    while i < n:
        if i < n - 1 and syms[i] == a and syms[i + 1] == b:
            out.append(merged)
            i += 2
        else:
            out.append(syms[i])
            i += 1
    return out


def _merge_words(words: Dict[Tuple[str, ...], int], pair: Tuple[str, str],
                 new_tok: str) -> Dict[Tuple[str, ...], int]:
    """Apply a merge across the whole word->freq table."""
    a, b = pair
    out: Dict[Tuple[str, ...], int] = {}
    for syms, freq in words.items():
        if a in syms:
            syms = tuple(_merge_seq(list(syms), a, b))
        out[syms] = out.get(syms, 0) + freq
    return out


# --------------------------------------------------------------------------- #
# loader dispatch + factory
# --------------------------------------------------------------------------- #
def load_tokenizer(path: str):
    with open(os.path.join(path, "tokenizer.json"), encoding="utf-8") as f:
        kind = json.load(f)["kind"]
    return CharTokenizer.load(path) if kind == "char" else BPETokenizer.load(path)


def build_tokenizer(scheme: dict, text: str, verbose: bool = False):
    """Build a tokenizer from a scheme dict (see corpora.yaml `schemes`)."""
    if scheme["kind"] == "char":
        return CharTokenizer.train(text)
    return BPETokenizer.train(text, scheme["vocab_size"], verbose=verbose)


# --------------------------------------------------------------------------- #
# self-test:  python src/tokenizer.py
# --------------------------------------------------------------------------- #
if __name__ == "__main__":
    sample = (
        "the quick brown fox jumps over the lazy dog. "
        "The QUICK brown Fox! thoughts, thinking, thinkers think.\n"
        "low lower lowest newer newest wider widest.\n"
    ) * 40

    print("== pre-tokenization is reversible ==")
    for t in ["Hello, world!", "  spaced   out\n\ttabs", "To be, or not to be",
              "café déjà", sample[:120]]:
        assert "".join(pretokenize(t)) == t, repr(t)
    print("ok")

    print("\n== char tokenizer ==")
    ct = CharTokenizer.train(sample)
    probe = "the lazy dog."
    assert ct.decode(ct.encode(probe)) == probe
    spans = ct.tokens_with_spans(probe)
    assert "".join(s["text"] for s in spans) == probe
    assert len(ct.encode(probe)) == len(probe)
    print(f"vocab={ct.vocab_size}  '{probe}' -> {len(ct.encode(probe))} tokens  (roundtrip ok)")

    for target in (512, 2048):
        print(f"\n== BPE-{target} ==")
        bpe = BPETokenizer.train(sample, target, verbose=True)
        # roundtrip on held-back-ish text
        for probe in ["the quick brown fox", "thinking thoughts", "lowest newer\n",
                      "café", "Unseen ZZZ chars ~`"]:
            dec = bpe.decode(bpe.encode(probe))
            # unseen chars (Z, ~, `, é if absent) map to unk and drop; only assert
            # roundtrip when every char was in the training vocab
            if all(ch in bpe.stoi for ch in probe):
                assert dec == probe, (probe, dec)
        # spans align
        sp = bpe.tokens_with_spans("the quick brown fox")
        assert "".join(s["text"] for s in sp) == "the quick brown fox"
        # vocab + history invariants
        assert len(bpe.merges_history) == len(bpe.merges)
        assert bpe.vocab_size <= target
        # BPE is no longer than char on the same text
        n_char = len(ct.encode("the quick brown fox jumps over the lazy dog"))
        n_bpe = len(bpe.encode("the quick brown fox jumps over the lazy dog"))
        print(f"vocab={bpe.vocab_size}  merges={len(bpe.merges)}  "
              f"'fox...dog': char={n_char} vs bpe={n_bpe} tokens")
        assert n_bpe <= n_char
        # watch-it-merge trace
        trace = bpe.bpe_trace("thoughts")
        print(f"  trace 'thoughts': {len(trace)} change-points; "
              f"final={trace[-1]['tokens']}")
        assert trace[0]["tokens"] == list("thoughts")
        assert trace[-1]["tokens"] == bpe._bpe_piece("thoughts")

    print("\nAll tokenizer self-tests passed.")