from __future__ import annotations

import asyncio

import pytest

from toki_socket.communicator import type_name_of
from toki_socket.packets.message_common_pb2 import TestData as ProtoTestData
from toki_socket.tcp_client import TcpClient, connect_tcp
from toki_socket.tcp_server import TcpServer


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


@pytest.mark.asyncio
async def test_tcp_send_receive():
    received = asyncio.get_running_loop().create_future()
    server = TcpServer(
        "127.0.0.1",
        0,
        lambda reader, writer: TcpClient(reader, writer, 0, 0, parser_map()),
    )
    server.on_client_connected = lambda client: client.communicator.add_listener(
        type_name_of(ProtoTestData),
        lambda message: received.set_result(message) if not received.done() else None,
    )
    await server.start()
    client = await connect_tcp("127.0.0.1", server.port, 0, 0, parser_map())

    await client.communicator.send(ProtoTestData(index=1, message="hello"))
    result = await asyncio.wait_for(received, 2)

    assert result.index == 1
    assert result.message == "hello"

    await client.close()
    await server.stop()


@pytest.mark.asyncio
async def test_tcp_request_response():
    server = TcpServer(
        "127.0.0.1",
        0,
        lambda reader, writer: TcpClient(reader, writer, 0, 0, parser_map()),
    )
    server.on_client_connected = lambda client: client.communicator.add_request_listener(
        type_name_of(ProtoTestData),
        lambda req: ProtoTestData(index=req.index * 2, message="echo: " + req.message),
    )
    await server.start()
    client = await connect_tcp("127.0.0.1", server.port, 0, 0, parser_map())

    result = await client.communicator.send_request(
        ProtoTestData(index=21, message="single request"),
        ProtoTestData,
        timeout=2,
    )

    assert result.index == 42
    assert result.message == "echo: single request"

    await client.close()
    await server.stop()