#!/usr/bin/env python3
"""WSL→Windows TCP relay for the Windows-side adb server.

Two modes:

  static:   wsl_tcp_relay.py <listen-port> <target-host> <target-port>
            Relays 127.0.0.1:<listen-port> -> <target-host>:<target-port>.

  dynamic:  wsl_tcp_relay.py --watch-adb <target-host>
            Polls `adb forward --list` (using the env-configured adb server),
            and for each host-side TCP port in the listing, ensures a relay
            127.0.0.1:<port> -> <target-host>:<port> exists. Adds/removes
            relays as forwards appear/disappear.

The dynamic mode handles `flutter run` choosing an ephemeral host port (and
sometimes rotating it across reconnect attempts) when adb is shared via
ADB_SERVER_SOCKET — flutter's adb-forward request lands on the Windows adb
server, which binds the listener on 0.0.0.0 of the Windows side. WSL-loopback
has nothing on that port, so flutter (connecting to 127.0.0.1:<port> on WSL)
fails. This relay fills that gap dynamically.
"""
from __future__ import annotations

import asyncio
import os
import sys
from typing import Dict, Set


async def _pipe(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    try:
        while data := await reader.read(65536):
            writer.write(data)
            await writer.drain()
    except (ConnectionResetError, BrokenPipeError, OSError):
        pass
    finally:
        try:
            writer.close()
            await writer.wait_closed()
        except Exception:
            pass


def _handler(target_host: str, target_port: int):
    async def handle(local_r, local_w):
        try:
            remote_r, remote_w = await asyncio.open_connection(target_host, target_port)
        except Exception as exc:
            print(f"[relay] connect {target_host}:{target_port} failed: {exc}", file=sys.stderr, flush=True)
            local_w.close()
            return
        await asyncio.gather(_pipe(local_r, remote_w), _pipe(remote_r, local_w))
    return handle


async def run_static(listen_port: int, target_host: str, target_port: int) -> None:
    server = await asyncio.start_server(_handler(target_host, target_port), "127.0.0.1", listen_port)
    print(f"[relay] static 127.0.0.1:{listen_port} -> {target_host}:{target_port}", flush=True)
    async with server:
        await server.serve_forever()


async def _adb_forward_ports() -> Set[int]:
    adb = os.environ.get("ADB", "adb")
    proc = await asyncio.create_subprocess_exec(
        adb, "forward", "--list",
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    out, _ = await proc.communicate()
    ports: Set[int] = set()
    for line in out.decode("utf-8", "replace").splitlines():
        parts = line.split()
        if len(parts) >= 2 and parts[1].startswith("tcp:"):
            try:
                ports.add(int(parts[1][4:]))
            except ValueError:
                pass
    return ports


async def run_dynamic(target_host: str, poll_seconds: float = 1.0) -> None:
    active: Dict[int, asyncio.AbstractServer] = {}
    print(f"[relay] dynamic watching adb forwards -> {target_host} (poll={poll_seconds}s)", flush=True)
    while True:
        try:
            target_ports = await _adb_forward_ports()
        except Exception as exc:
            print(f"[relay] adb listing failed: {exc}", file=sys.stderr, flush=True)
            await asyncio.sleep(poll_seconds)
            continue

        for port in target_ports - set(active):
            try:
                srv = await asyncio.start_server(_handler(target_host, port), "127.0.0.1", port)
                active[port] = srv
                print(f"[relay] + 127.0.0.1:{port} -> {target_host}:{port}", flush=True)
            except OSError as exc:
                print(f"[relay] bind 127.0.0.1:{port} failed: {exc}", file=sys.stderr, flush=True)

        for port in set(active) - target_ports:
            srv = active.pop(port)
            srv.close()
            try:
                await srv.wait_closed()
            except Exception:
                pass
            print(f"[relay] - 127.0.0.1:{port}", flush=True)

        await asyncio.sleep(poll_seconds)


def _usage() -> None:
    print(__doc__, file=sys.stderr)
    sys.exit(2)


async def _main() -> None:
    argv = sys.argv[1:]
    if not argv:
        _usage()
    if argv[0] == "--watch-adb":
        if len(argv) != 2:
            _usage()
        await run_dynamic(argv[1])
    else:
        if len(argv) != 3:
            _usage()
        await run_static(int(argv[0]), argv[1], int(argv[2]))


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