"""Train a ByteLevel BPE tokenizer from UTF-8 .txt files."""
import argparse
from pathlib import Path
from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers, processors

SPECIAL = ["<|pad|>", "<|unk|>", "<|bos|>", "<|eos|>", "<|system|>", "<|user|>", "<|assistant|>", "<|end|>"]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", nargs="+", required=True, help="Text files or folders containing .txt")
    ap.add_argument("--output", default="tokenizer.json")
    ap.add_argument("--vocab-size", type=int, default=32000)
    args = ap.parse_args()
    files = []
    for f in args.input:
        p = Path(f)
        files.extend(sorted(p.rglob("*.txt")) if p.is_dir() else [p])
    if not files:
        raise ValueError("No .txt files found")
    tok = Tokenizer(models.BPE(unk_token="<|unk|>"))
    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tok.decoder = decoders.ByteLevel()
    tok.train([str(f) for f in files], trainers.BpeTrainer(
        vocab_size=args.vocab_size, min_frequency=2,
        special_tokens=SPECIAL, initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
        show_progress=True,
    ))
    # For tiny corpora the trainer may produce fewer tokens than requested.
    if tok.get_vocab_size() != args.vocab_size:
        raise ValueError(f"Tokenizer has {tok.get_vocab_size()} tokens instead of {args.vocab_size}. Add more text or adjust config.json vocab_size.")
    tok.save(args.output)
    print(f"Saved {args.output}; vocab={tok.get_vocab_size()}")


if __name__ == "__main__":
    main()