"""Verify that the exported HF tokenizer matches the training-time tokenizer wrapper."""

from pathlib import Path
import sys

import sentencepiece as spm

from tokenization_taonet import TaoNetTokenizer


SAMPLES = [
    "Explain why compact language models can still be useful.",
    "Fruit is now expensive so we should",
    "Hello world",
    "<user>",
    "<assistant>",
    "\n",
]


def main():
    repo_dir = Path(__file__).resolve().parent
    sys.path.insert(0, str(repo_dir / "src"))
    from taoTrain.data.tokenizer import SentencePieceTokenizerWrapper, load_special_token_metadata

    tokenizer_model = repo_dir / "tokenizer" / "tokenizer.model"
    if not tokenizer_model.exists():
        tokenizer_model = repo_dir / "tokenizer.model"

    sp = spm.SentencePieceProcessor()
    sp.Load(str(tokenizer_model))
    special_token_ids = load_special_token_metadata(tokenizer_model)

    train_tokenizer = SentencePieceTokenizerWrapper(sp, special_token_ids=special_token_ids)
    hf_tokenizer = TaoNetTokenizer.from_pretrained(str(repo_dir))

    print(f"train vocab_size: {train_tokenizer.vocab_size}")
    print(f"hf vocab_size:    {hf_tokenizer.vocab_size}")

    for token in ["<UNK>", "<BOS>", "<EOS>", "<PAD>", "<think>", "<user>", "<assistant>", "<image>", "\n"]:
        train_id = train_tokenizer.get_special_token_id(token)
        hf_id = hf_tokenizer.get_special_token_id(token)
        print(f"{token!r}: train={train_id}, hf={hf_id}")
        if train_id != hf_id:
            raise SystemExit(f"Special token mismatch for {token}: train={train_id}, hf={hf_id}")

    print("\nChecking ID -> token mapping...")
    for token_id in range(sp.vocab_size()):
        train_piece = sp.id_to_piece(token_id)
        hf_piece = hf_tokenizer._convert_id_to_token(token_id)
        if token_id in special_token_ids.values():
            expected = next(token for token, value in special_token_ids.items() if value == token_id)
            if hf_piece != expected:
                raise SystemExit(
                    f"HF id->token mismatch at id={token_id}: expected special token {expected!r}, got {hf_piece!r}"
                )
        else:
            if hf_piece != train_piece:
                raise SystemExit(
                    f"HF id->token mismatch at id={token_id}: train={train_piece!r}, hf={hf_piece!r}"
                )

    print("ID -> token mapping matches.")

    print("\nChecking sample encodes/decodes...")
    for sample in SAMPLES:
        train_ids = train_tokenizer(sample, return_attention_mask=True)
        hf_ids = hf_tokenizer(sample, return_attention_mask=True)
        print(f"sample: {sample!r}")
        print(f"  train ids: {train_ids['input_ids']}")
        print(f"  hf ids:    {hf_ids['input_ids']}")
        if train_ids["input_ids"] != hf_ids["input_ids"]:
            raise SystemExit(f"Encoding mismatch for sample {sample!r}")

        train_decoded = train_tokenizer.decode(train_ids["input_ids"], skip_special_tokens=True)
        hf_decoded = hf_tokenizer.decode(hf_ids["input_ids"], skip_special_tokens=True)
        print(f"  train decode: {train_decoded!r}")
        print(f"  hf decode:    {hf_decoded!r}")
        if train_decoded != hf_decoded:
            raise SystemExit(f"Decode mismatch for sample {sample!r}")

    prompt = "Explain why compact language models can still be useful."
    chat_ids = [
        train_tokenizer.get_special_token_id("<user>"),
        *train_tokenizer(prompt)["input_ids"],
        train_tokenizer.get_special_token_id("<assistant>"),
    ]
    hf_chat = hf_tokenizer.build_chat_inputs(prompt)
    print("\nChecking chat prompt construction...")
    print(f"  train-style chat ids: {chat_ids}")
    print(f"  hf chat ids:          {hf_chat['input_ids']}")
    if chat_ids != hf_chat["input_ids"]:
        raise SystemExit("Chat prompt IDs do not match training-time construction.")

    print("\nTokenizer verification passed.")


if __name__ == "__main__":
    main()