import itertools

import pytest

from ... import _core
from ...testing import assert_checkpoints, wait_all_tasks_blocked

pytestmark = pytest.mark.filterwarnings(
    "ignore:.*UnboundedQueue:trio.TrioDeprecationWarning"
)


async def test_UnboundedQueue_basic():
    q = _core.UnboundedQueue()
    q.put_nowait("hi")
    assert await q.get_batch() == ["hi"]
    with pytest.raises(_core.WouldBlock):
        q.get_batch_nowait()
    q.put_nowait(1)
    q.put_nowait(2)
    q.put_nowait(3)
    assert q.get_batch_nowait() == [1, 2, 3]

    assert q.empty()
    assert q.qsize() == 0
    q.put_nowait(None)
    assert not q.empty()
    assert q.qsize() == 1

    stats = q.statistics()
    assert stats.qsize == 1
    assert stats.tasks_waiting == 0

    # smoke test
    repr(q)


async def test_UnboundedQueue_blocking():
    record = []
    q = _core.UnboundedQueue()

    async def get_batch_consumer():
        while True:
            batch = await q.get_batch()
            assert batch
            record.append(batch)

    async def aiter_consumer():
        async for batch in q:
            assert batch
            record.append(batch)

    for consumer in (get_batch_consumer, aiter_consumer):
        record.clear()
        async with _core.open_nursery() as nursery:
            nursery.start_soon(consumer)
            await _core.wait_all_tasks_blocked()
            stats = q.statistics()
            assert stats.qsize == 0
            assert stats.tasks_waiting == 1
            q.put_nowait(10)
            q.put_nowait(11)
            await _core.wait_all_tasks_blocked()
            q.put_nowait(12)
            await _core.wait_all_tasks_blocked()
            assert record == [[10, 11], [12]]
            nursery.cancel_scope.cancel()


async def test_UnboundedQueue_fairness():
    q = _core.UnboundedQueue()

    # If there's no-one else around, we can put stuff in and take it out
    # again, no problem
    q.put_nowait(1)
    q.put_nowait(2)
    assert q.get_batch_nowait() == [1, 2]

    result = None

    async def get_batch(q):
        nonlocal result
        result = await q.get_batch()

    # But if someone else is waiting to read, then they get dibs
    async with _core.open_nursery() as nursery:
        nursery.start_soon(get_batch, q)
        await _core.wait_all_tasks_blocked()
        q.put_nowait(3)
        q.put_nowait(4)
        with pytest.raises(_core.WouldBlock):
            q.get_batch_nowait()
    assert result == [3, 4]

    # If two tasks are trying to read, they alternate
    record = []

    async def reader(name):
        while True:
            record.append((name, await q.get_batch()))

    async with _core.open_nursery() as nursery:
        nursery.start_soon(reader, "a")
        await _core.wait_all_tasks_blocked()
        nursery.start_soon(reader, "b")
        await _core.wait_all_tasks_blocked()

        for i in range(20):
            q.put_nowait(i)
            await _core.wait_all_tasks_blocked()

        nursery.cancel_scope.cancel()

    assert record == list(zip(itertools.cycle("ab"), [[i] for i in range(20)]))


async def test_UnboundedQueue_trivial_yields():
    q = _core.UnboundedQueue()

    q.put_nowait(None)
    with assert_checkpoints():
        await q.get_batch()

    q.put_nowait(None)
    with assert_checkpoints():
        async for _ in q:  # noqa # pragma: no branch
            break


async def test_UnboundedQueue_no_spurious_wakeups():
    # If we have two tasks waiting, and put two items into the queue... then
    # only one task wakes up
    record = []

    async def getter(q, i):
        got = await q.get_batch()
        record.append((i, got))

    async with _core.open_nursery() as nursery:
        q = _core.UnboundedQueue()
        nursery.start_soon(getter, q, 1)
        await wait_all_tasks_blocked()
        nursery.start_soon(getter, q, 2)
        await wait_all_tasks_blocked()

        for i in range(10):
            q.put_nowait(i)
        await wait_all_tasks_blocked()

        assert record == [(1, list(range(10)))]

        nursery.cancel_scope.cancel()