#!/usr/bin/env python3
"""Train SentencePiece model and save a Hugging Face–compatible fast tokenizer."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from weather_llm.config_loader import repo_root_from
from weather_llm.tokenization.hf_tokenizer import load_llama_tokenizer_from_spm
from weather_llm.tokenization.spm_trainer import train_sentencepiece
from weather_llm.tokenization.validate import validate_tokenizer, write_validation_report


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--corpus", type=Path, required=True, help="JSONL with `text` field")
    ap.add_argument("--out_dir", type=Path, default=None)
    ap.add_argument("--vocab_size", type=int, default=32000)
    args = ap.parse_args()
    root = repo_root_from(ROOT)
    corpus = args.corpus if args.corpus.is_absolute() else (root / args.corpus)
    out_dir = args.out_dir or (root / "artifacts/tokenizer")
    out_dir.mkdir(parents=True, exist_ok=True)
    prefix = out_dir / "weather_spm"
    model_file = train_sentencepiece(
        corpus,
        prefix,
        vocab_size=args.vocab_size,
        sample_docs=None,
    )
    tok = load_llama_tokenizer_from_spm(model_file)
    tok.save_pretrained(str(out_dir))
    report = validate_tokenizer(tok)
    write_validation_report(out_dir / "validation_report.txt", report)
    if report["issues"]:
        print("Tokenizer validation warnings:")
        for i in report["issues"]:
            print(" -", i)
    print("Saved tokenizer to", out_dir)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())