import pytest

from edsnlp.utils.batching import (

    DATASET_END_SENTINEL,

    BatchSizeArg,

    FragmentEndSentinel,

    StreamSentinel,

    batchify,

    batchify_by_dataset,

    batchify_by_fragment,

    batchify_by_length_sum,

    batchify_by_padded,

    stat_batchify,

)

class MockStreamSentinel(StreamSentinel):

    pass

# Tests for BatchSizeArg

def test_batch_size_arg_validate():

    # Valid inputs

    assert BatchSizeArg.validate("10 samples") == (10, "samples")

    assert BatchSizeArg.validate("20 words") == (20, "words")

    assert BatchSizeArg.validate(15) == (15, "docs")

    assert BatchSizeArg.validate("docs") == (None, "docs")

    assert BatchSizeArg.validate("tokens") == (None, "tokens")

    assert BatchSizeArg.validate("25") == (25, "docs")

    # Invalid inputs

    with pytest.raises(Exception):

        BatchSizeArg.validate("invalid input")

    with pytest.raises(Exception):

        BatchSizeArg.validate("10 invalid input")

    with pytest.raises(Exception):

        BatchSizeArg.validate("invalid input 10")

# Tests for batchify function

def test_batchify_simple():

    data = [1, 2, 3, 4, 5]

    batches = list(batchify(data, batch_size=2))

    assert batches == [[1, 2], [3, 4], [5]]

def test_batchify_drop_last():

    data = [1, 2, 3, 4, 5]

    batches = list(batchify(data, batch_size=2, drop_last=True))

    assert batches == [[1, 2], [3, 4]]

def test_batchify_sentinel_drop():

    data = [1, 2, MockStreamSentinel(), 3, 4]

    batches = list(batchify(data, batch_size=2, sentinel_mode="drop"))

    assert batches == [[1, 2], [3, 4]]

def test_batchify_sentinel_keep():

    sentinel = MockStreamSentinel()

    data = [1, 2, sentinel, 3, 4]

    batches = list(batchify(data, batch_size=2, sentinel_mode="keep"))

    assert batches == [[1, 2, sentinel], [3, 4]]

def test_batchify_sentinel_split():

    sentinel = MockStreamSentinel()

    data = [1, 2, sentinel, 3, 4]

    batches = list(batchify(data, batch_size=2, sentinel_mode="split"))

    assert batches == [[1, 2], sentinel, [3, 4]]

# Tests for batchify_by_length_sum

def test_batchify_by_length_sum_simple():

    data = ["a", "bb", "ccc", "dddd", "eeeee"]

    batches = list(batchify_by_length_sum(data, batch_size=5))

    assert batches == [["a", "bb"], ["ccc"], ["dddd"], ["eeeee"]]

def test_batchify_by_length_sum_drop_last():

    data = ["a", "bb", "ccc", "dddd", "eeeee"]

    batches = list(batchify_by_length_sum(data, batch_size=5, drop_last=True))

    assert batches == [["a", "bb"], ["ccc"], ["dddd"]]

# Tests for batchify_by_length_sum

def test_batchify_by_length_sum_split():

    sentinel = MockStreamSentinel()

    data = ["aa", "bb", sentinel, "ccc", "dddd", "eeeee"]

    batches = list(batchify_by_length_sum(data, batch_size=7, sentinel_mode="split"))

    assert batches == [["aa", "bb"], sentinel, ["ccc", "dddd"], ["eeeee"]]

# Tests for batchify_by_length_sum

def test_batchify_by_length_sum_keep():

    sentinel = MockStreamSentinel()

    data = ["aa", "bb", sentinel, "ccc", "dddd", "eeeee"]

    batches = list(batchify_by_length_sum(data, batch_size=7, sentinel_mode="keep"))

    assert batches == [["aa", "bb", sentinel, "ccc"], ["dddd"], ["eeeee"]]

# Tests for batchify_by_padded

def test_batchify_by_padded_simple():

    data = ["a", "bb", "ccc", "dddd"]

    batches = list(batchify_by_padded(data, batch_size=6))

    assert batches == [["a", "bb"], ["ccc"], ["dddd"]]

def test_batchify_by_padded_drop_last():

    data = ["a", "bb", "ccc", "dddd"]

    batches = list(batchify_by_padded(data, batch_size=6, drop_last=True))

    assert batches == [["a", "bb"], ["ccc"]]

def test_batchify_by_padded_sentinel_keep():

    sentinel = MockStreamSentinel()

    data = ["a", sentinel, "bb", "ccc"]

    batches = list(batchify_by_padded(data, batch_size=6, sentinel_mode="keep"))

    assert batches == [["a", sentinel, "bb"], ["ccc"]]

def test_batchify_by_padded_sentinel_split():

    sentinel = MockStreamSentinel()

    data = ["a", sentinel, "bb", "ccc"]

    batches = list(batchify_by_padded(data, batch_size=5, sentinel_mode="split"))

    assert batches == [["a"], sentinel, ["bb"], ["ccc"]]

# Tests for batchify_by_dataset

def test_batchify_by_dataset_simple():

    data = [

        "item1",

        "item2",

        DATASET_END_SENTINEL,

        "item3",

        DATASET_END_SENTINEL,

        "item4",

        "item5",

    ]

    batches = list(batchify_by_dataset(data))

    assert batches == [

        ["item1", "item2"],

        DATASET_END_SENTINEL,

        ["item3"],

        DATASET_END_SENTINEL,

        ["item4", "item5"],

    ]

def test_batchify_by_dataset_sentinel_split():

    sentinel = MockStreamSentinel()

    data = ["item1", sentinel, "item2", DATASET_END_SENTINEL, "item3"]

    batches = list(batchify_by_dataset(data, sentinel_mode="split"))

    assert batches == [["item1"], sentinel, ["item2"], DATASET_END_SENTINEL, ["item3"]]

def test_batchify_by_dataset_sentinel_keep():

    sentinel = MockStreamSentinel()

    data = ["item1", sentinel, "item2", DATASET_END_SENTINEL, "item3"]

    batches = list(batchify_by_dataset(data, sentinel_mode="keep"))

    assert batches == [["item1", sentinel, "item2"], DATASET_END_SENTINEL, ["item3"]]

def test_batchify_by_dataset_sentinel_drop():

    sentinel = MockStreamSentinel()

    data = ["item1", sentinel, "item2", DATASET_END_SENTINEL, "item3"]

    batches = list(batchify_by_dataset(data, sentinel_mode="drop"))

    assert batches == [["item1", "item2"], DATASET_END_SENTINEL, ["item3"]]

def test_batchify_by_dataset_drop_last():

    data = ["item1", "item2", DATASET_END_SENTINEL, "item3"]

    batches = list(batchify_by_dataset(data, drop_last=True))

    assert batches == [["item1", "item2"], DATASET_END_SENTINEL]

# Tests for batchify_by_fragment

def test_batchify_by_fragment_simple():

    fragment_end_1 = FragmentEndSentinel("fragment1")

    fragment_end_2 = FragmentEndSentinel("fragment2")

    data = ["item1", "item2", fragment_end_1, "item3", fragment_end_2, "item4"]

    batches = list(batchify_by_fragment(data))

    assert batches == [

        ["item1", "item2"],

        fragment_end_1,

        ["item3"],

        fragment_end_2,

        ["item4"],

    ]

def test_batchify_by_fragment_sentinel_split():

    sentinel = MockStreamSentinel()

    fragment_end = FragmentEndSentinel("fragment")

    data = ["item1", sentinel, "item2", fragment_end]

    batches = list(batchify_by_fragment(data, sentinel_mode="split"))

    assert batches == [["item1"], sentinel, ["item2"], fragment_end]

def test_batchify_by_fragment_sentinel_keep():

    sentinel = MockStreamSentinel()

    fragment_end = FragmentEndSentinel("fragment")

    data = ["item1", sentinel, "item2", fragment_end]

    batches = list(batchify_by_fragment(data, sentinel_mode="keep"))

    assert batches == [["item1", sentinel, "item2"], fragment_end]

def test_batchify_by_fragment_sentinel_drop():

    sentinel = MockStreamSentinel()

    fragment_end = FragmentEndSentinel("fragment")

    data = ["item1", sentinel, "item2", fragment_end]

    batches = list(batchify_by_fragment(data, sentinel_mode="drop"))

    assert batches == [["item1", "item2"], fragment_end]

def test_batchify_by_fragment_drop_last():

    fragment_end = FragmentEndSentinel("fragment")

    data = ["item1", "item2", fragment_end]

    batches = list(batchify_by_fragment(data, sentinel_mode="split", drop_last=True))

    assert batches == [["item1", "item2"], fragment_end]

# Tests for stat_batchify

def test_stat_batchify_simple():

    data = [

        {"/stats/length": 2, "text": "aa"},

        {"/stats/length": 3, "text": "bbb"},

        {"/stats/length": 4, "text": "cccc"},

        {"/stats/length": 2, "text": "dd"},

    ]

    batch_fn = stat_batchify("length")

    batches = list(batch_fn(data, batch_size=5))

    assert batches == [

        [data[0], data[1]],  # Total length: 5

        [data[2]],  # Total length: 4

        [data[3]],  # Total length: 2

    ]

def test_stat_batchify_invalid_key():

    data = [{"text": "aaa"}]

    batch_fn = stat_batchify("length")

    with pytest.raises(ValueError):

        list(batch_fn(data, batch_size=5))

def test_stat_batchify_sentinel_split():

    sentinel = MockStreamSentinel()

    data = [

        {"/stats/length": 2, "text": "aa"},

        sentinel,

        {"/stats/length": 3, "text": "bbb"},

    ]

    batch_fn = stat_batchify("length")

    batches = list(batch_fn(data, batch_size=5, sentinel_mode="split"))

    assert batches == [

        [data[0]],

        sentinel,

        [data[2]],

    ]

def test_stat_batchify_sentinel_keep():

    sentinel = MockStreamSentinel()

    data = [

        {"/stats/length": 2, "text": "aa"},

        sentinel,

        {"/stats/length": 4, "text": "bbbb"},

    ]

    batch_fn = stat_batchify("length")

    batches = list(batch_fn(data, batch_size=5, sentinel_mode="keep"))

    assert batches == [

        [data[0], sentinel],

        [data[2]],

    ]

def test_stat_batchify_drop_last():

    data = [

        {"/stats/length": 2, "text": "aa"},

        {"/stats/length": 3, "text": "bbb"},

        {"/stats/length": 4, "text": "cccc"},

    ]

    batch_fn = stat_batchify("length")

    batches = list(batch_fn(data, batch_size=6, drop_last=True))

    assert batches == [

        [data[0], data[1]],  # Total length: 5

    ]  # Last batch is dropped because total length is 4

# Additional tests to ensure full coverage

def test_batchify_empty_iterable():

    data = []

    batches = list(batchify(data, batch_size=2))

    assert batches == []

def test_batchify_by_length_sum_empty_iterable():

    data = []

    batches = list(batchify_by_length_sum(data, batch_size=5))

    assert batches == []

def test_batchify_by_padded_empty_iterable():

    data = []

    batches = list(batchify_by_padded(data, batch_size=6))

    assert batches == []

def test_batchify_by_dataset_empty_iterable():

    data = []

    batches = list(batchify_by_dataset(data))

    assert batches == []

def test_batchify_by_fragment_empty_iterable():

    data = []

    batches = list(batchify_by_fragment(data))

    assert batches == []

def test_stat_batchify_empty_iterable():

    data = []

    batch_fn = stat_batchify("length")

    batches = list(batch_fn(data, batch_size=5))

    assert batches == []

def test_batchify_invalid_sentinel_mode():

    data = [1, 2, 3]

    with pytest.raises(AssertionError):

        list(batchify(data, batch_size=2, sentinel_mode="invalid_mode"))