"""Generate BPE tokenizer trained on textbook corpus.

Vocab layout: 0-50 special tokens, 51-306 byte-level chars, 307-4095 BPE merges.
"""
import json
import os
from tokenizers import Tokenizer, models, pre_tokenizers, trainers, decoders

TOKENIZER_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "tokenizer")
BOOKS_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "books")

os.makedirs(TOKENIZER_DIR, exist_ok=True)

special_tokens = [
    "<unk>", "<s>", "</s>", "<pad>",
    "<|system|>", "<|user|>", "<|assistant|>",
    "<think>", "</think>",
    "[INST]", "[/INST]",
    "<|begin_of_thought|>", "<|end_of_thought|>",
    "<|reflect|>", "<|revise|>", "<|verify|>",
    "<|code|>", "<|text|>", "<|math|>",
    "<|think|>", "<|answer|>", "<|step|>", "<|reason|>",
    "<|check|>", "<|output|>", "<|plan|>", "<|solve|>",
    "<|analyze|>", "<|conclude|>", "<|approach|>", "<|alternative|>",
    "<|summary|>", "<|question|>", "<|hint|>", "<|example|>",
    "<|correct|>", "<|incorrect|>", "<|feedback|>",
    "<|start|>", "<|end|>", "<|sep|>", "<|cls|>",
    "<|tool|>", "<|function|>", "<|result|>", "<|input|>",
    "<|detect|>", "<|context|>", "<|proof|>", "<|lemma|>", "<|theorem|>",
]

special_map = {t: i for i, t in enumerate(special_tokens)}
NUM_SPECIAL = len(special_tokens)

def find_book_files():
    if not os.path.isdir(BOOKS_DIR):
        print(f"Warning: {BOOKS_DIR} not found, using empty corpus")
        return []
    files = []
    for fname in os.listdir(BOOKS_DIR):
        if fname.endswith(".txt"):
            fpath = os.path.join(BOOKS_DIR, fname)
            if os.path.getsize(fpath) > 100:
                files.append(fpath)
    files.sort()
    print(f"Found {len(files)} book files in {BOOKS_DIR}")
    return files

def build_tokenizer():
    # Create BPE tokenizer with ByteLevel pre-tokenizer
    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tokenizer.decoder = decoders.ByteLevel(add_prefix_space=False)
    tokenizer.add_special_tokens(special_tokens)

    book_files = find_book_files()

    print(f"Training BPE with vocab_size=4096, {len(book_files)} files...")
    trainer = trainers.BpeTrainer(
        vocab_size=4096,
        min_frequency=2,
        special_tokens=special_tokens,
        show_progress=True,
        initial_alphabet=[],
    )
    tokenizer.train(book_files, trainer)
    print("BPE training done.")

    output_path = os.path.join(TOKENIZER_DIR, "tokenizer.json")
    tokenizer.save(output_path)
    print(f"Saved to {output_path}")

    # Post-process: ensure byte_fallback=True
    with open(output_path) as f:
        data = json.load(f)

    data["model"]["byte_fallback"] = True
    data["model"]["dropout"] = None

    with open(output_path, "w") as f:
        json.dump(data, f, ensure_ascii=False)

    # Verify
    with open(output_path) as f:
        data = json.load(f)
    vocab = data["model"]["vocab"]
    merges = data["model"]["merges"]
    sorted_vocab = sorted(vocab.items(), key=lambda x: x[1])
    print(f"Vocab size: {len(vocab)}")
    print(f"Merges: {len(merges)}")
    print(f"First tokens: {sorted_vocab[:7]}")
    print(f"Tokens 50-55: {sorted_vocab[50:55]}")
    print(f"Last tokens: {sorted_vocab[-5:]}")
    if merges:
        print(f"First merges: {merges[:5]}")

if __name__ == "__main__":
    build_tokenizer()