import gpstime
import numpy
import pytest

from ...block import SeriesBlock, Time
from ...channel import Channel
from ...mux import BlockMuxStream, TimedQueue

STRIDE = 62_500_000  # 62.5ms in nanoseconds
TIMEOUT = 2 * Time.s
GPS_T0 = 1_400_000_000 * Time.s
GPS_T0_DATETIME = gpstime.gpstime.fromgps(GPS_T0 / Time.s)


def make_channel(name="H1:FKE-TEST_CHANNEL", sample_rate=128, **kwargs):
    defaults = {
        "data_type": "float64",
        "partition_id": "A",
        "partition_index": 0,
        "stride": STRIDE,
        "max_latency": TIMEOUT,
    }
    defaults.update(kwargs)
    return Channel(name, sample_rate=sample_rate, **defaults)


def make_block(time_ns, channels, *, gap=False):
    """Create a SeriesBlock at the given time for the given channels."""
    if isinstance(channels, Channel):
        channels = [channels]
    data = {}
    chan_dict = {}
    for ch in channels:
        size = int(ch.sample_rate * STRIDE) // Time.s
        if gap:
            data[ch.name] = numpy.ma.masked_all(size, dtype=ch.data_type)
        else:
            data[ch.name] = numpy.random.randint(256, size=size).astype(ch.data_type)
        chan_dict[ch.name] = ch
    return SeriesBlock(time_ns, data, chan_dict)


def is_gap_block(block):
    """True if all series in the block are fully masked."""
    return all(
        isinstance(arr, numpy.ma.MaskedArray) and arr.mask.all()
        for arr in block.data.values()
    )


class FakeStreamReader:
    """A controllable StreamReader test double.

    Push (stream_id, block) pairs via `feed()`, then `read()` yields them.
    Call `finish()` to make `done()` return True.
    """

    def __init__(self, streams_map):
        self._streams = streams_map
        self._buffer = []
        self._done = False

    @property
    def streams(self):
        return self._streams

    def feed(self, stream_id, block):
        self._buffer.append((stream_id, block))

    def finish(self):
        self._done = True

    def enter(self):
        pass

    def __enter__(self):
        self.enter()
        return self

    def read(self, *, convert_blocks=False):
        buf = list(self._buffer)
        self._buffer.clear()
        yield from buf

    def done(self):
        return self._done and len(self._buffer) == 0

    def close(self):
        pass

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()
        return False


class ContinuousStreamReader:
    """A reader that yields blocks continuously from a generator.

    read() calls the generator to produce the next batch of blocks.
    done() always returns False — the only way to stop is via end_ns.
    Tracks how many blocks have been yielded and the max queue depth
    observed externally (set by test code).
    """

    def __init__(self, streams_map, block_generator):
        self._streams = streams_map
        self._gen = block_generator
        self.blocks_yielded = 0

    @property
    def streams(self):
        return self._streams

    def enter(self):
        pass

    def __enter__(self):
        self.enter()
        return self

    def read(self, *, convert_blocks=False):
        for item in self._gen:
            self.blocks_yielded += 1
            yield item

    def done(self):
        return False

    def close(self):
        pass

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()
        return False


class TestTimedQueueBasic:
    def test_push_pull_sequential(self):
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=GPS_T0)
        t1 = GPS_T0
        t2 = GPS_T0 + STRIDE

        q.push(t1, "block_1")
        q.push(t2, "block_2")

        items = list(q.pull())
        assert len(items) == 2
        assert items[0] == (t1, "block_1")
        assert items[1] == (t2, "block_2")

    def test_drain_empties_queue(self):
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=GPS_T0)
        q.push(GPS_T0, "a")
        q.push(GPS_T0 + STRIDE, "b")
        list(q.pull())
        assert len(q) == 0

    def test_backfill_gaps(self):
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=GPS_T0)
        q.push(GPS_T0 + 2 * STRIDE, "after_gap")

        items = list(q.pull())
        assert len(items) == 3
        assert items[0] == (GPS_T0, None)
        assert items[1] == (GPS_T0 + STRIDE, None)
        assert items[2] == (GPS_T0 + 2 * STRIDE, "after_gap")

    def test_drop_old_timestamps(self):
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=GPS_T0)
        q.push(GPS_T0 + STRIDE, "newer")
        with pytest.warns(UserWarning, match="timestamp is too old"):
            q.push(GPS_T0, "older")

        items = list(q.pull())
        assert len(items) == 2
        assert items[0] == (GPS_T0, None)
        assert items[1] == (GPS_T0 + STRIDE, "newer")

    def test_drop_equal_timestamp(self):
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=GPS_T0)
        q.push(GPS_T0, "first")
        with pytest.warns(UserWarning, match="timestamp is too old"):
            q.push(GPS_T0, "duplicate")

        items = list(q.pull())
        assert len(items) == 1
        assert items[0] == (GPS_T0, "first")

    def test_stride_alignment_assertion(self):
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=GPS_T0)
        with pytest.raises(AssertionError, match="not a multiple"):
            q.push(GPS_T0 + 1, "misaligned")

    def test_ready_with_enough_elements(self):
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=GPS_T0)
        duration = 2 * STRIDE
        q.push(GPS_T0, "a")
        q.push(GPS_T0 + STRIDE, "b")
        assert q.ready(duration) == 2

    def test_ready_when_empty(self):
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=GPS_T0)
        assert q.ready(STRIDE) is None

    def test_pull_with_duration(self):
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=GPS_T0)
        duration = 2 * STRIDE
        for i in range(4):
            q.push(GPS_T0 + i * STRIDE, f"block_{i}")

        items = list(q.pull(duration))
        assert len(items) == 2
        assert items[0] == (GPS_T0, "block_0")
        assert items[1] == (GPS_T0 + STRIDE, "block_1")
        assert len(q) == 2

    def test_pull_with_duration_not_enough(self):
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=GPS_T0)
        q.push(GPS_T0, "only_one")
        items = list(q.pull(2 * STRIDE))
        assert len(items) == 0
        assert len(q) == 1


class TestTimedQueueLive:
    """Tests for start=None (live stream) behavior using freezer."""

    def test_initialized_reset_on_first_real_block(self, freezer):
        freezer.move_to(GPS_T0_DATETIME)
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=None)
        assert not q._initialized

        q.push(GPS_T0, "real_block")
        assert q._initialized
        assert q.last_time == GPS_T0

    def test_gap_push_does_not_set_initialized(self, freezer):
        freezer.move_to(GPS_T0_DATETIME)
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=None)
        gap_time = (q.last_time // STRIDE + 1) * STRIDE
        q.push(gap_time, None, on_drop="ignore")
        assert not q._initialized

    def test_first_real_block_clears_queued_gaps(self, freezer):
        """After timeout gaps have accumulated, first real block resets
        last_time, effectively discarding the stale gaps."""
        freezer.move_to(GPS_T0_DATETIME)
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=None)

        # simulate timeout pushing a few gaps
        gap_time = (q.last_time // STRIDE + 1) * STRIDE
        q.push(gap_time, None, on_drop="ignore")
        q.push(gap_time + STRIDE, None, on_drop="ignore")
        assert len(q) == 2

        # first real block arrives — resets last_time
        real_time = GPS_T0
        q.push(real_time, "real_block")
        assert q._initialized
        assert q.last_time == real_time

    def test_pull_drains_with_frozen_time(self, freezer):
        """With time frozen near data timestamps, pull works normally."""
        freezer.move_to(GPS_T0_DATETIME)
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=None)

        q.push(GPS_T0, "first")
        q.push(GPS_T0 + STRIDE, "second")
        q.push(GPS_T0 + 2 * STRIDE, "third")

        items = list(q.pull())
        assert len(items) == 3
        assert all(item[1] is not None for item in items)

    def test_timeout_pushes_gaps_when_no_data(self, freezer):
        """When time advances past timeout with no data, gaps appear."""
        freezer.move_to(GPS_T0_DATETIME)
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=None)

        # push one real block to initialize
        q.push(GPS_T0, "real_block")
        initial_len = len(q)

        # advance time past timeout
        freezer.tick(TIMEOUT / Time.s + 0.1)

        # _update_timeout (called by ready/pull) should push gap(s)
        q._update_timeout()
        assert len(q) > initial_len

    def test_timeout_does_not_outrun_real_data(self, freezer):
        """When real data arrives on time, timeout gaps should be dropped."""
        freezer.move_to(GPS_T0_DATETIME)
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=None)

        # push blocks up to current time
        n_strides = int(TIMEOUT / STRIDE) + 2
        for i in range(n_strides):
            t = GPS_T0 + i * STRIDE
            q.push(t, f"block_{i}")

        # advance time by one stride
        freezer.tick(STRIDE / Time.s)
        q._update_timeout()

        # timeout gap should be dropped (data is ahead of timeout)
        items = list(q.pull())
        assert all(item[1] is not None for item in items)

    def test_recovery_after_timeout_gaps(self, freezer):
        """After timeout gaps, real data arriving should be usable."""
        freezer.move_to(GPS_T0_DATETIME)
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=None)

        # push initial data
        q.push(GPS_T0, "initial")

        # advance time to create some timeout gaps
        freezer.tick(TIMEOUT / Time.s + STRIDE / Time.s)
        q._update_timeout()
        gap_count = len(q) - 1  # subtract the initial real block
        assert gap_count > 0

        # now push new real data at a time near "now"
        now_ns = int(gpstime.gpsnow() * Time.s)
        now_aligned = (now_ns // STRIDE) * STRIDE
        q.push(now_aligned, "recovery")

        items = list(q.pull())
        real_items = [item for item in items if item[1] is not None]
        assert len(real_items) == 2  # initial + recovery

    def test_late_data_dropped_after_timeout_advances(self, freezer):
        """Data arriving after timeout has advanced past it gets dropped."""
        freezer.move_to(GPS_T0_DATETIME)
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=None)

        q.push(GPS_T0, "on_time")

        # advance time well past the initial data
        freezer.tick(TIMEOUT / Time.s + 1.0)
        q._update_timeout()
        # drain everything
        list(q.pull())

        # try to push data at the old time — should be dropped
        with pytest.warns(UserWarning, match="timestamp is too old"):
            q.push(GPS_T0 + STRIDE, "late_arrival")


class TestTimedQueueFrontAndDrain:
    def test_front_time_empty(self):
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=GPS_T0)
        assert q.front_time() is None

    def test_front_time(self):
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=GPS_T0)
        q.push(GPS_T0, "a")
        q.push(GPS_T0 + STRIDE, "b")
        assert q.front_time() == GPS_T0

    def test_drain_until(self):
        q = TimedQueue(stride=STRIDE, timeout=TIMEOUT, start=GPS_T0)
        for i in range(5):
            q.push(GPS_T0 + i * STRIDE, f"block_{i}")
        target = GPS_T0 + 3 * STRIDE
        q.drain_until(target)
        assert q.front_time() == target
        assert len(q) == 2


class TestBlockMuxStreamSingleStream:
    def _make_reader_and_channel(self):
        ch = make_channel()
        stream_id = "test_stream"
        reader = FakeStreamReader({stream_id: [ch]})
        return reader, stream_id, ch

    def test_push_ready_pull(self):
        reader, stream_id, ch = self._make_reader_and_channel()
        mux = BlockMuxStream(reader, start=GPS_T0)

        block = make_block(GPS_T0, ch)
        mux.push(stream_id, block)

        assert mux.ready()
        result = mux.pull()
        assert result.time_ns == GPS_T0

    def test_gap_block_for_none_elements(self):
        reader, stream_id, ch = self._make_reader_and_channel()
        mux = BlockMuxStream(reader, start=GPS_T0)

        block = make_block(GPS_T0 + STRIDE, ch)
        mux.push(stream_id, block)

        assert mux.ready()
        result = mux.pull()
        assert result.time_ns == GPS_T0
        assert is_gap_block(result)

    def test_stream_yields_blocks(self):
        reader, stream_id, ch = self._make_reader_and_channel()

        n_blocks = 4
        for i in range(n_blocks):
            reader.feed(stream_id, make_block(GPS_T0 + i * STRIDE, ch))
        reader.finish()

        end_ns = GPS_T0 + n_blocks * STRIDE
        mux = BlockMuxStream(reader, start=GPS_T0)

        blocks = list(mux.stream(end_ns))
        assert len(blocks) > 0
        for b in blocks:
            assert ch.name in b.channels

    def test_stream_stops_at_end(self):
        reader, stream_id, ch = self._make_reader_and_channel()

        for i in range(10):
            reader.feed(stream_id, make_block(GPS_T0 + i * STRIDE, ch))
        reader.finish()

        end_ns = GPS_T0 + 3 * STRIDE
        mux = BlockMuxStream(reader, start=GPS_T0)
        blocks = list(mux.stream(end_ns))

        for b in blocks:
            assert b.time_ns <= end_ns

    def test_stream_start_none_immediate_data(self, freezer):
        """With start=None and data arriving immediately, blocks are yielded."""
        freezer.move_to(GPS_T0_DATETIME)
        reader, stream_id, ch = self._make_reader_and_channel()

        n_blocks = 4
        for i in range(n_blocks):
            reader.feed(stream_id, make_block(GPS_T0 + i * STRIDE, ch))
        reader.finish()

        end_ns = GPS_T0 + n_blocks * STRIDE
        mux = BlockMuxStream(reader, start=None)

        blocks = list(mux.stream(end_ns))
        assert len(blocks) > 0
        for b in blocks:
            assert ch.name in b.channels


class TestBlockMuxStreamTwoStreams:
    def _make_reader_and_channels(self):
        ch1 = make_channel(
            name="H1:FKE-CHANNEL_A",
            partition_id="A",
            partition_index=0,
        )
        ch2 = make_channel(
            name="H1:FKE-CHANNEL_B",
            partition_id="B",
            partition_index=1,
        )
        stream_a = "stream_A"
        stream_b = "stream_B"
        reader = FakeStreamReader({stream_a: [ch1], stream_b: [ch2]})
        return reader, stream_a, stream_b, ch1, ch2

    def test_combines_two_streams(self):
        reader, sa, sb, ch1, ch2 = self._make_reader_and_channels()
        mux = BlockMuxStream(reader, start=GPS_T0)

        mux.push(sa, make_block(GPS_T0, ch1))
        mux.push(sb, make_block(GPS_T0, ch2))

        assert mux.ready()
        result = mux.pull()
        assert ch1.name in result.channels
        assert ch2.name in result.channels

    def test_partial_stream_produces_gap(self):
        reader, sa, sb, ch1, ch2 = self._make_reader_and_channels()
        mux = BlockMuxStream(reader, start=GPS_T0)

        mux.push(sa, make_block(GPS_T0, ch1))
        mux[sb].push(GPS_T0, None, on_drop="ignore")

        assert mux.ready()
        result = mux.pull()
        assert ch1.name in result.channels
        assert ch2.name in result.channels
        assert numpy.ma.is_masked(result.data[ch2.name])

    def test_stream_with_two_streams_interleaved(self):
        reader, sa, sb, ch1, ch2 = self._make_reader_and_channels()
        n_blocks = 3
        for i in range(n_blocks):
            reader.feed(sa, make_block(GPS_T0 + i * STRIDE, ch1))
            reader.feed(sb, make_block(GPS_T0 + i * STRIDE, ch2))
        reader.finish()

        end_ns = GPS_T0 + n_blocks * STRIDE
        mux = BlockMuxStream(reader, start=GPS_T0)
        blocks = list(mux.stream(end_ns))

        assert len(blocks) > 0
        for b in blocks:
            assert ch1.name in b.channels
            assert ch2.name in b.channels


class TestBlockMuxStreamContinuousRead:
    @pytest.mark.timeout(5)
    def test_continuous_single_stream_yields_blocks(self):
        """stream() must yield blocks even when read() never stops."""
        ch = make_channel()
        stream_id = "test_stream"
        n_total = 20
        end_ns = GPS_T0 + n_total * STRIDE

        def gen():
            for i in range(n_total):
                yield stream_id, make_block(GPS_T0 + i * STRIDE, ch)

        reader = ContinuousStreamReader({stream_id: [ch]}, gen())
        mux = BlockMuxStream(reader, start=GPS_T0)
        blocks = list(mux.stream(end_ns))

        assert len(blocks) > 0
        for b in blocks:
            assert ch.name in b.channels

    @pytest.mark.timeout(5)
    def test_continuous_read_does_not_buffer_all_before_yielding(self):
        """stream() should yield blocks incrementally, not buffer everything
        from read() first."""
        ch = make_channel()
        stream_id = "test_stream"
        n_total = 100
        end_ns = GPS_T0 + n_total * STRIDE

        def gen():
            for i in range(n_total):
                yield stream_id, make_block(GPS_T0 + i * STRIDE, ch)

        reader = ContinuousStreamReader({stream_id: [ch]}, gen())
        mux = BlockMuxStream(reader, start=GPS_T0)

        max_queue_depth = 0
        for _block in mux.stream(end_ns):
            for q in mux._queues.values():
                max_queue_depth = max(max_queue_depth, len(q))

        assert max_queue_depth < n_total, (
            f"queue depth {max_queue_depth} suggests all blocks were "
            f"buffered before any were yielded"
        )

    @pytest.mark.timeout(5)
    def test_continuous_read_end_ns_stops_early(self):
        """stream() should stop at end_ns even with unlimited data."""
        ch = make_channel()
        stream_id = "test_stream"
        stop_at = 5
        end_ns = GPS_T0 + stop_at * STRIDE

        def gen():
            i = 0
            while True:
                yield stream_id, make_block(GPS_T0 + i * STRIDE, ch)
                i += 1

        reader = ContinuousStreamReader({stream_id: [ch]}, gen())
        mux = BlockMuxStream(reader, start=GPS_T0)
        blocks = list(mux.stream(end_ns))

        assert len(blocks) > 0
        for b in blocks:
            assert b.time_ns <= end_ns

    @pytest.mark.timeout(5)
    def test_continuous_two_streams_yields_combined(self):
        """With two continuous streams, stream() yields combined blocks."""
        ch1 = make_channel(
            name="H1:FKE-CHANNEL_A",
            partition_id="A",
            partition_index=0,
        )
        ch2 = make_channel(
            name="H1:FKE-CHANNEL_B",
            partition_id="B",
            partition_index=1,
        )
        sa, sb = "stream_A", "stream_B"
        n_total = 10
        end_ns = GPS_T0 + n_total * STRIDE

        def gen():
            for i in range(n_total):
                t = GPS_T0 + i * STRIDE
                yield sa, make_block(t, ch1)
                yield sb, make_block(t, ch2)

        reader = ContinuousStreamReader({sa: [ch1], sb: [ch2]}, gen())
        mux = BlockMuxStream(reader, start=GPS_T0)
        blocks = list(mux.stream(end_ns))

        assert len(blocks) > 0
        for b in blocks:
            assert ch1.name in b.channels
            assert ch2.name in b.channels

    @pytest.mark.timeout(5)
    def test_continuous_read_many_blocks_per_call(self):
        """When read() yields many blocks at once, stream() should still
        yield incrementally."""
        ch = make_channel()
        stream_id = "test_stream"
        n_total = 50
        end_ns = GPS_T0 + n_total * STRIDE

        def gen():
            for i in range(n_total):
                yield stream_id, make_block(GPS_T0 + i * STRIDE, ch)

        reader = ContinuousStreamReader({stream_id: [ch]}, gen())
        mux = BlockMuxStream(reader, start=GPS_T0)

        block_count = 0
        for _block in mux.stream(end_ns):
            block_count += 1
            if block_count == 1:
                assert reader.blocks_yielded < n_total, (
                    "all blocks consumed from reader before first yield"
                )

        assert block_count > 0

    @pytest.mark.timeout(5)
    def test_continuous_two_streams_one_ahead(self):
        """When one stream is ahead of another, stream() should still
        yield when both queues have enough data."""
        ch1 = make_channel(
            name="H1:FKE-CHANNEL_A",
            partition_id="A",
            partition_index=0,
        )
        ch2 = make_channel(
            name="H1:FKE-CHANNEL_B",
            partition_id="B",
            partition_index=1,
        )
        sa, sb = "stream_A", "stream_B"
        n_total = 10
        end_ns = GPS_T0 + n_total * STRIDE

        def gen():
            # stream A sends 3 blocks before stream B sends any
            for i in range(3):
                yield sa, make_block(GPS_T0 + i * STRIDE, ch1)
            # then they alternate
            for i in range(n_total):
                t = GPS_T0 + i * STRIDE
                yield sb, make_block(t, ch2)
                if i + 3 < n_total:
                    yield sa, make_block(GPS_T0 + (i + 3) * STRIDE, ch1)

        reader = ContinuousStreamReader({sa: [ch1], sb: [ch2]}, gen())
        mux = BlockMuxStream(reader, start=GPS_T0)
        blocks = list(mux.stream(end_ns))

        assert len(blocks) > 0
        for b in blocks:
            assert ch1.name in b.channels
            assert ch2.name in b.channels
