from __future__ import annotations

import asyncio
import ssl
from pathlib import Path

import pytest

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


def _make_server_ssl_context() -> ssl.SSLContext:
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    certs_dir = Path(__file__).parent / "certs"
    ctx.load_cert_chain(certs_dir / "server.crt", certs_dir / "server.key")
    return ctx


def _make_client_ssl_context() -> ssl.SSLContext:
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    certs_dir = Path(__file__).parent / "certs"
    ctx.load_verify_locations(certs_dir / "server.crt")
    ctx.check_hostname = False
    return ctx


async def _wait_for_clients(server: TcpServer, count: int = 1, timeout: float = 1.0) -> None:
    loop = asyncio.get_running_loop()
    deadline = loop.time() + timeout
    while loop.time() < deadline:
        if len(server.clients()) >= count:
            return
        await asyncio.sleep(0.01)
    raise TimeoutError("server did not see expected client count")


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()


@pytest.mark.asyncio
async def test_tcp_tls_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()),
        ssl_context=_make_server_ssl_context(),
    )
    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_tls(
        "127.0.0.1",
        server.port,
        _make_client_ssl_context(),
        0,
        0,
        parser_map(),
    )

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

    assert result.index == 3
    assert result.message == "hello over tls"

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


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

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

    assert result.index == 48
    assert result.message == "tls echo: secure request"

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


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

    await _wait_for_clients(server)
    await server.broadcast(ProtoTestData(index=200, message="push from python server"))
    result = await asyncio.wait_for(pushed, 2)

    assert result.index == 200
    assert result.message == "push from python server"

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


@pytest.mark.asyncio
async def test_tcp_concurrent_requests():
    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())

    async def do_request(i: int) -> None:
        index = 30 + i
        message = f"request {i}"
        result = await client.communicator.send_request(
            ProtoTestData(index=index, message=message),
            ProtoTestData,
            timeout=2,
        )
        assert result.index == index * 2
        assert result.message == f"echo: {message}"

    await asyncio.gather(*[do_request(i) for i in range(5)])
    await client.close()
    await server.stop()


@pytest.mark.asyncio
async def test_tcp_inbound_backpressure_pauses_read():
    server_client = None
    
    def on_server_connect(client):
        nonlocal server_client
        server_client = client

    server = TcpServer(
        "127.0.0.1",
        0,
        lambda reader, writer: TcpClient(reader, writer, 0, 0, parser_map()),
    )
    server.on_client_connected = on_server_connect
    await server.start()
    
    client = await connect_tcp("127.0.0.1", server.port, 0, 0, parser_map())
    await _wait_for_clients(server, count=1, timeout=2.0)
    
    handler_promise = asyncio.Event()

    async def blocking_handler(msg, nonce=0):
        await handler_promise.wait()
        return ProtoTestData(index=100)

    server_client.communicator.add_request_listener(type_name_of(ProtoTestData), blocking_handler)

    first_req_task = asyncio.create_task(
        client.communicator.send_request(
            ProtoTestData(index=1, message="first"),
            ProtoTestData,
            timeout=5.0,
        )
    )
    await asyncio.sleep(0.05)

    for i in range(2, 66):
        await client.communicator.send(ProtoTestData(index=i, message="heavy"))

    await client.communicator.send(ProtoTestData(index=66, message="overflow"))
    await asyncio.sleep(0.1)

    assert not server_client._read_task.done(), "read loop should not be finished"
    
    handler_promise.set()
    await asyncio.sleep(0.05)

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