import pytest

from functools import partial
import errno

import attr

import trio
from trio.testing import memory_stream_pair, wait_all_tasks_blocked


@attr.s(hash=False, eq=False)
class MemoryListener(trio.abc.Listener):
    closed = attr.ib(default=False)
    accepted_streams = attr.ib(factory=list)
    queued_streams = attr.ib(factory=(lambda: trio.open_memory_channel(1)))
    accept_hook = attr.ib(default=None)

    async def connect(self):
        assert not self.closed
        client, server = memory_stream_pair()
        await self.queued_streams[0].send(server)
        return client

    async def accept(self):
        await trio.lowlevel.checkpoint()
        assert not self.closed
        if self.accept_hook is not None:
            await self.accept_hook()
        stream = await self.queued_streams[1].receive()
        self.accepted_streams.append(stream)
        return stream

    async def aclose(self):
        self.closed = True
        await trio.lowlevel.checkpoint()


async def test_serve_listeners_basic():
    listeners = [MemoryListener(), MemoryListener()]

    record = []

    def close_hook():
        # Make sure this is a forceful close
        assert trio.current_effective_deadline() == float("-inf")
        record.append("closed")

    async def handler(stream):
        await stream.send_all(b"123")
        assert await stream.receive_some(10) == b"456"
        stream.send_stream.close_hook = close_hook
        stream.receive_stream.close_hook = close_hook

    async def client(listener):
        s = await listener.connect()
        assert await s.receive_some(10) == b"123"
        await s.send_all(b"456")

    async def do_tests(parent_nursery):
        async with trio.open_nursery() as nursery:
            for listener in listeners:
                for _ in range(3):
                    nursery.start_soon(client, listener)

        await wait_all_tasks_blocked()

        # verifies that all 6 streams x 2 directions each were closed ok
        assert len(record) == 12

        parent_nursery.cancel_scope.cancel()

    async with trio.open_nursery() as nursery:
        l2 = await nursery.start(trio.serve_listeners, handler, listeners)
        assert l2 == listeners
        # This is just split into another function because gh-136 isn't
        # implemented yet
        nursery.start_soon(do_tests, nursery)

    for listener in listeners:
        assert listener.closed


async def test_serve_listeners_accept_unrecognized_error():
    for error in [KeyError(), OSError(errno.ECONNABORTED, "ECONNABORTED")]:
        listener = MemoryListener()

        async def raise_error():
            raise error

        listener.accept_hook = raise_error

        with pytest.raises(type(error)) as excinfo:
            await trio.serve_listeners(None, [listener])
        assert excinfo.value is error


async def test_serve_listeners_accept_capacity_error(autojump_clock, caplog):
    listener = MemoryListener()

    async def raise_EMFILE():
        raise OSError(errno.EMFILE, "out of file descriptors")

    listener.accept_hook = raise_EMFILE

    # It retries every 100 ms, so in 950 ms it will retry at 0, 100, ..., 900
    # = 10 times total
    with trio.move_on_after(0.950):
        await trio.serve_listeners(None, [listener])

    assert len(caplog.records) == 10
    for record in caplog.records:
        assert "retrying" in record.msg
        assert record.exc_info[1].errno == errno.EMFILE


async def test_serve_listeners_connection_nursery(autojump_clock):
    listener = MemoryListener()

    async def handler(stream):
        await trio.sleep(1)

    class Done(Exception):
        pass

    async def connection_watcher(*, task_status=trio.TASK_STATUS_IGNORED):
        async with trio.open_nursery() as nursery:
            task_status.started(nursery)
            await wait_all_tasks_blocked()
            assert len(nursery.child_tasks) == 10
            raise Done

    with pytest.raises(Done):
        async with trio.open_nursery() as nursery:
            handler_nursery = await nursery.start(connection_watcher)
            await nursery.start(
                partial(
                    trio.serve_listeners,
                    handler,
                    [listener],
                    handler_nursery=handler_nursery,
                )
            )
            for _ in range(10):
                nursery.start_soon(listener.connect)