from collections import namedtuple
from unittest.mock import patch

import pyarrow

from ... import Channel
from ...kafka import KafkaReader

DummyChannel = namedtuple("DummyChannel", ["name", "partition_id", "partition_index"])


def build_test_rb1() -> pyarrow.RecordBatch:
    time = [1000, 1000]
    ids = pyarrow.array([42, 5], type=pyarrow.uint32())
    data = [
        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
        [
            2,
        ],
    ]

    names = ["time", "id", "data"]
    return pyarrow.RecordBatch.from_arrays([time, ids, data], names=names)


def build_test_rb1_with_extra() -> pyarrow.RecordBatch:
    time = [1000, 1000, 1000, 1000]
    ids = pyarrow.array([42, 43, 5, 6], type=pyarrow.uint32())
    data = [
        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
        [43],
        [
            2,
        ],
        [6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6],
    ]

    names = ["time", "id", "data"]
    return pyarrow.RecordBatch.from_arrays([time, ids, data], names=names)


def build_rb1_channels() -> list[Channel]:
    return {
        channel.name: channel
        for channel in [
            Channel(
                name="X1:TST-SUB_1",
                data_type="int64",
                sample_rate=256,
                partition_id="PART1",
                partition_index=42,
                stride=1_000_000_000,
                max_latency=2,
            ),
            Channel(
                name="X1:TST-SUB_2",
                data_type="int64",
                sample_rate=16,
                partition_id="PART1",
                partition_index=5,
                stride=1_000_000_000,
                max_latency=2,
            ),
        ]
    }


def build_rb1_partition_mapping() -> dict[str, str]:
    return {
        "X1:TST-SUB_1": "PART1",
        "X1:TST-SUB_1A": "PART1",
        "X1:TST-SUB_2": "PART1",
        "X1:TST-SUB_2A": "PART1",
    }


def test_smart_filter_batch_fast():
    batch = build_test_rb1()
    expected_output = build_test_rb1()
    partition = "PART1"
    channels = build_rb1_channels()

    with patch("confluent_kafka.Consumer"):
        conn = KafkaReader(
            url="grpc://non-existant:31206",
            metadata=channels,
        )

        actual = conn._smart_filter_batch(batch, partition)
    assert actual == expected_output


def test_smart_filter_batch():
    batch = build_test_rb1_with_extra()
    expected_output = build_test_rb1()
    partition = "PART1"
    channels = build_rb1_channels()

    conn = KafkaReader(
        url="grpc://non-existant:31206",
        metadata=channels,
    )

    actual = conn._smart_filter_batch(batch, partition)
    assert actual == expected_output
