#!/usr/bin/env python3
"""frp-tunnel CLI — 连接隧道服务器，转发本地端口 (async并发版, 二进制帧, WebSocket relay, uvloop+aiohttp优化)"""

import argparse
import asyncio
import json
import sys
import urllib.request
import urllib.error

try:
    import uvloop
    uvloop.install()
except ImportError:
    pass

try:
    import aiohttp
except ImportError:
    print("Installing aiohttp...", flush=True)
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "aiohttp", "-q"])
    import aiohttp

try:
    import websockets
except ImportError:
    print("Installing websockets...", flush=True)
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "websockets", "-q"])
    import websockets


DOMAIN = "fxtunnel.qzz.io"
DEFAULT_SERVER = "fxtunnel.qzz.io"
MAX_CONCURRENT = 16  # 最大并发请求数

# ── WebSocket relay state ──
# Track local WS connections by conn_id
_local_ws: dict = {}  # conn_id → websockets.WebSocketClientProtocol
_local_ws_tasks: dict = {}  # conn_id → asyncio.Task


def api_call(server, token, method, path, body=None):
    proto = "http" if "127.0.0.1" in server or "localhost" in server or ":" in server else "https"
    if proto == "http" and ":" not in server.split("://")[-1]:
        server = f"{server}:8095"
    url = f"{proto}://{server}{path}"
    data = json.dumps(body).encode() if body else None
    req = urllib.request.Request(url, data=data, method=method)
    req.add_header("User-Agent", "Mozilla/5.0 (compatible; fxtunnel-cli/1.0)")
    req.add_header("Authorization", f"Bearer {token}")
    if body:
        req.add_header("Content-Type", "application/json")
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            return json.loads(resp.read())
    except urllib.error.HTTPError as e:
        body_bytes = e.read()
        try:
            detail = json.loads(body_bytes).get("detail", str(e.reason))
        except Exception:
            detail = str(e.reason)
        raise RuntimeError(detail)


# 全局 aiohttp 连接池（复用TCP连接，省握手开销）
_http_session: aiohttp.ClientSession = None

def _get_session():
    global _http_session
    if _http_session is None or _http_session.closed:
        connector = aiohttp.TCPConnector(limit=32, ttl_dns_cache=300, keepalive_timeout=30)
        _http_session = aiohttp.ClientSession(connector=connector)
    return _http_session


async def _forward_one(ws, meta, req_body, local_port, sem):
    """转发单个HTTP请求到本地服务（aiohttp异步版）"""
    async with sem:
        rid = meta.get("req_id", "")
        method = meta["method"]
        path = meta["path"]
        headers = meta.get("headers", {})

        try:
            session = _get_session()
            local_url = f"http://127.0.0.1:{local_port}{path}"
            fwd_headers = {k: v for k, v in headers.items() if k.lower() not in ("host", "transfer-encoding")}

            async with session.request(
                method, local_url,
                headers=fwd_headers,
                data=req_body if req_body else None,
                timeout=aiohttp.ClientTimeout(total=30),
            ) as resp:
                resp_body = await resp.read()
                resp_headers = {k: v for k, v in resp.headers.items()}

            # Send binary response: [2 bytes meta_len][meta JSON][body]
            resp_meta = json.dumps({
                "req_id": rid,
                "status": resp.status, "headers": resp_headers,
            }).encode()
            frame = len(resp_meta).to_bytes(2, "big") + resp_meta + resp_body
            await ws.send(frame)
        except Exception as e:
            err_body = f"Local error: {e}".encode()
            try:
                resp_meta = json.dumps({
                    "req_id": rid,
                    "status": 502, "headers": {"Content-Type": "text/plain"},
                }).encode()
                frame = len(resp_meta).to_bytes(2, "big") + resp_meta + err_body
                await ws.send(frame)
            except Exception:
                pass


async def _ws_relay(server_ws, conn_id, local_port, path, headers, subprotocol=""):
    """建立本地WS连接并双向relay"""
    local_url = f"ws://127.0.0.1:{local_port}{path}"

    # Use subprotocol from server (browser requested)
    subprotocols = [subprotocol] if subprotocol else None

    # Filter headers for local connection
    skip = {"host", "transfer-encoding", "upgrade", "connection",
            "sec-websocket-key", "sec-websocket-version", "sec-websocket-extensions",
            "sec-websocket-protocol"}
    extra = {k: v for k, v in headers.items() if k.lower() not in skip}

    local_ws = None
    try:
        connect_kwargs = {"close_timeout": 3}
        if subprotocols:
            connect_kwargs["subprotocols"] = subprotocols
        if extra:
            try:
                local_ws = await websockets.connect(local_url, additional_headers=extra, **connect_kwargs)
            except TypeError:
                # websockets < 11 uses extra_headers
                local_ws = await websockets.connect(local_url, extra_headers=extra, **connect_kwargs)
        else:
            local_ws = await websockets.connect(local_url, **connect_kwargs)

        _local_ws[conn_id] = local_ws
        print(f"  [ws] {conn_id} → localhost:{local_port}{path}", flush=True)

        # Relay: local → server
        async for msg in local_ws:
            if isinstance(msg, str):
                await server_ws.send(json.dumps({
                    "type": "ws_text", "conn_id": conn_id, "data": msg
                }))
            elif isinstance(msg, bytes):
                meta = json.dumps({"type": "ws_bin", "conn_id": conn_id}).encode()
                frame = len(meta).to_bytes(2, "big") + meta + msg
                await server_ws.send(frame)

    except websockets.exceptions.ConnectionClosed:
        pass
    except Exception as e:
        print(f"  [ws] {conn_id} error: {e}", flush=True)
    finally:
        _local_ws.pop(conn_id, None)
        if local_ws:
            try:
                await local_ws.close()
            except Exception:
                pass
        try:
            await server_ws.send(json.dumps({"type": "ws_close", "conn_id": conn_id}))
        except Exception:
            pass
        print(f"  [ws] {conn_id} closed", flush=True)


async def run_tunnel(server, token, subdomain, local_port):
    proto = "ws" if "127.0.0.1" in server or "localhost" in server or ":" in server else "wss"
    if proto == "ws" and ":" not in server.split("://")[-1]:
        server = f"{server}:8095"
    ws_url = f"{proto}://{server}/ws/register?subdomain={subdomain}&token={token}"
    ssl_ctx = None if proto == "ws" else True

    print(f"Connecting to {server}...", flush=True)
    print(f"   Tunnel: https://{subdomain}.{DOMAIN} -> localhost:{local_port}", flush=True)
    print(f"   Max concurrency: {MAX_CONCURRENT}", flush=True)
    print(f"   Press Ctrl+C to stop\n", flush=True)

    sem = asyncio.Semaphore(MAX_CONCURRENT)

    while True:
        try:
            async with websockets.connect(
                ws_url, ssl=ssl_ctx, proxy=None,
                ping_interval=25, ping_timeout=10,
                close_timeout=5,
            ) as ws:
                print("Connected!", flush=True)

                tasks = set()

                async def heartbeat():
                    while True:
                        await asyncio.sleep(20)
                        try:
                            await ws.send("ping")
                        except Exception:
                            break

                hb = asyncio.create_task(heartbeat())

                try:
                    async for msg in ws:
                        # Text frame: ping/pong or WS relay control
                        if isinstance(msg, str):
                            if msg == "pong":
                                continue
                            # Try parse as WS relay message
                            try:
                                d = json.loads(msg)
                                mtype = d.get("type")
                                cid = d.get("conn_id")

                                if mtype == "ws_open":
                                    # New browser WS connection → open local WS
                                    task = asyncio.create_task(
                                        _ws_relay(ws, cid, local_port, d["path"], d.get("headers", {}), d.get("subprotocol", ""))
                                    )
                                    _local_ws_tasks[cid] = task
                                    task.add_done_callback(lambda t, c=cid: _local_ws_tasks.pop(c, None))

                                elif mtype == "ws_text" and cid in _local_ws:
                                    await _local_ws[cid].send(d.get("data", ""))

                                elif mtype == "ws_bin" and cid in _local_ws:
                                    # Should not happen (binary via text frame), but handle it
                                    pass

                                elif mtype == "ws_close" and cid in _local_ws:
                                    try:
                                        await _local_ws[cid].close(code=1000)
                                    except Exception:
                                        pass
                                    _local_ws.pop(cid, None)

                            except (json.JSONDecodeError, KeyError):
                                pass

                        # Binary frame: [2 bytes meta_len][meta JSON][body]
                        elif isinstance(msg, bytes) and len(msg) >= 2:
                            meta_len = int.from_bytes(msg[:2], "big")
                            if 2 + meta_len > len(msg):
                                continue
                            meta = json.loads(msg[2:2 + meta_len])
                            body = msg[2 + meta_len:]

                            if meta.get("type") == "ws_bin":
                                # Binary WS data from browser → forward to local WS
                                cid = meta.get("conn_id")
                                if cid and cid in _local_ws:
                                    try:
                                        await _local_ws[cid].send(body)
                                    except Exception:
                                        pass
                            else:
                                # HTTP response (existing logic)
                                req_body = body
                                task = asyncio.create_task(_forward_one(ws, meta, req_body, local_port, sem))
                                tasks.add(task)
                                task.add_done_callback(tasks.discard)
                finally:
                    hb.cancel()
                    # Close all local WS connections
                    for cid, lws in list(_local_ws.items()):
                        try:
                            await lws.close(code=1001)
                        except Exception:
                            pass
                    _local_ws.clear()
                    # Cancel all WS relay tasks
                    for cid, task in list(_local_ws_tasks.items()):
                        task.cancel()
                    _local_ws_tasks.clear()
                    if tasks:
                        await asyncio.gather(*tasks, return_exceptions=True)

        except websockets.exceptions.ConnectionClosed as e:
            print(f"Disconnected ({e}). Reconnecting in 3s...", flush=True)
        except Exception as e:
            print(f"Error: {e}. Reconnecting in 3s...", flush=True)
        await asyncio.sleep(3)

    # Cleanup aiohttp session on exit
    global _http_session
    if _http_session and not _http_session.closed:
        await _http_session.close()
        _http_session = None


def cmd_expose(args):
    server = args.server
    token = args.token
    subdomain = args.sub
    local_port = args.local_port
    name = args.name or subdomain or ""

    body = {"port": local_port, "name": name}
    if subdomain:
        body["subdomain"] = subdomain

    try:
        result = api_call(server, token, "POST", "/api/tunnels", body)
        subdomain = result["subdomain"]
        print(f"Tunnel created: {result['url']}", flush=True)
    except RuntimeError as e:
        if "Taken" in str(e):
            if not subdomain:
                print("Error: no subdomain and auto-generate failed", flush=True)
                return
            print(f"Tunnel {subdomain} already exists, connecting...", flush=True)
        else:
            print(f"Error: {e}", flush=True)
            return

    try:
        asyncio.run(run_tunnel(server, token, subdomain, local_port))
    except KeyboardInterrupt:
        print("\nStopped", flush=True)
    except Exception as e:
        print(f"\nError: {e}", flush=True)


def cmd_list(args):
    try:
        result = api_call(args.server, args.token, "GET", "/api/tunnels")
    except RuntimeError as e:
        print(f"Error: {e}", flush=True)
        return
    tunnels = result["tunnels"]
    if not tunnels:
        print("No tunnels", flush=True)
        return
    print(f"{'ID':>4}  {'Status':<8}  {'Subdomain':<20}  {'Port':<6}  {'Name'}", flush=True)
    print("-" * 60, flush=True)
    for t in tunnels:
        print(f"{t['id']:>4}  {t['status']:<8}  {t['subdomain']}.{DOMAIN:<14}  {t['port']:<6}  {t['name']}", flush=True)


def cmd_delete(args):
    try:
        api_call(args.server, args.token, "DELETE", f"/api/tunnels/{args.tunnel_id}")
        print(f"Tunnel #{args.tunnel_id} deleted", flush=True)
    except RuntimeError as e:
        print(f"Error: {e}", flush=True)


def main():
    parser = argparse.ArgumentParser(description="frp-tunnel CLI (async, binary, ws-relay)")
    parser.add_argument("--server", default=DEFAULT_SERVER)
    parser.add_argument("--token", default="fx_tourist")

    sub = parser.add_subparsers(dest="command")

    p_expose = sub.add_parser("expose", help="Expose local port")
    p_expose.add_argument("local_port", type=int)
    p_expose.add_argument("--sub", help="Subdomain (auto if empty)")
    p_expose.add_argument("--name", default="")
    p_expose.set_defaults(func=cmd_expose)

    p_list = sub.add_parser("list", help="List tunnels")
    p_list.set_defaults(func=cmd_list)

    p_del = sub.add_parser("delete", help="Delete tunnel")
    p_del.add_argument("tunnel_id", type=int)
    p_del.set_defaults(func=cmd_delete)

    args = parser.parse_args()
    if not args.command:
        parser.print_help()
        sys.exit(1)
    args.func(args)


if __name__ == "__main__":
    main()