from __future__ import annotations

import argparse
import asyncio
import ssl
import sys
from contextlib import suppress
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from toki_socket.communicator import type_name_of
from toki_socket.packets.message_common_pb2 import TestData
from toki_socket.tcp_client import connect_tcp, connect_tcp_tls
from toki_socket.ws_client import connect_ws, connect_wss

HOST = "127.0.0.1"
WS_PATH = "/"
CONNECT_WINDOW = 3.0
REQUEST_WINDOW = 2.0
SERVER_READY_DELAY = 0.05


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


async def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=["tcp", "ws", "tls", "wss"], default="tcp")
    parser.add_argument("--port", type=int, required=True)
    parser.add_argument("--phase", choices=["send-push", "requests"], default="send-push")
    parser.add_argument("--cert", default=None)
    args = parser.parse_args()

    print(f"INFO typeName python={type_name_of(TestData)}")

    try:
        client = await dial_with_retry(args.mode, args.port, args.cert)
    except Exception as exc:
        fail("setup", str(exc))
        raise SystemExit(1) from exc
    await asyncio.sleep(SERVER_READY_DELAY)

    try:
        if args.phase == "send-push":
            ok = await run_send_push(client)
        else:
            ok = await run_requests(client)
    finally:
        with suppress(Exception):
            await asyncio.wait_for(client.close(), 1.0)

    if not ok:
        raise SystemExit(1)


async def dial_with_retry(mode: str, port: int, cert: str | None):
    deadline = asyncio.get_running_loop().time() + CONNECT_WINDOW
    last_error: Exception | None = None
    while asyncio.get_running_loop().time() < deadline:
        try:
            return await asyncio.wait_for(dial(mode, port, cert), 0.3)
        except Exception as exc:
            last_error = exc
            await asyncio.sleep(0.1)
    raise TimeoutError(f"connect {mode}:{port} timed out: {last_error}")


async def dial(mode: str, port: int, cert: str | None):
    if mode == "tcp":
        return await connect_tcp(HOST, port, 0, 0, parser_map())
    if mode == "ws":
        return await connect_ws(HOST, port, WS_PATH, 0, 0, parser_map())
    if mode == "tls":
        if cert is None:
            raise ValueError("--cert is required for tls mode")
        ssl_ctx = ssl.create_default_context(cafile=cert)
        ssl_ctx.check_hostname = False
        return await connect_tcp_tls(HOST, port, ssl_ctx, 0, 0, parser_map())
    if mode == "wss":
        if cert is None:
            raise ValueError("--cert is required for wss mode")
        ssl_ctx = ssl.create_default_context(cafile=cert)
        ssl_ctx.check_hostname = False
        return await connect_wss(HOST, port, WS_PATH, ssl_ctx, 0, 0, parser_map())
    raise ValueError(f"unknown mode {mode!r}")


async def run_send_push(client) -> bool:
    loop = asyncio.get_running_loop()
    push = loop.create_future()
    client.communicator.add_listener(
        type_name_of(TestData),
        lambda message: push.set_result(message) if not push.done() else None,
    )

    try:
        await client.communicator.send(TestData(index=101, message="fire from python client"))
        passed("1", "fire-and-forget sent")

        message = await asyncio.wait_for(push, REQUEST_WINDOW)
        if message.index != 200 or message.message != "push from go server":
            fail("2", f"unexpected push index={message.index} message={message.message!r}")
            return False
        passed("2", "received push from go server")
        return True
    except Exception as exc:
        fail("2", str(exc))
        return False


async def run_requests(client) -> bool:
    try:
        single = await client.communicator.send_request(
            TestData(index=21, message="single request from python"),
            TestData,
            timeout=REQUEST_WINDOW,
        )
        if single.index != 42 or single.message != "echo: single request from python":
            fail("3", f"unexpected response index={single.index} message={single.message!r}")
            return False
        passed("3", "single request response matched")
    except Exception as exc:
        fail("3", str(exc))
        return False

    async def request_one(i: int) -> None:
        index = 30 + i
        message = f"multi request {i} from python"
        res = await client.communicator.send_request(
            TestData(index=index, message=message),
            TestData,
            timeout=REQUEST_WINDOW,
        )
        if res.index != index * 2 or res.message != "echo: " + message:
            raise AssertionError(f"request {i} got index={res.index} message={res.message!r}")

    try:
        await asyncio.gather(*(request_one(i) for i in range(5)))
        passed("4", "concurrent request responses matched")
        return True
    except Exception as exc:
        fail("4", str(exc))
        return False


def passed(scenario: str, detail: str) -> None:
    print(f"PASS scenario={scenario} detail={detail}")


def fail(scenario: str, error: str) -> None:
    print(f"FAIL scenario={scenario} error={error}")


if __name__ == "__main__":
    asyncio.run(main())