#!/usr/bin/env python3
"""
Test_Tokenizers.py - on-device validation for the tokenizers wheel.

Exercises the Rust/PyO3 binding: import, version, BPE train on a tiny
corpus + encode/decode roundtrip, models/normalizers/pre-tokenizers.
Exit code 0 = all tests passed, 1 = any FAIL.

Generated by RIMI
"""
import sys

RESULTS = []


def test(name, fn):
    try:
        fn()
        RESULTS.append(("PASS", name))
    except NotImplementedError:
        RESULTS.append(("SKIP", name))
    except Exception as e:
        RESULTS.append(("FAIL", name, str(e)))


def section(title):
    print("\n===== %s =====" % title)


def check(cond, msg):
    if not cond:
        raise AssertionError(msg)


# ---------------------------------------------------------------------------
# 1. imports + versions
# ---------------------------------------------------------------------------
def test_import_tokenizers():
    import tokenizers
    check(hasattr(tokenizers, "__version__"), "no __version__")
    print("  tokenizers version:", tokenizers.__version__)
    check(tokenizers.__version__ == "0.23.2", "version != 0.23.2")


def test_import_submodules():
    import tokenizers.models
    import tokenizers.trainers
    import tokenizers.pre_tokenizers
    import tokenizers.normalizers
    import tokenizers.processors
    import tokenizers.decoders
    print("  submodules: models/trainers/pre_tokenizers/normalizers/processors/decoders OK")


def test_import_tokenizer_class():
    from tokenizers import Tokenizer
    check(callable(Tokenizer), "Tokenizer not callable")
    print("  Tokenizer class OK")


# ---------------------------------------------------------------------------
# 2. BPE train on tiny corpus + encode/decode roundtrip
# ---------------------------------------------------------------------------
_TINY_CORPUS = [
    "Hello world, this is a test.",
    "Tokenizers are fast and versatile.",
    "Hello again, another test sentence.",
    "BPE training on a tiny corpus.",
    "The quick brown fox jumps over the lazy dog.",
]

_TRAIN_FILES = ["/tmp/tok_train.txt"]


def _write_corpus():
    # Scripts dir on device is writable; fall back to current dir
    import os
    for cand in ("/tmp/tok_train.txt", "tok_train.txt"):
        try:
            with open(cand, "w", encoding="utf-8") as fh:
                for line in _TINY_CORPUS:
                    fh.write(line + "\n")
            return cand
        except OSError:
            continue
    raise AssertionError("cannot write training corpus")


def test_bpe_train():
    from tokenizers import Tokenizer
    from tokenizers.models import BPE
    from tokenizers.trainers import BpeTrainer
    from tokenizers.pre_tokenizers import Whitespace

    path = _write_corpus()
    tok = Tokenizer(BPE(unk_token="[UNK]"))
    tok.pre_tokenizer = Whitespace()
    trainer = BpeTrainer(vocab_size=200, special_tokens=["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"])
    tok.train([path], trainer)
    vs = tok.get_vocab_size()
    check(vs > 0, "vocab size 0")
    print("  BPE trained, vocab size:", vs)


def test_encode_decode_roundtrip():
    from tokenizers import Tokenizer
    from tokenizers.models import BPE
    from tokenizers.trainers import BpeTrainer
    from tokenizers.pre_tokenizers import Whitespace

    path = _write_corpus()
    tok = Tokenizer(BPE(unk_token="[UNK]"))
    tok.pre_tokenizer = Whitespace()
    trainer = BpeTrainer(vocab_size=200, special_tokens=["[UNK]"])
    tok.train([path], trainer)
    text = "Hello world, BPE roundtrip test."
    enc = tok.encode(text)
    check(len(enc.ids) > 0, "no ids")
    check(len(enc.tokens) > 0, "no tokens")
    dec = tok.decode(enc.ids)
    check(isinstance(dec, str) and len(dec) > 0, "empty decode")
    # roundtrip: decoded text must contain the key words (whitespace split)
    check("Hello" in dec, "roundtrip lost 'Hello': %r" % dec)
    print("  ids:", enc.ids[:10])
    print("  tokens:", enc.tokens[:10])
    print("  decoded:", dec)


def test_encode_batch():
    from tokenizers import Tokenizer
    from tokenizers.models import BPE
    from tokenizers.trainers import BpeTrainer
    from tokenizers.pre_tokenizers import Whitespace

    path = _write_corpus()
    tok = Tokenizer(BPE(unk_token="[UNK]"))
    tok.pre_tokenizer = Whitespace()
    tok.train([path], BpeTrainer(vocab_size=200, special_tokens=["[UNK]"]))
    encs = tok.encode_batch(_TINY_CORPUS[:3])
    check(len(encs) == 3, "batch len")
    check(all(len(e.ids) > 0 for e in encs), "empty batch ids")
    print("  batch ok:", [len(e.ids) for e in encs])


# ---------------------------------------------------------------------------
# 3. WordLevel + save/load roundtrip
# ---------------------------------------------------------------------------
def test_wordlevel():
    from tokenizers import Tokenizer
    from tokenizers.models import WordLevel
    from tokenizers.pre_tokenizers import WhitespaceSplit

    tok = Tokenizer(WordLevel(vocab={"hello": 0, "world": 1, "[UNK]": 2}, unk_token="[UNK]"))
    tok.pre_tokenizer = WhitespaceSplit()
    enc = tok.encode("hello world")
    check(enc.ids == [0, 1], "wordlevel ids %r" % (enc.ids,))
    print("  WordLevel ids:", enc.ids)


def test_save_load():
    import os
    import tempfile
    from tokenizers import Tokenizer
    from tokenizers.models import BPE
    from tokenizers.trainers import BpeTrainer
    from tokenizers.pre_tokenizers import Whitespace

    path = _write_corpus()
    tok = Tokenizer(BPE(unk_token="[UNK]"))
    tok.pre_tokenizer = Whitespace()
    tok.train([path], BpeTrainer(vocab_size=200, special_tokens=["[UNK]"]))
    tmpd = tempfile.mkdtemp()
    fp = os.path.join(tmpd, "tok.json")
    tok.save(fp)
    check(os.path.isfile(fp), "save missing")
    tok2 = Tokenizer.from_file(fp)
    check(tok2.get_vocab_size() == tok.get_vocab_size(), "vocab mismatch after load")
    print("  save/load vocab:", tok2.get_vocab_size())


# ---------------------------------------------------------------------------
# 4. normalizers / pre-tokenizers / processors / decoders
# ---------------------------------------------------------------------------
def test_normalizer():
    from tokenizers import Tokenizer
    from tokenizers.models import WordLevel
    from tokenizers.normalizers import Lowercase
    from tokenizers.pre_tokenizers import Whitespace

    tok = Tokenizer(WordLevel(vocab={"hello": 0, "world": 1, "[UNK]": 2}, unk_token="[UNK]"))
    tok.normalizer = Lowercase()
    tok.pre_tokenizer = Whitespace()
    enc = tok.encode("HELLO WORLD")
    check(enc.ids == [0, 1], "lowercase ids %r" % (enc.ids,))
    print("  Lowercase normalizer OK")


def test_bert_processing():
    from tokenizers import Tokenizer
    from tokenizers.models import WordPiece
    from tokenizers.processors import BertProcessing

    tok = Tokenizer(WordPiece(vocab={"hello": 0, "world": 1, "[UNK]": 2, "[CLS]": 3, "[SEP]": 4}, unk_token="[UNK]"))
    tok.post_processor = BertProcessing(("[SEP]", 4), ("[CLS]", 3))
    enc = tok.encode("hello world")
    check(enc.ids[0] == 3 and enc.ids[-1] == 4, "bert ids %r" % (enc.ids,))
    print("  BertProcessing ids:", enc.ids)


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------
def main():
    section("1. imports + versions")
    test("import tokenizers", test_import_tokenizers)
    test("import submodules", test_import_submodules)
    test("Tokenizer class", test_import_tokenizer_class)

    section("2. BPE train + roundtrip")
    test("BPE train tiny corpus", test_bpe_train)
    test("encode/decode roundtrip", test_encode_decode_roundtrip)
    test("encode_batch", test_encode_batch)

    section("3. models + serialization")
    test("WordLevel", test_wordlevel)
    test("save/load", test_save_load)

    section("4. pipeline pieces")
    test("Lowercase normalizer", test_normalizer)
    test("BertProcessing", test_bert_processing)

    section("RESULT")
    n_ok = n_fail = n_skip = 0
    for r in RESULTS:
        status = r[0]
        if status == "PASS":
            n_ok += 1
            print("  OK    %s" % r[1])
        elif status == "SKIP":
            n_skip += 1
            print("  SKIP  %s" % r[1])
        else:
            n_fail += 1
            print("  FAIL  %s: %s" % (r[1], r[2]))
    print("RESULT: %d ok, %d failed, %d skipped" % (n_ok, n_fail, n_skip))
    sys.exit(1 if n_fail else 0)


if __name__ == "__main__":
    main()