"""Unit tests for Publisher functionality."""

from __future__ import annotations

from datetime import timedelta
from unittest.mock import MagicMock, patch

import numpy
import pyarrow
import pytest

from arrakis.block import SeriesBlock, Time
from arrakis.channel import Channel
from arrakis.publish import Publisher, serialize_batch


@pytest.fixture
def channels(arrakis_channels):
    """FAKE_H1 channels from the pytest-arrakis channel config."""
    return {
        name: ch for name, ch in arrakis_channels.items() if ch.publisher == "FAKE_H1"
    }


@pytest.fixture
def mock_client():
    """Patch arrakis.publish.Client; yield the *instance* mock."""
    with patch("arrakis.publish.Client") as cls:
        yield cls.return_value


@pytest.fixture
def fake_producer():
    """A plain MagicMock standing in for a Kafka Producer."""
    return MagicMock()


@pytest.fixture
def publisher(channels, mock_client):
    """A Publisher that has been registered (but not entered)."""
    mock_client.find.return_value = list(channels.values())
    pub = Publisher("FAKE_H1")
    pub.register()
    return pub


@pytest.fixture
def ready_publisher(publisher, fake_producer):
    """A Publisher ready to publish, with a fake producer injected."""
    publisher._producer = fake_producer
    yield publisher, fake_producer
    publisher.close()


@pytest.fixture
def series_block(channels):
    """A real SeriesBlock with one second of data for all FAKE_H1 channels."""
    time_ns = 1_000_000_000 * Time.SECONDS
    data = {
        name: numpy.zeros(int(ch.sample_rate), dtype=ch.data_type)
        for name, ch in channels.items()
    }
    return SeriesBlock(time_ns, data, channels)


def test_serialize_batch():
    schema = pyarrow.schema(
        [
            pyarrow.field("time", pyarrow.int64()),
            pyarrow.field("data", pyarrow.float64()),
        ]
    )
    batch = pyarrow.RecordBatch.from_arrays(
        [
            pyarrow.array([1_000_000_000, 1_000_000_001]),
            pyarrow.array([1.0, 2.0]),
        ],
        schema=schema,
    )

    buf = serialize_batch(batch)
    assert isinstance(buf, pyarrow.Buffer)

    roundtrip = pyarrow.ipc.open_stream(buf).read_next_batch()
    assert roundtrip.equals(batch)


def test_register_success(channels, mock_client):
    mock_client.find.return_value = list(channels.values())
    pub = Publisher("FAKE_H1")
    result = pub.register()

    assert result is pub
    assert pub.channels == channels
    mock_client.find.assert_called_once_with(publisher="FAKE_H1")


def test_register_unknown_publisher(mock_client):
    mock_client.find.return_value = []
    pub = Publisher("unknown")
    with pytest.raises(ValueError, match="unknown publisher ID 'unknown'"):
        pub.register()


def test_register_channel_without_partition_id(mock_client):
    ch = Channel(
        name="H1:TEST-NO_PART",
        data_type="float64",
        sample_rate=16384,
        publisher="FAKE_H1",
        partition_id=None,
    )
    mock_client.find.return_value = [ch]
    pub = Publisher("FAKE_H1")
    with pytest.raises(AssertionError, match="missing partition_id"):
        pub.register()


def test_register_channel_without_partition_index(mock_client):
    ch = Channel(
        name="H1:TEST-NO_IDX",
        data_type="float64",
        sample_rate=16384,
        publisher="FAKE_H1",
        partition_id="part_1",
        partition_index=None,
    )
    mock_client.find.return_value = [ch]
    pub = Publisher("FAKE_H1")
    with pytest.raises(AssertionError, match="missing partition_index"):
        pub.register()


def test_register_inconsistent_stride(mock_client):
    ch1 = Channel(
        name="H1:TEST-A",
        data_type="float64",
        sample_rate=16384,
        publisher="FAKE_H1",
        partition_id="p1",
        partition_index=0,
        stride=62_500_000,
    )
    ch2 = Channel(
        name="H1:TEST-B",
        data_type="float32",
        sample_rate=512,
        publisher="FAKE_H1",
        partition_id="p2",
        partition_index=0,
        stride=125_000_000,
    )
    mock_client.find.return_value = [ch1, ch2]
    pub = Publisher("FAKE_H1")
    with pytest.raises(AssertionError, match="inconsistent stride"):
        pub.register()


def test_register_self_register_not_supported():
    """Passing channels= to register() raises NotImplementedError."""
    ch = Channel(name="H1:TEST-X", data_type="float64", sample_rate=16384)
    pub = Publisher("FAKE_H1")
    with pytest.raises(NotImplementedError, match="self-register"):
        pub.register(channels=[ch])


def test_publish_success(ready_publisher, series_block, channels):
    publisher, producer = ready_publisher
    publisher.publish(series_block)

    n_partitions = len({ch.partition_id for ch in channels.values()})
    assert producer.produce.call_count == n_partitions
    topics = {call.kwargs["topic"] for call in producer.produce.call_args_list}
    expected_topics = {f"arrakis-{ch.partition_id}" for ch in channels.values()}
    assert topics == expected_topics
    assert producer.flush.call_count == n_partitions


def test_publish_without_setup(series_block):
    """Publishing without registering or entering raises RuntimeError."""
    pub = Publisher("FAKE_H1")
    with pytest.raises(RuntimeError, match="publication interface not initialized"):
        pub.publish(series_block)


def test_publish_not_entered(publisher, series_block):
    """Registered but not entered raises RuntimeError."""
    with pytest.raises(RuntimeError, match="publication interface not initialized"):
        publisher.publish(series_block)


def test_publish_unknown_channel(ready_publisher):
    publisher, _ = ready_publisher
    unknown = Channel(
        name="H1:UNKNOWN-CHAN",
        data_type="float64",
        sample_rate=16384,
        publisher="FAKE_H1",
        partition_id="part_1",
    )
    block = MagicMock(spec=SeriesBlock)
    block.channels = {"H1:UNKNOWN-CHAN": unknown}
    with pytest.raises(KeyError):
        publisher.publish(block)


def test_publish_modified_channel(ready_publisher, channels):
    publisher, _ = ready_publisher
    name, original = next(iter(channels.items()))
    wrong_dtype = "int32" if original.data_type != "int32" else "float64"
    modified = Channel(
        name=name,
        data_type=wrong_dtype,
        sample_rate=original.sample_rate,
        publisher=original.publisher,
        partition_id=original.partition_id,
    )
    block = MagicMock(spec=SeriesBlock)
    block.channels = {name: modified}
    with pytest.raises(ValueError, match="invalid channel"):
        publisher.publish(block)


def test_publish_with_timeout(ready_publisher, series_block, channels):
    publisher, producer = ready_publisher
    publisher.publish(series_block, timeout=timedelta(seconds=5))
    n_partitions = len({ch.partition_id for ch in channels.values()})
    assert producer.produce.call_count == n_partitions


def test_publish_multiple_blocks(ready_publisher, series_block, channels):
    publisher, producer = ready_publisher
    publisher.publish(series_block)
    publisher.publish(series_block)
    publisher.publish(series_block)

    n_partitions = len({ch.partition_id for ch in channels.values()})
    assert producer.produce.call_count == 3 * n_partitions
    assert producer.flush.call_count == 3 * n_partitions


def test_publish_single_channel(ready_publisher, channels):
    publisher, producer = ready_publisher
    name, ch = next(iter(channels.items()))
    time_ns = 1_000_000_000 * Time.SECONDS
    data = {name: numpy.zeros(int(ch.sample_rate), dtype=ch.data_type)}
    block = SeriesBlock(time_ns, data, {name: ch})

    publisher.publish(block)
    assert producer.produce.call_count == 1
    assert producer.produce.call_args.kwargs["topic"] == f"arrakis-{ch.partition_id}"


def test_publish_data_is_deserializable(ready_publisher, series_block):
    """Verify the data passed to produce() can be deserialized."""
    publisher, producer = ready_publisher
    publisher.publish(series_block)

    for call in producer.produce.call_args_list:
        buf = call.kwargs["value"]
        assert isinstance(buf, pyarrow.Buffer)
        batch = pyarrow.ipc.open_stream(buf).read_next_batch()
        assert batch.num_columns == 3  # time, id, data
        assert batch.column("time")[0].as_py() == series_block.time_ns


def test_context_manager(publisher, fake_producer):
    """__enter__ calls enter(), __exit__ calls close()."""
    publisher.enter = MagicMock()
    publisher._producer = fake_producer

    with publisher as pub:
        assert pub is publisher
        publisher.enter.assert_called_once()

    fake_producer.flush.assert_called()


def test_close_flushes_producer(publisher, fake_producer):
    publisher._producer = fake_producer
    publisher.close()
    fake_producer.flush.assert_called_once()


def test_close_suppresses_exceptions(publisher):
    """close() should not raise even if flush() fails."""
    bad_producer = MagicMock()
    bad_producer.flush.side_effect = RuntimeError("broker down")
    publisher._producer = bad_producer
    publisher.close()  # should not raise
