"""纯 Python WebSocket 隧道后端 - 基于 aiohttp，无需外部二进制"""
from __future__ import annotations

import asyncio
import json
import uuid

from rich import print


async def _ws_server_async(port: int, host: str, auth: str):
    """WebSocket 隧道服务端核心逻辑"""
    import aiohttp
    from aiohttp import web

    # tunnel_id -> (ws, asyncio.Queue)
    tunnels: dict[str, tuple[web.WebSocketResponse, asyncio.Queue]] = {}
    # remote_port -> tcp_server
    tcp_servers: dict[int, asyncio.AbstractServer] = {}
    # 当前连接的客户端 ws
    client_ws: web.WebSocketResponse | None = None

    async def handle_ws(request):
        nonlocal client_ws
        ws = web.WebSocketResponse(max_msg_size=16 * 1024 * 1024)
        await ws.prepare(request)

        registered = False
        async for msg in ws:
            if msg.type == aiohttp.WSMsgType.TEXT:
                data = json.loads(msg.data)

                if data["type"] == "register":
                    if auth and data.get("auth") != auth:
                        await ws.send_json({"type": "error", "msg": "认证失败"})
                        await ws.close()
                        return ws
                    client_ws = ws
                    registered = True
                    print("[green]客户端已注册[/green]")

                    # 为每个 mapping 启动 TCP 监听
                    for m in data["mappings"]:
                        rport = m["remote_port"]
                        if rport in tcp_servers:
                            tcp_servers[rport].close()
                            await tcp_servers[rport].wait_closed()

                        async def on_tcp_connect(
                            reader, writer, _rport=rport
                        ):
                            tid = str(uuid.uuid4())
                            q: asyncio.Queue = asyncio.Queue()
                            tunnels[tid] = (ws, q)
                            await ws.send_json(
                                {
                                    "type": "connect",
                                    "tunnel_id": tid,
                                    "remote_port": _rport,
                                }
                            )

                            async def tcp_to_ws():
                                try:
                                    while True:
                                        chunk = await reader.read(65536)
                                        if not chunk:
                                            break
                                        await ws.send_bytes(
                                            tid.encode() + chunk
                                        )
                                except (ConnectionError, asyncio.CancelledError):
                                    pass

                            async def ws_to_tcp():
                                try:
                                    while True:
                                        payload = await q.get()
                                        if payload is None:
                                            break
                                        writer.write(payload)
                                        await writer.drain()
                                except (ConnectionError, asyncio.CancelledError):
                                    pass

                            t1 = asyncio.create_task(tcp_to_ws())
                            t2 = asyncio.create_task(ws_to_tcp())
                            try:
                                await asyncio.gather(t1, t2, return_exceptions=True)
                            finally:
                                tunnels.pop(tid, None)
                                writer.close()
                                try:
                                    await ws.send_json(
                                        {"type": "close", "tunnel_id": tid}
                                    )
                                except Exception:
                                    pass

                        srv = await asyncio.start_server(
                            on_tcp_connect, "0.0.0.0", rport
                        )
                        tcp_servers[rport] = srv
                        print(f"  [blue]TCP 监听 :{rport}[/blue]")

                    await ws.send_json({"type": "registered"})

            elif msg.type == aiohttp.WSMsgType.BINARY:
                # tunnel_id(36B) + payload
                tid = msg.data[:36].decode()
                payload = msg.data[36:]
                if tid in tunnels:
                    await tunnels[tid][1].put(payload)

            elif msg.type in (
                aiohttp.WSMsgType.ERROR,
                aiohttp.WSMsgType.CLOSE,
            ):
                break

        # 客户端断开，清理
        if registered:
            client_ws = None
            for tid, (_, q) in list(tunnels.items()):
                await q.put(None)
            tunnels.clear()
            for srv in tcp_servers.values():
                srv.close()
            tcp_servers.clear()
            print("[yellow]客户端已断开，TCP 监听已清理[/yellow]")

        return ws

    app = web.Application()
    app.router.add_get("/ws", handle_ws)

    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, host, port)
    await site.start()

    print(f"[green]Python WebSocket 隧道服务端，监听 {host}:{port}[/green]")
    print(f"[dim]客户端连接: maque relay <本机IP>:{port} <远程端口>:<本地端口> --method=python[/dim]")
    print("[dim]按 Ctrl+C 停止[/dim]\n")

    try:
        await asyncio.Event().wait()
    except asyncio.CancelledError:
        pass
    finally:
        await runner.cleanup()


def ws_server(port: int = 25565, host: str = "0.0.0.0", auth: str = None):
    """启动 Python WebSocket 隧道服务端"""
    try:
        asyncio.run(_ws_server_async(port, host, auth))
    except KeyboardInterrupt:
        print("\n[yellow]服务已停止[/yellow]")


async def _ws_client_async(
    server: str, mappings: list[tuple[int, str, int]], auth: str
):
    """WebSocket 隧道客户端核心逻辑"""
    import aiohttp

    # tunnel_id -> (reader, writer)
    local_conns: dict[str, tuple[asyncio.StreamReader, asyncio.StreamWriter]] = {}
    # remote_port -> (target_host, target_port)
    mapping_map = {m[0]: (m[1], m[2]) for m in mappings}

    url = f"ws://{server}/ws"
    print(f"[green]连接到 {url}[/green]")

    async with aiohttp.ClientSession() as session:
        async with session.ws_connect(url, max_msg_size=16 * 1024 * 1024) as ws:
            # 发送注册
            reg = {
                "type": "register",
                "mappings": [{"remote_port": rp} for rp, _, _ in mappings],
            }
            if auth:
                reg["auth"] = auth
            await ws.send_json(reg)

            server_host = server.split(":")[0]
            for rp, th, tp in mappings:
                print(f"  [blue]{server_host}:{rp}[/blue] -> [cyan]{th}:{tp}[/cyan]")
            print("[dim]按 Ctrl+C 断开[/dim]\n")

            async for msg in ws:
                if msg.type == aiohttp.WSMsgType.TEXT:
                    data = json.loads(msg.data)

                    if data["type"] == "registered":
                        print("[green]注册成功，等待连接...[/green]")

                    elif data["type"] == "error":
                        print(f"[red]服务端错误: {data['msg']}[/red]")
                        break

                    elif data["type"] == "connect":
                        tid = data["tunnel_id"]
                        rport = data["remote_port"]
                        target_host, target_port = mapping_map[rport]

                        try:
                            reader, writer = await asyncio.open_connection(
                                target_host, target_port
                            )
                            local_conns[tid] = (reader, writer)

                            async def forward_local(
                                _tid=tid, _reader=reader
                            ):
                                try:
                                    while True:
                                        chunk = await _reader.read(65536)
                                        if not chunk:
                                            break
                                        await ws.send_bytes(
                                            _tid.encode() + chunk
                                        )
                                except (ConnectionError, asyncio.CancelledError):
                                    pass
                                finally:
                                    local_conns.pop(_tid, None)

                            asyncio.create_task(forward_local())
                        except OSError as e:
                            print(
                                f"[red]连接 {target_host}:{target_port} 失败: {e}[/red]"
                            )
                            await ws.send_json(
                                {"type": "close", "tunnel_id": tid}
                            )

                    elif data["type"] == "close":
                        tid = data["tunnel_id"]
                        conn = local_conns.pop(tid, None)
                        if conn:
                            conn[1].close()

                elif msg.type == aiohttp.WSMsgType.BINARY:
                    tid = msg.data[:36].decode()
                    payload = msg.data[36:]
                    conn = local_conns.get(tid)
                    if conn:
                        conn[1].write(payload)
                        await conn[1].drain()

                elif msg.type in (
                    aiohttp.WSMsgType.ERROR,
                    aiohttp.WSMsgType.CLOSE,
                ):
                    break

    # 清理
    for _, writer in local_conns.values():
        writer.close()
    local_conns.clear()


def ws_client(
    server: str, mappings: list[tuple[int, str, int]], auth: str = None
):
    """启动 Python WebSocket 隧道客户端"""
    try:
        asyncio.run(_ws_client_async(server, mappings, auth))
    except KeyboardInterrupt:
        print("\n[yellow]中转已断开[/yellow]")
