"""Test Multiplexer queue."""

import asyncio
import os

import pytest

from snitun.multiplexer.message import (
    CHANNEL_FLOW_DATA,
    HEADER_SIZE,
    MultiplexerChannelId,
    MultiplexerMessage,
)
from snitun.multiplexer.queue import (
    MultiplexerMultiChannelQueue,
    MultiplexerSingleChannelQueue,
)

MOCK_MSG_SIZE = 4


def _make_mock_channel_id() -> MultiplexerChannelId:
    return MultiplexerChannelId(os.urandom(16))


def _make_mock_message(
    channel_id: MultiplexerChannelId,
    size: int = MOCK_MSG_SIZE,
) -> MultiplexerMessage:
    return MultiplexerMessage(channel_id, CHANNEL_FLOW_DATA, os.urandom(size))


async def test_get_non_existent_channels() -> None:
    """Test MultiplexerMultiChannelQueue get on non-existent channel."""
    queue = MultiplexerMultiChannelQueue(100000, 10, 1000)
    assert queue.empty(_make_mock_channel_id())
    assert not queue.full(_make_mock_channel_id())
    assert queue.size(_make_mock_channel_id()) == 0
    # Make sure defaultdict does not leak
    assert not queue._channels


async def test_single_channel_queue() -> None:
    """Test MultiplexerSingleChannelQueue."""
    queue = MultiplexerSingleChannelQueue(100, 10, 50, lambda _: None)
    channel_id = _make_mock_channel_id()
    msg = _make_mock_message(channel_id)
    assert queue.qsize() == 0
    queue.put_nowait(msg)
    assert queue.qsize() == len(msg.data) + HEADER_SIZE
    assert queue.get_nowait() == msg
    assert queue.qsize() == 0
    queue.put_nowait(None)
    assert queue.qsize() == 0


async def test_multi_channel_queue_full() -> None:
    """Test MultiplexerMultiChannelQueue getting full."""
    msg_size = MOCK_MSG_SIZE + HEADER_SIZE
    # Max two mock messages per channel
    queue = MultiplexerMultiChannelQueue(msg_size * 2, msg_size, msg_size * 2)

    channel_one_id = _make_mock_channel_id()
    channel_two_id = _make_mock_channel_id()
    queue.create_channel(channel_one_id, lambda _: None)
    queue.create_channel(channel_two_id, lambda _: None)

    channel_one_msg = _make_mock_message(channel_one_id)
    channel_two_msg = _make_mock_message(channel_two_id)

    queue.put_nowait(channel_one_id, channel_one_msg)
    queue.put_nowait(channel_one_id, channel_one_msg)
    with pytest.raises(asyncio.QueueFull):
        queue.put_nowait(channel_one_id, channel_one_msg)
    queue.put_nowait(channel_two_id, channel_two_msg)
    queue.put_nowait(channel_two_id, channel_two_msg)
    with pytest.raises(asyncio.QueueFull):
        queue.put_nowait(channel_two_id, channel_two_msg)

    with pytest.raises(TimeoutError):
        async with asyncio.timeout(0.1):
            await queue.put(channel_one_id, channel_one_msg)

    assert queue.size(channel_one_id) == msg_size * 2

    add_task = asyncio.create_task(queue.put(channel_one_id, channel_one_msg))
    await asyncio.sleep(0)
    assert not add_task.done()
    assert queue.get_nowait() == channel_one_msg
    await asyncio.sleep(0)
    assert add_task.done()
    assert queue.get_nowait() == channel_two_msg


async def test_multi_channel_queue_force_message_on_full() -> None:
    """Test MultiplexerMultiChannelQueue getting full and forcing a message in."""
    msg_size = MOCK_MSG_SIZE + HEADER_SIZE
    # Max two mock messages per channel
    queue = MultiplexerMultiChannelQueue(msg_size * 2, msg_size, msg_size * 2)

    channel_one_id = _make_mock_channel_id()
    queue.create_channel(channel_one_id, lambda _: None)

    channel_one_msg = _make_mock_message(channel_one_id)

    queue.put_nowait(channel_one_id, channel_one_msg)
    queue.put_nowait(channel_one_id, channel_one_msg)
    with pytest.raises(asyncio.QueueFull):
        queue.put_nowait(channel_one_id, channel_one_msg)

    queue.put_nowait_force(channel_one_id, channel_one_msg)
    queue.put_nowait_force(channel_one_id, None)

    assert queue.size(channel_one_id) == msg_size * 3

    assert queue.get_nowait() == channel_one_msg
    assert queue.get_nowait() == channel_one_msg
    assert queue.get_nowait() == channel_one_msg
    assert queue.get_nowait() is None

    queue.delete_channel(channel_one_id)
    with pytest.raises(RuntimeError, match="does not exist or already closed"):
        queue.put_nowait_force(channel_one_id, channel_one_msg)


async def test_multi_channel_queue_round_robin_get() -> None:
    """Test MultiplexerMultiChannelQueue round robin get."""
    msg_size = MOCK_MSG_SIZE + HEADER_SIZE
    # Max two mock messages per channel
    queue = MultiplexerMultiChannelQueue(msg_size * 2, msg_size, msg_size * 2)
    channel_one_id = _make_mock_channel_id()
    channel_two_id = _make_mock_channel_id()
    channel_three_id = _make_mock_channel_id()
    queue.create_channel(channel_one_id, lambda _: None)
    queue.create_channel(channel_two_id, lambda _: None)
    queue.create_channel(channel_three_id, lambda _: None)

    channel_one_msg = _make_mock_message(channel_one_id)
    assert queue.empty(channel_one_id)
    await queue.put(channel_one_id, channel_one_msg)
    assert not queue.empty(channel_one_id)
    assert queue.size(channel_one_id) == len(channel_one_msg.data) + HEADER_SIZE

    channel_two_msg = _make_mock_message(channel_two_id)
    assert queue.empty(channel_two_id)
    await queue.put(channel_two_id, channel_two_msg)
    assert not queue.empty(channel_two_id)
    assert queue.size(channel_two_id) == len(channel_two_msg.data) + HEADER_SIZE

    channel_three_msg = _make_mock_message(channel_three_id)
    assert queue.empty(channel_three_id)
    queue.put_nowait(channel_three_id, channel_three_msg)
    assert not queue.empty(channel_three_id)
    assert queue.size(channel_three_id) == len(channel_three_msg.data) + HEADER_SIZE

    assert queue.get_nowait() == channel_one_msg
    assert queue.empty(channel_one_id)
    assert queue.size(channel_one_id) == 0

    assert queue.get_nowait() == channel_two_msg
    assert queue.empty(channel_two_id)
    assert queue.size(channel_two_id) == 0

    assert queue.get_nowait() == channel_three_msg
    assert queue.empty(channel_three_id)
    assert queue.size(channel_three_id) == 0

    with pytest.raises(asyncio.QueueEmpty):
        queue.get_nowait()

    with pytest.raises(TimeoutError):
        async with asyncio.timeout(0.1):
            await queue.get()

    queue.put_nowait(channel_two_id, channel_two_msg)
    queue.put_nowait(channel_three_id, channel_three_msg)
    queue.put_nowait(channel_one_id, channel_one_msg)
    queue.put_nowait(channel_one_id, channel_one_msg)
    queue.put_nowait(channel_three_id, channel_three_msg)
    queue.put_nowait(channel_two_id, channel_two_msg)

    msgs = [queue.get_nowait() for _ in range(6)]
    # Queue should be fair regardless of the order of the messages
    # coming in
    assert msgs == [
        channel_two_msg,
        channel_three_msg,
        channel_one_msg,
        channel_two_msg,
        channel_three_msg,
        channel_one_msg,
    ]

    with pytest.raises(asyncio.QueueEmpty):
        queue.get_nowait()


async def test_concurrent_get() -> None:
    """Test MultiplexerMultiChannelQueue concurrent get."""
    msg_size = MOCK_MSG_SIZE + HEADER_SIZE
    # Max two mock messages per channel
    queue = MultiplexerMultiChannelQueue(msg_size * 2, msg_size, msg_size * 2)
    channel_one_id = _make_mock_channel_id()
    channel_two_id = _make_mock_channel_id()
    channel_three_id = _make_mock_channel_id()

    queue.create_channel(channel_one_id, lambda _: None)
    queue.create_channel(channel_two_id, lambda _: None)
    queue.create_channel(channel_three_id, lambda _: None)

    channel_one_msg = _make_mock_message(channel_one_id)
    channel_two_msg = _make_mock_message(channel_two_id)
    channel_three_msg = _make_mock_message(channel_three_id)

    fetch_tasks = [asyncio.create_task(queue.get()) for _ in range(3)]

    await queue.put(channel_one_id, channel_one_msg)
    await queue.put(channel_two_id, channel_two_msg)
    await queue.put(channel_three_id, channel_three_msg)

    fetched_msgs = await asyncio.gather(*fetch_tasks)

    assert channel_one_msg in fetched_msgs
    assert channel_two_msg in fetched_msgs
    assert channel_three_msg in fetched_msgs

    with pytest.raises(asyncio.QueueEmpty):
        queue.get_nowait()


async def test_cancel_one_get() -> None:
    """Test the cancellation of a single `get` operation on multiplexer queue."""
    queue = MultiplexerMultiChannelQueue(100000, 10, 10000)
    reader = asyncio.create_task(queue.get())
    channel_one_id = _make_mock_channel_id()
    queue.create_channel(channel_one_id, lambda _: None)

    channel_one_msg1 = _make_mock_message(channel_one_id)
    channel_one_msg2 = _make_mock_message(channel_one_id)

    await asyncio.sleep(0)

    queue.put_nowait(channel_one_id, channel_one_msg1)
    queue.put_nowait(channel_one_id, channel_one_msg2)
    reader.cancel()

    with pytest.raises(asyncio.CancelledError):
        await reader

    assert await queue.get() == channel_one_msg1


async def test_reader_cancellation() -> None:
    """
    Test behavior of the MultiplexerMultiChannelQueue when a reader task is cancelled.

     Assertions:
        - The cancelled reader task raises asyncio.CancelledError.
        - The remaining reader tasks retrieve the messages from the queue in any order.
    """
    queue = MultiplexerMultiChannelQueue(100000, 10, 10000)
    channel_one_id = _make_mock_channel_id()
    queue.create_channel(channel_one_id, lambda _: None)
    channel_one_msg1 = _make_mock_message(channel_one_id)
    channel_one_msg2 = _make_mock_message(channel_one_id)

    async with asyncio.TaskGroup() as tg:
        reader1 = tg.create_task(queue.get())
        reader2 = tg.create_task(queue.get())
        reader3 = tg.create_task(queue.get())

        await asyncio.sleep(0)

        queue.put_nowait(channel_one_id, channel_one_msg1)
        queue.put_nowait(channel_one_id, channel_one_msg2)
        reader1.cancel()

        with pytest.raises(asyncio.CancelledError):
            await reader1

        await reader3

    # Any order is fine as long as we get both messages
    # since task order is not guaranteed
    assert {reader2.result(), reader3.result()} == {channel_one_msg1, channel_one_msg2}


async def test_put_cancel_race() -> None:
    """Test race between putting messages and cancelling the put operation."""
    msg_size = MOCK_MSG_SIZE + HEADER_SIZE
    # Max one message
    queue = MultiplexerMultiChannelQueue(msg_size, msg_size, msg_size)
    channel_one_id = _make_mock_channel_id()
    queue.create_channel(channel_one_id, lambda _: None)

    channel_one_msg_1 = _make_mock_message(channel_one_id)
    channel_one_msg_2 = _make_mock_message(channel_one_id)
    channel_one_msg_3 = _make_mock_message(channel_one_id)

    queue.put_nowait(channel_one_id, channel_one_msg_1)
    assert queue.get_nowait() == channel_one_msg_1
    assert queue.empty(channel_one_id)

    put_1 = asyncio.create_task(queue.put(channel_one_id, channel_one_msg_1))
    put_2 = asyncio.create_task(queue.put(channel_one_id, channel_one_msg_2))
    put_3 = asyncio.create_task(queue.put(channel_one_id, channel_one_msg_3))

    await asyncio.sleep(0)
    assert put_1.done()
    assert not put_2.done()
    assert not put_3.done()

    put_3.cancel()
    await asyncio.sleep(0)
    assert put_3.done()
    assert queue.get_nowait() == channel_one_msg_1
    await asyncio.sleep(0)
    assert queue.get_nowait() == channel_one_msg_2

    await put_2


async def test_putters_cleaned_up_correctly_on_cancellation() -> None:
    """Test that putters are cleaned up correctly when a put operation is canceled."""
    msg_size = MOCK_MSG_SIZE + HEADER_SIZE
    # Max one message
    queue = MultiplexerMultiChannelQueue(msg_size, msg_size, msg_size)
    channel_one_id = _make_mock_channel_id()
    queue.create_channel(channel_one_id, lambda _: None)
    channel_one_msg_1 = _make_mock_message(channel_one_id)
    channel_one_msg_2 = _make_mock_message(channel_one_id)

    queue.put_nowait(channel_one_id, channel_one_msg_1)

    put_task = asyncio.create_task(queue.put(channel_one_id, channel_one_msg_2))
    await asyncio.sleep(0)

    # Check that the putter is correctly removed from channel putters
    # the task is canceled.
    assert len(queue._channels[channel_one_id].putters) == 1
    put_task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await put_task
    assert len(queue._channels[channel_one_id].putters) == 0


async def test_getters_cleaned_up_correctly_on_cancellation() -> None:
    """Test getters are cleaned up correctly when a get operation is canceled."""
    msg_size = MOCK_MSG_SIZE + HEADER_SIZE
    # Max one message
    queue = MultiplexerMultiChannelQueue(msg_size, msg_size, msg_size)
    with pytest.raises(TimeoutError):
        async with asyncio.timeout(0.1):
            await queue.get()

    assert len(queue._getters) == 0


async def test_cancelled_when_putter_already_removed() -> None:
    """Test put operation is correctly cancelled when the putter is already removed."""
    msg_size = MOCK_MSG_SIZE + HEADER_SIZE
    # Max one message
    queue = MultiplexerMultiChannelQueue(msg_size, msg_size, msg_size)
    channel_one_id = _make_mock_channel_id()
    queue.create_channel(channel_one_id, lambda _: None)
    channel_one_msg_1 = _make_mock_message(channel_one_id)

    queue.put_nowait(channel_one_id, channel_one_msg_1)
    put_task = asyncio.create_task(queue.put(channel_one_id, channel_one_msg_1))
    await asyncio.sleep(0)

    queue.get_nowait()
    put_task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await put_task


async def test_multiple_getters_waiting_multiple_putters() -> None:
    """Test that multiple getters and putters are correctly handled."""
    msg_size = MOCK_MSG_SIZE + HEADER_SIZE
    # Max one message
    queue = MultiplexerMultiChannelQueue(msg_size, msg_size, msg_size)
    channel_one_id = _make_mock_channel_id()
    queue.create_channel(channel_one_id, lambda _: None)
    channel_one_msg_1 = _make_mock_message(channel_one_id)
    channel_one_msg_2 = _make_mock_message(channel_one_id)
    t1 = asyncio.create_task(queue.put(channel_one_id, channel_one_msg_1))
    t2 = asyncio.create_task(queue.put(channel_one_id, channel_one_msg_2))
    assert await queue.get() == channel_one_msg_1
    assert await queue.get() == channel_one_msg_2
    await t1
    await t2


async def test_get_cancelled_race() -> None:
    """Test cancelling a get operation while another get operation is in progress."""
    queue = MultiplexerMultiChannelQueue(10000000, 10, 10000)
    channel_one_id = _make_mock_channel_id()
    queue.create_channel(channel_one_id, lambda _: None)
    channel_one_msg_1 = _make_mock_message(channel_one_id)

    t1 = asyncio.create_task(queue.get())
    t2 = asyncio.create_task(queue.get())

    await asyncio.sleep(0)
    t1.cancel()
    await asyncio.sleep(0)
    assert t1.done()
    await queue.put(channel_one_id, channel_one_msg_1)
    await asyncio.sleep(0)
    assert await t2 == channel_one_msg_1


async def test_get_with_other_putters() -> None:
    """Test that a get operation is correctly handled when other putters are waiting."""
    loop = asyncio.get_running_loop()
    queue = MultiplexerMultiChannelQueue(10000000, 10, 10000)
    channel_one_id = _make_mock_channel_id()
    queue.create_channel(channel_one_id, lambda _: None)
    channel_one_msg_1 = _make_mock_message(channel_one_id)

    queue.put_nowait(channel_one_id, channel_one_msg_1)
    other_putter = loop.create_future()
    queue._channels[channel_one_id].putters.append(other_putter)

    assert await queue.get() == channel_one_msg_1
    assert other_putter.done()
    assert await other_putter is None

    await queue.put(channel_one_id, channel_one_msg_1)
    assert queue.get_nowait() == channel_one_msg_1


async def test_get_with_other_putter_already_one() -> None:
    """Test that a get operation is correctly handled when other putters are waiting."""
    loop = asyncio.get_running_loop()
    queue = MultiplexerMultiChannelQueue(10000000, 10, 10000)
    channel_one_id = _make_mock_channel_id()
    queue.create_channel(channel_one_id, lambda _: None)
    channel_one_msg_1 = _make_mock_message(channel_one_id)

    queue.put_nowait(channel_one_id, channel_one_msg_1)
    other_putter = loop.create_future()
    other_putter.set_result(None)
    queue._channels[channel_one_id].putters.append(other_putter)

    assert await queue.get() == channel_one_msg_1
    assert other_putter.done()
    assert await other_putter is None

    await queue.put(channel_one_id, channel_one_msg_1)
    assert queue.get_nowait() == channel_one_msg_1


async def test_single_channel_queue_under_water() -> None:
    """Test MultiplexerSingleChannelQueue under water."""
    msg_size = MOCK_MSG_SIZE + HEADER_SIZE
    under_water_callbacks: list[bool] = []

    def on_under_water(under_water: bool) -> None:
        under_water_callbacks.append(under_water)

    queue = MultiplexerSingleChannelQueue(
        msg_size * 10,
        msg_size * 2,
        msg_size * 4,
        on_under_water,
    )
    channel_id = _make_mock_channel_id()
    msg = _make_mock_message(channel_id)
    assert queue.qsize() == 0
    queue.put_nowait(msg)
    assert queue.qsize() == len(msg.data) + HEADER_SIZE
    assert not under_water_callbacks
    queue.put_nowait(msg)  # now 2 messages
    assert not under_water_callbacks
    queue.put_nowait(msg)  # now 3 messages
    assert not under_water_callbacks
    queue.put_nowait(msg)  # now 4 messages -- under water
    assert under_water_callbacks == [True]
    queue.put_nowait(msg)  # now 5 messages -- still under water
    assert under_water_callbacks == [True]
    queue.get_nowait()  # now 4 messages -- have not reached low watermark
    assert under_water_callbacks == [True]
    queue.get_nowait()  # now 3 messages -- have not reached low watermark
    assert under_water_callbacks == [True]
    queue.get_nowait()  # now 2 messages -- reached low watermark
    assert under_water_callbacks == [True, False]
    queue.get_nowait()  # now 1 message -- still below low watermark
    assert under_water_callbacks == [True, False]
    queue.get_nowait()  # now 0 messages -- empty
    assert under_water_callbacks == [True, False]
    queue.put_nowait(msg)  # now 1 message -- below high watermark
    assert under_water_callbacks == [True, False]
    queue.put_nowait(msg)  # now 2 messages -- still below high watermark
    assert under_water_callbacks == [True, False]
    queue.put_nowait(msg)  # now 3 messages -- still below high watermark
    assert under_water_callbacks == [True, False]
    queue.put_nowait(msg)  # now 4 messages -- reached high watermark
    assert under_water_callbacks == [True, False, True]
    queue.get_nowait()  # now 3 messages -- below high watermark, but still above low watermark
    assert under_water_callbacks == [True, False, True]
    queue.get_nowait()  # now 2 messages -- below high watermark and below low watermark
    assert under_water_callbacks == [True, False, True, False]


async def test_multi_channel_queue_under_water() -> None:
    """Test MultiplexerMultiChannelQueue under water."""
    msg_size = MOCK_MSG_SIZE + HEADER_SIZE
    under_water_callbacks: list[bool] = []

    def on_under_water(under_water: bool) -> None:
        under_water_callbacks.append(under_water)

    queue = MultiplexerMultiChannelQueue(
        msg_size * 10,
        msg_size * 2,
        msg_size * 4,
    )
    channel_id = _make_mock_channel_id()
    queue.create_channel(channel_id, on_under_water)
    msg = _make_mock_message(channel_id)
    assert queue.empty(channel_id)
    queue.put_nowait(channel_id, msg)
    assert not under_water_callbacks
    queue.put_nowait(channel_id, msg)  # now 2 messages
    assert not under_water_callbacks
    queue.put_nowait(channel_id, msg)  # now 3 messages
    assert not under_water_callbacks
    queue.put_nowait(channel_id, msg)  # now 4 messages -- under water
    assert under_water_callbacks == [True]
    queue.put_nowait(channel_id, msg)  # now 5 messages -- still under water
    assert under_water_callbacks == [True]
    queue.get_nowait()  # now 4 messages -- have not reached low watermark
    assert under_water_callbacks == [True]
    queue.get_nowait()  # now 3 messages -- have not reached low watermark
    assert under_water_callbacks == [True]
    queue.get_nowait()  # now 2 messages -- reached low watermark
    assert under_water_callbacks == [True, False]
    queue.get_nowait()  # now 1 message -- still below low watermark
    assert under_water_callbacks == [True, False]
    queue.get_nowait()  # now 0 messages -- empty
    assert under_water_callbacks == [True, False]
    queue.put_nowait(channel_id, msg)  # now 1 message -- below high watermark
    assert under_water_callbacks == [True, False]
    queue.put_nowait(channel_id, msg)  # now 2 messages -- still below high watermark
    assert under_water_callbacks == [True, False]
    queue.put_nowait(channel_id, msg)  # now 3 messages -- still below high watermark
    assert under_water_callbacks == [True, False]
    queue.put_nowait(channel_id, msg)  # now 4 messages -- reached high watermark
    assert under_water_callbacks == [True, False, True]
    queue.get_nowait()  # now 3 messages -- below high watermark, but still above low watermark
    assert under_water_callbacks == [True, False, True]
    queue.get_nowait()  # now 2 messages -- below high watermark and below low watermark
    assert under_water_callbacks == [True, False, True, False]


async def test_put_nowait_to_non_existent_multi_channel_queue() -> None:
    """Test writing to a non-existent channel."""
    queue = MultiplexerMultiChannelQueue(100000, 10, 1000)
    channel_id = _make_mock_channel_id()
    msg = _make_mock_message(channel_id)
    with pytest.raises(RuntimeError, match=f"Channel {channel_id} does not exist"):
        queue.put_nowait(channel_id, msg)


async def test_put_to_non_existent_multi_channel_queue() -> None:
    """Test writing to a non-existent channel."""
    queue = MultiplexerMultiChannelQueue(100000, 10, 1000)
    channel_id = _make_mock_channel_id()
    msg = _make_mock_message(channel_id)
    with pytest.raises(RuntimeError, match=f"Channel {channel_id} does not exist"):
        await queue.put(channel_id, msg)


async def test_multiple_delete_channel_is_forgiving() -> None:
    """Test a channel can be deleted multiple times."""
    queue = MultiplexerMultiChannelQueue(100000, 10, 1000)
    channel_id = _make_mock_channel_id()
    queue.create_channel(channel_id, lambda _: None)
    queue.delete_channel(channel_id)
    queue.delete_channel(channel_id)


async def test_delete_channel_when_queue_is_not_empty() -> None:
    """Test a channel can be deleted when its queue is not empty."""
    queue = MultiplexerMultiChannelQueue(100000, 10, 1000)
    channel_id = _make_mock_channel_id()
    queue.create_channel(channel_id, lambda _: None)
    queue.put_nowait(channel_id, _make_mock_message(channel_id))
    queue.delete_channel(channel_id)
    assert not queue.empty(channel_id)
    assert queue.get_nowait() is not None
    queue.delete_channel(channel_id)
    assert queue.empty(channel_id)


async def test_multiple_create_channel_raises() -> None:
    """Test the same channel can only be created once."""
    queue = MultiplexerMultiChannelQueue(100000, 10, 1000)
    channel_id = _make_mock_channel_id()
    queue.create_channel(channel_id, lambda _: None)
    with pytest.raises(RuntimeError, match=f"Channel {channel_id} already exists"):
        queue.create_channel(channel_id, lambda _: None)