"""Integration tests for the publish path.

Requires the pytest-arrakis publish extra (embedded Kafka broker).
"""

from __future__ import annotations

import numpy
import pytest
from gpstime import gpsnow

from ... import Client
from ...block import SeriesBlock, Time


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


def _make_block(channels, time_ns):
    data = {
        name: numpy.full(
            int(ch.sample_rate * ch.stride / Time.SECONDS),
            time_ns % 1000,
            dtype=ch.data_type,
        )
        for name, ch in channels.items()
    }
    return SeriesBlock(time_ns, data, channels)


@pytest.mark.integration
def test_publish_block(arrakis_publisher, channels):
    """Publish a real SeriesBlock through the full enter/publish/close path."""
    block = _make_block(channels, 1_000_000_000 * Time.SECONDS)
    arrakis_publisher.publish(block)


@pytest.mark.integration
def test_publish_multiple_blocks(arrakis_publisher, channels):
    """Publish multiple blocks in sequence."""
    t0 = 1_000_000_000 * Time.SECONDS
    stride = channels[next(iter(channels))].stride
    for i in range(3):
        block = _make_block(channels, t0 + i * stride)
        arrakis_publisher.publish(block)


@pytest.mark.integration
def test_publish_roundtrip(
    arrakis_publisher, arrakis_publish_server, kafka_broker, channels
):
    """Publish blocks, then consume them back via Client.stream and verify."""
    stride = channels[next(iter(channels))].stride
    t0 = int(gpsnow()) * Time.SECONDS
    n_blocks = 3

    published_times = []
    for i in range(n_blocks):
        time_ns = t0 + i * stride
        published_times.append(time_ns)
        block = _make_block(channels, time_ns)
        arrakis_publisher.publish(block)

    start = t0 / Time.SECONDS
    end = (t0 + n_blocks * stride) / Time.SECONDS
    channel_names = list(channels.keys())

    client = Client(url=arrakis_publish_server.url)
    consumed_times = []
    for block in client.stream(
        channel_names,
        start=start,
        end=end,
        kafka_url=kafka_broker.bootstrap_server,
    ):
        consumed_times.append(block.time_ns)
        for name in block.channels:
            assert name in channels

    assert sorted(consumed_times) == sorted(published_times)
