from __future__ import annotations

import asyncio
import os
import re
import ssl
import sys
from pathlib import Path

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

from proto_socket.communicator import type_name_of
from proto_socket.packets.message_common_pb2 import TestData
from proto_socket.tcp_server import TcpServer
from proto_socket.ws_server import WsServer

HOST = "127.0.0.1"
PYTHON_DART_TCP_PORT = 29494
PYTHON_DART_WS_PORT = 29496
TLS_TCP_PORT = 29498
WSS_PORT = 29500
WS_PATH = "/"
PROCESS_TIMEOUT = 20.0
SERVER_OBSERVATION_WINDOW = 0.2
CERT_PATH = Path(__file__).resolve().parents[1] / "test" / "certs" / "server.crt"
KEY_PATH = Path(__file__).resolve().parents[1] / "test" / "certs" / "server.key"


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


async def main() -> None:
    print(f"INFO typeName python={type_name_of(TestData)}")
    selected_transport = os.environ.get("PROTO_SOCKET_CROSSTEST_TRANSPORT")
    try:
        if selected_transport is None or selected_transport == "tcp":
            await run_tcp_send_push()
            await asyncio.sleep(0.15)
            await run_tcp_requests()
            await asyncio.sleep(0.15)
            await run_tcp_legacy_simple_receive()
            await asyncio.sleep(0.15)
        if selected_transport is None or selected_transport == "ws":
            await run_ws_send_push()
            await asyncio.sleep(0.15)
            await run_ws_requests()
            await asyncio.sleep(0.15)
        if selected_transport is None or selected_transport in {"tls", "wss"}:
            await run_tls(selected_transport)
    except Exception as exc:
        print(f"FAIL crosstest error={exc}", file=sys.stderr)
        raise SystemExit(1) from exc
    print("PASS all python-server/dart-client crosstests passed")


async def run_tcp_send_push() -> None:
    received = asyncio.get_running_loop().create_future()
    server = TcpServer(HOST, PYTHON_DART_TCP_PORT, interval_sec=0, wait_sec=0, parser_map=parser_map())

    def on_connected(client):
        def on_message(data):
            print(f"SERVER_RECEIVED index={data.index} message={data.message}")
            valid = data.index == 101 and data.message == "fire from dart client"
            if not received.done():
                received.set_result(valid)
            if valid:
                asyncio.create_task(
                    client.communicator.send(TestData(index=200, message="push from python server"))
                )

        client.communicator.add_listener(type_name_of(TestData), on_message)

    server.on_client_connected = on_connected
    await server.start()
    try:
        await run_dart_client("tcp", PYTHON_DART_TCP_PORT, "send-push", {"1", "2"})
        ok = await asyncio.wait_for(received, SERVER_OBSERVATION_WINDOW)
        if not ok:
            raise RuntimeError("TCP send-push server received unexpected data")
    finally:
        await server.stop()


async def run_tcp_requests() -> None:
    server = TcpServer(HOST, PYTHON_DART_TCP_PORT, interval_sec=0, wait_sec=0, parser_map=parser_map())
    server.on_client_connected = lambda client: client.communicator.add_request_listener(
        type_name_of(TestData),
        lambda req: TestData(index=req.index * 2, message="echo: " + req.message),
    )
    await server.start()
    try:
        await run_dart_client("tcp", PYTHON_DART_TCP_PORT, "requests", {"3", "4"})
    finally:
        await server.stop()


async def run_tcp_legacy_simple_receive() -> None:
    received: asyncio.Future[bool] = asyncio.get_running_loop().create_future()
    server = TcpServer(HOST, PYTHON_DART_TCP_PORT, interval_sec=0, wait_sec=0, parser_map=parser_map())

    def on_connected(client):
        def on_message(data):
            if not received.done():
                received.set_result(data.index == 101)
        client.communicator.add_listener(type_name_of(TestData), on_message)

    server.on_client_connected = on_connected
    await server.start()
    try:
        await run_dart_client("tcp", PYTHON_DART_TCP_PORT, "legacy-receive", set())
        try:
            ok = await asyncio.wait_for(received, SERVER_OBSERVATION_WINDOW)
        except asyncio.TimeoutError:
            raise RuntimeError("TCP legacy simple-name server did not receive packet")
        if not ok:
            raise RuntimeError("TCP legacy simple-name server received unexpected data")
        print("PASS scenario=tcp-legacy-simple-receive detail=server-received-legacy-alias")
    finally:
        await server.stop()


async def run_ws_send_push() -> None:
    received = asyncio.get_running_loop().create_future()
    server = WsServer(HOST, PYTHON_DART_WS_PORT, WS_PATH, interval_sec=0, wait_sec=0, parser_map=parser_map())

    def on_connected(client):
        def on_message(data):
            print(f"SERVER_RECEIVED index={data.index} message={data.message}")
            valid = data.index == 101 and data.message == "fire from dart client"
            if not received.done():
                received.set_result(valid)
            if valid:
                asyncio.create_task(
                    client.communicator.send(TestData(index=200, message="push from python server"))
                )

        client.communicator.add_listener(type_name_of(TestData), on_message)

    server.on_client_connected = on_connected
    await server.start()
    try:
        await run_dart_client("ws", PYTHON_DART_WS_PORT, "send-push", {"1", "2"})
        ok = await asyncio.wait_for(received, SERVER_OBSERVATION_WINDOW)
        if not ok:
            raise RuntimeError("WS send-push server received unexpected data")
    finally:
        await server.stop()


async def run_ws_requests() -> None:
    server = WsServer(HOST, PYTHON_DART_WS_PORT, WS_PATH, interval_sec=0, wait_sec=0, parser_map=parser_map())
    server.on_client_connected = lambda client: client.communicator.add_request_listener(
        type_name_of(TestData),
        lambda req: TestData(index=req.index * 2, message="echo: " + req.message),
    )
    await server.start()
    try:
        await run_dart_client("ws", PYTHON_DART_WS_PORT, "requests", {"3", "4"})
    finally:
        await server.stop()


def _make_server_ssl_context() -> ssl.SSLContext:
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ctx.load_cert_chain(CERT_PATH, KEY_PATH)
    return ctx


async def run_tls(selected_transport: str | None) -> None:
    ssl_context = _make_server_ssl_context()
    if selected_transport is None or selected_transport == "tls":
        await run_tls_tcp_send_push(ssl_context)
        await asyncio.sleep(0.15)
        await run_tls_tcp_requests(ssl_context)
        await asyncio.sleep(0.15)
    if selected_transport is None or selected_transport == "wss":
        await run_wss_send_push(ssl_context)
        await asyncio.sleep(0.15)
        await run_wss_requests(ssl_context)


async def run_tls_tcp_send_push(ssl_context: ssl.SSLContext) -> None:
    received = asyncio.get_running_loop().create_future()
    server = TcpServer(
        HOST,
        TLS_TCP_PORT,
        interval_sec=0,
        wait_sec=0,
        parser_map=parser_map(),
        ssl_context=ssl_context,
    )

    def on_connected(client):
        def on_message(data):
            print(f"SERVER_RECEIVED index={data.index} message={data.message}")
            valid = data.index == 101 and data.message == "fire from dart client"
            if not received.done():
                received.set_result(valid)
            if valid:
                asyncio.create_task(
                    client.communicator.send(TestData(index=200, message="push from python server"))
                )

        client.communicator.add_listener(type_name_of(TestData), on_message)

    server.on_client_connected = on_connected
    await server.start()
    try:
        await run_dart_client("tls", TLS_TCP_PORT, "send-push", {"1", "2"}, CERT_PATH)
        ok = await asyncio.wait_for(received, SERVER_OBSERVATION_WINDOW)
        if not ok:
            raise RuntimeError("TLS TCP send-push server received unexpected data")
    finally:
        await server.stop()


async def run_tls_tcp_requests(ssl_context: ssl.SSLContext) -> None:
    server = TcpServer(
        HOST,
        TLS_TCP_PORT,
        interval_sec=0,
        wait_sec=0,
        parser_map=parser_map(),
        ssl_context=ssl_context,
    )
    server.on_client_connected = lambda client: client.communicator.add_request_listener(
        type_name_of(TestData),
        lambda req: TestData(index=req.index * 2, message="echo: " + req.message),
    )
    await server.start()
    try:
        await run_dart_client("tls", TLS_TCP_PORT, "requests", {"3", "4"}, CERT_PATH)
    finally:
        await server.stop()


async def run_wss_send_push(ssl_context: ssl.SSLContext) -> None:
    received = asyncio.get_running_loop().create_future()
    server = WsServer(
        HOST,
        WSS_PORT,
        WS_PATH,
        interval_sec=0,
        wait_sec=0,
        parser_map=parser_map(),
        ssl_context=ssl_context,
    )

    def on_connected(client):
        def on_message(data):
            print(f"SERVER_RECEIVED index={data.index} message={data.message}")
            valid = data.index == 101 and data.message == "fire from dart client"
            if not received.done():
                received.set_result(valid)
            if valid:
                asyncio.create_task(
                    client.communicator.send(TestData(index=200, message="push from python server"))
                )

        client.communicator.add_listener(type_name_of(TestData), on_message)

    server.on_client_connected = on_connected
    await server.start()
    try:
        await run_dart_client("wss", WSS_PORT, "send-push", {"1", "2"}, CERT_PATH)
        ok = await asyncio.wait_for(received, SERVER_OBSERVATION_WINDOW)
        if not ok:
            raise RuntimeError("WSS send-push server received unexpected data")
    finally:
        await server.stop()


async def run_wss_requests(ssl_context: ssl.SSLContext) -> None:
    server = WsServer(
        HOST,
        WSS_PORT,
        WS_PATH,
        interval_sec=0,
        wait_sec=0,
        parser_map=parser_map(),
        ssl_context=ssl_context,
    )
    server.on_client_connected = lambda client: client.communicator.add_request_listener(
        type_name_of(TestData),
        lambda req: TestData(index=req.index * 2, message="echo: " + req.message),
    )
    await server.start()
    try:
        await run_dart_client("wss", WSS_PORT, "requests", {"3", "4"}, CERT_PATH)
    finally:
        await server.stop()


async def run_dart_client(
    mode: str,
    port: int,
    phase: str,
    expected: set[str],
    cert: Path | None = None,
) -> None:
    dart_dir = dart_package_dir()
    args = [
        "run",
        "crosstest/python_dart_client.dart",
        f"--mode={mode}",
        f"--port={port}",
        f"--phase={phase}",
    ]
    if cert is not None:
        args.append(f"--cert={cert}")
    process = await asyncio.create_subprocess_exec(
        "dart",
        *args,
        cwd=dart_dir,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    try:
        stdout, stderr = await asyncio.wait_for(process.communicate(), PROCESS_TIMEOUT)
    except TimeoutError as exc:
        process.kill()
        await process.wait()
        raise TimeoutError(f"dart client {mode}/{phase} timed out") from exc

    result_lines: list[str] = []
    for line in stdout.decode().splitlines():
        print(line)
        if line.startswith(("PASS ", "FAIL ")):
            result_lines.append(line)
    for line in stderr.decode().splitlines():
        print(line, file=sys.stderr)

    validate_result_lines(f"dart-client {mode}/{phase}", process.returncode, result_lines, expected)


def validate_result_lines(label: str, return_code: int | None, lines: list[str], expected: set[str]) -> None:
    failed = [line for line in lines if line.startswith("FAIL ")]
    passed: set[str] = set()
    for line in lines:
        if line.startswith("PASS "):
            match = re.search(r"scenario=([^ ]+)", line)
            if match:
                passed.add(match.group(1))
    missing = expected - passed
    if return_code or failed or missing:
        raise RuntimeError(f"{label} failed returnCode={return_code} failed={failed} missing={sorted(missing)}")


def dart_package_dir() -> Path:
    repo_root = Path(__file__).resolve().parents[2]
    candidate = repo_root / "dart"
    if (candidate / "pubspec.yaml").is_file():
        return candidate
    raise RuntimeError(f"cannot resolve dart package directory from {candidate}")


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