"""Dataset + corpus-encoding tests on synthetic data (no real corpus needed).

Run directly: `.venv/bin/python tests/test_dataset.py`
"""

import sys
import tempfile
from pathlib import Path

import numpy as np
import torch

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "src"))
sys.path.insert(0, str(ROOT / "scripts"))

from dataset import TOKEN_DTYPE, get_batch, load_tokens  # noqa: E402
from tokenizer import BPETokenizer  # noqa: E402
from encode_corpus import encode_file  # noqa: E402

EOT = "<|endoftext|>"


def test_get_batch_shapes_and_shift():
    # Tokens 0..999 in order, so the shift property is easy to verify.
    data = np.arange(1000, dtype=TOKEN_DTYPE)
    gen = torch.Generator().manual_seed(0)
    x, y = get_batch(data, batch_size=8, context_length=16, device="cpu", generator=gen)
    assert x.shape == (8, 16) and y.shape == (8, 16)
    assert x.dtype == torch.long
    # y is x shifted by one; with sequential data each y == x + 1.
    assert torch.equal(y, x + 1)


def test_get_batch_is_reproducible():
    data = np.arange(1000, dtype=TOKEN_DTYPE)
    a = get_batch(data, 4, 8, "cpu", torch.Generator().manual_seed(42))
    b = get_batch(data, 4, 8, "cpu", torch.Generator().manual_seed(42))
    assert torch.equal(a[0], b[0]) and torch.equal(a[1], b[1])


def test_get_batch_in_bounds():
    data = np.arange(50, dtype=TOKEN_DTYPE)
    x, y = get_batch(data, 32, 8, "cpu", torch.Generator().manual_seed(1))
    assert int(x.max()) < 50 and int(y.max()) < 50


def _tiny_tokenizer() -> BPETokenizer:
    corpus = ("The cat sat. Tom ran fast. Lily played.\n" + EOT + "\n") * 40
    return BPETokenizer.train(corpus, vocab_size=350, special_tokens=[EOT])


def test_encode_file_roundtrip_and_separators():
    tok = _tiny_tokenizer()
    eot_id = tok.special_tokens[EOT]
    with tempfile.TemporaryDirectory() as d:
        src = Path(d) / "mini.txt"
        dst = Path(d) / "mini.bin"
        # Two stories, separator on its own line, with stray blank lines.
        src.write_text(
            "The cat sat.\nTom ran fast.\n" + EOT + "\n\n"
            "Lily played.\n" + EOT + "\n"
        )
        n_stories, n_tokens = encode_file(tok, src, dst, progress=False)

        assert n_stories == 2
        ids = load_tokens(dst).tolist()
        assert len(ids) == n_tokens
        # Exactly one EOT per story, and each story ends with one.
        assert ids.count(eot_id) == 2
        assert ids[-1] == eot_id

        # Content between separators decodes back to the original stories.
        first = ids[: ids.index(eot_id)]
        assert tok.decode(first) == "The cat sat.\nTom ran fast."


def test_encoded_ids_fit_uint16():
    tok = _tiny_tokenizer()
    with tempfile.TemporaryDirectory() as d:
        src = Path(d) / "m.txt"
        dst = Path(d) / "m.bin"
        src.write_text("The cat sat.\n" + EOT + "\n")
        encode_file(tok, src, dst, progress=False)
        arr = load_tokens(dst)
        assert arr.dtype == TOKEN_DTYPE  # stored as uint16
        assert int(arr.max()) < 65536


def main() -> None:
    tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
    for t in tests:
        t()
        print(f"ok  {t.__name__}")
    print(f"\n{len(tests)} passed")


if __name__ == "__main__":
    main()