from __future__ import annotations

import asyncio

import pytest

from toki_socket.communicator import Communicator, type_name_of
from toki_socket.packets.message_common_pb2 import PacketBase, TestData as ProtoTestData


class FakeTransport:
    def __init__(self) -> None:
        self.packets: list[PacketBase] = []
        self.closed = False

    async def write_packet(self, base: PacketBase) -> None:
        self.packets.append(base)

    async def close(self) -> None:
        self.closed = True


def parser_map():
    return {type_name_of(ProtoTestData): ProtoTestData.FromString}


@pytest.mark.asyncio
async def test_add_listener_exclusivity():
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())
    communicator.add_listener(type_name_of(ProtoTestData), lambda _: None)

    with pytest.raises(ValueError):
        communicator.add_request_listener(type_name_of(ProtoTestData), lambda req: req)

    await communicator.close()


@pytest.mark.asyncio
async def test_send_and_receive():
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())
    received = asyncio.get_running_loop().create_future()
    communicator.add_listener(type_name_of(ProtoTestData), received.set_result)

    await communicator.send(ProtoTestData(index=7, message="hello"))
    assert len(transport.packets) == 1
    packet = transport.packets[0]
    assert packet.typeName == "TestData"
    assert packet.nonce == 1

    communicator.on_received_data(packet.typeName, packet.data, packet.nonce, 0)
    result = await asyncio.wait_for(received, 1)
    assert result.index == 7
    assert result.message == "hello"

    await communicator.close()


@pytest.mark.asyncio
async def test_send_request_response():
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())

    task = asyncio.create_task(
        communicator.send_request(ProtoTestData(index=5, message="request"), ProtoTestData, timeout=1)
    )
    for _ in range(20):
        if transport.packets:
            break
        await asyncio.sleep(0.01)

    assert len(transport.packets) == 1
    request = transport.packets[0]
    response = ProtoTestData(index=10, message="echo: request")
    communicator.on_received_data(
        type_name_of(ProtoTestData),
        response.SerializeToString(),
        nonce=2,
        response_nonce=request.nonce,
    )

    result = await task
    assert result.index == 10
    assert result.message == "echo: request"

    await communicator.close()


@pytest.mark.asyncio
async def test_shutdown():
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())
    await communicator.close()

    assert not communicator.is_alive()
    assert transport.closed


@pytest.mark.asyncio
async def test_shutdown_public_method():
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())
    communicator.shutdown()
    await communicator.wait_closed()
    await asyncio.sleep(0)

    assert not communicator.is_alive()