"""Train a SentencePiece tokenizer on a Chinese Wikipedia subset.

用法:
    uv run python kda/training/train_tokenizer.py \
        --out data/spm_4k --vocab-size 4096 --limit 20000

产出:
    data/spm_4k.model / data/spm_4k.vocab  (BPE/unigram, 中文小语料)
"""
from __future__ import annotations

import argparse

import sentencepiece as spm

from .data import fetch_wiki_texts


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--out", default="data/spm_4k", help="输出前缀 (model/vocab 文件)")
    p.add_argument("--vocab-size", type=int, default=8192)
    p.add_argument("--limit", type=int, default=20000, help="用于训练的 wiki 文章数")
    p.add_argument("--model-type", default="unigram", choices=["unigram", "bpe"])
    p.add_argument("--character-coverage", type=float, default=0.9995)
    args = p.parse_args()

    texts = fetch_wiki_texts(args.limit)
    corpus = "".join(texts)
    tmp = args.out + ".corpus.txt"
    with open(tmp, "w", encoding="utf-8") as f:
        f.write(corpus)
    print(f"corpus: {len(corpus):,} chars from {len(texts)} articles")

    spm.SentencePieceTrainer.train(
        input=tmp,
        model_prefix=args.out,
        vocab_size=args.vocab_size,
        model_type=args.model_type,
        character_coverage=args.character_coverage,
        unk_id=0,
        pad_id=1,
        bos_id=-1,
        eos_id=-1,
        num_threads=4,
    )
    print(f"tokenizer saved: {args.out}.model / {args.out}.vocab")


if __name__ == "__main__":
    main()