# ─────────────────────────────────────────────────────────
# EGen Labs | ErebusTN
# AXIOM-1 — EGen-Core/AXIOM-1
# File: axiom1/tokenizer/build_tokenizer.py
# ─────────────────────────────────────────────────────────

"""
Script to train the AXIOM-1 tokenizer from a text corpus or HuggingFace dataset.

Usage:
    # From a text file:
    python -m axiom1.tokenizer.build_tokenizer --source corpus.txt

    # From the HuggingFace dataset:
    python -m axiom1.tokenizer.build_tokenizer --hf-dataset ErebusTN/axiom1

    # From multiple files:
    python -m axiom1.tokenizer.build_tokenizer --source file1.txt file2.txt
"""

import argparse
from pathlib import Path

from axiom1.tokenizer.tokenizer import AXIOM1Tokenizer


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Train the AXIOM-1 byte-level BPE tokenizer"
    )
    parser.add_argument(
        "--source", nargs="+", type=str, default=None,
        help="Path(s) to text files to train on"
    )
    parser.add_argument(
        "--hf-dataset", type=str, default=None,
        help="HuggingFace dataset name (e.g., ErebusTN/axiom1)"
    )
    parser.add_argument(
        "--text-column", type=str, default="text",
        help="Column name containing text in the HF dataset"
    )
    parser.add_argument(
        "--save-dir", type=str, default="axiom1/tokenizer",
        help="Directory to save the trained tokenizer"
    )
    parser.add_argument(
        "--vocab-size", type=int, default=132000,
        help="Target vocabulary size (default: 132000)"
    )
    args = parser.parse_args()

    if args.hf_dataset:
        from datasets import load_dataset
        print(f"[AXIOM-1 Tokenizer] Loading dataset {args.hf_dataset}...")
        ds = load_dataset(args.hf_dataset, split="train")
        texts = ds[args.text_column]
        print(f"[AXIOM-1 Tokenizer] Training on {len(texts)} examples...")
        AXIOM1Tokenizer.train(
            texts=texts,
            save_dir=args.save_dir,
            vocab_size=args.vocab_size,
        )
    elif args.source:
        files = [Path(f) for f in args.source]
        for f in files:
            if not f.exists():
                raise FileNotFoundError(f"File not found: {f}")
        print(f"[AXIOM-1 Tokenizer] Training on {len(files)} file(s)...")
        AXIOM1Tokenizer.train(
            files=files,
            save_dir=args.save_dir,
            vocab_size=args.vocab_size,
        )
    else:
        raise ValueError("Provide either --source or --hf-dataset")


if __name__ == "__main__":
    main()
