import pytest

from datasets import Dataset

from dlb.data import build_owt_split, pack_tokens, preprocess_split


class TinyBatchTokenizer:
    def __call__(self, texts, **kwargs):
        assert kwargs == {
            "add_special_tokens": False,
            "return_attention_mask": False,
            "return_token_type_ids": False,
        }
        table = {"first": [10, 11], "second": [12, 13, 14]}
        return {"input_ids": [table[text] for text in texts]}


def test_owt_split_is_last_100k_documents():
    """Catch shuffling or selecting any validation slice except the final 100k."""

    split = build_owt_split(total_documents=8_013_769)

    assert split.train == "train[:-100000]"
    assert split.validation == "train[-100000:]"


def test_owt_split_rejects_a_source_too_small_for_both_splits():
    """Catch silently producing an empty training split from a truncated snapshot."""

    with pytest.raises(ValueError, match="more than 100,000"):
        build_owt_split(total_documents=100_000)


def test_pack_tokens_reserves_boundaries():
    """Catch a packed block ending in document content rather than EOS."""

    blocks = list(
        pack_tokens([[10, 11], [12, 13, 14]], length=6, bos_id=101, eos_id=102)
    )

    assert blocks[0] == [101, 10, 11, 102, 12, 102]
    assert all(len(block) == 6 for block in blocks)


def test_pack_tokens_carries_overflow_and_drops_only_the_incomplete_tail():
    """Catch token loss when one document crosses a fixed-length boundary."""

    blocks = list(
        pack_tokens([[10, 11, 12, 13, 14, 15]], length=6, bos_id=101, eos_id=102)
    )

    assert blocks == [[101, 10, 11, 12, 13, 102]]


def test_preprocess_split_packs_real_dataset_rows_in_source_order(tmp_path):
    """Catch preprocessing that reorders documents or packs each batch separately."""

    source = Dataset.from_dict({"text": ["first", "second"]})

    processed = preprocess_split(
        source,
        tokenizer=TinyBatchTokenizer(),
        length=6,
        bos_id=101,
        eos_id=102,
        cache_dir=tmp_path,
        batch_documents=1,
    )

    assert processed.to_list() == [
        {"input_ids": [101, 10, 11, 102, 12, 102]},
    ]