"""
MeshBay Node POC — Spike 4: NAT Traversal (v2)

Tests:
  A. STUN  — external address discovery (pure UDP, no lib)
  B. NAT type probe — does external port change per destination? (cone vs symmetric)
  C. Bidirectional UDP hole punching:
       node → meshbay.org:19002  (creates NAT entry for that dest)
       meshbay.org receives, echoes to the source addr it saw
       node receives echo  → bidirectional UDP confirmed
  D. UPnP  — attempt port mapping on SFR box
  E. Update hub endpoint_hint
"""

import asyncio
import json
import os
import socket
import struct
from pathlib import Path

import httpx

PASS = "✓"; FAIL = "✗"; SKIP = "–"

LOCAL_PORT      = 19000
MESHBAY_IP      = "164.132.246.44"   # meshbay.org resolved
MESHBAY_ECHO_PORT = 19002
STATE_FILE      = Path("node_state.json")
HUB_URL         = "http://meshbay.org"
MESH_SERVER     = "cbesson@meshbay.org"

STUN_SERVERS = [
    ("stun.cloudflare.com", 3478),
    ("stun.l.google.com",   19302),
]

# ── STUN (pure UDP) ────────────────────────────────────────────────────────────

def stun_query(local_port: int, stun_host: str, stun_port: int) -> tuple[str|None, int|None]:
    """Single STUN query from local_port. Returns (ext_ip, ext_port) or (None, None)."""
    MAGIC = 0x2112A442
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.settimeout(3)
        sock.bind(('', local_port))
        txn = os.urandom(12)
        sock.sendto(struct.pack('>HHI12s', 0x0001, 0, MAGIC, txn), (stun_host, stun_port))
        data, _ = sock.recvfrom(1024)
        sock.close()
        msg_type, msg_len, magic, _ = struct.unpack_from('>HHI12s', data)
        if msg_type != 0x0101 or magic != MAGIC:
            return None, None
        offset = 20
        while offset < 20 + msg_len:
            atype, alen = struct.unpack_from('>HH', data, offset)
            offset += 4
            if atype == 0x0020:
                family = struct.unpack_from('>xB', data, offset)[0]
                if family == 0x01:
                    xport = struct.unpack_from('>H', data, offset + 2)[0]
                    xaddr = struct.unpack_from('>I', data, offset + 4)[0]
                    return (socket.inet_ntoa(struct.pack('>I', xaddr ^ MAGIC)),
                            xport ^ (MAGIC >> 16))
            offset += alen + (4 - alen % 4) % 4
    except Exception:
        pass
    return None, None


# ── UPnP ──────────────────────────────────────────────────────────────────────

def try_upnp(port: int) -> tuple[str|None, int|None]:
    try:
        import miniupnpc
        u = miniupnpc.UPnP()
        u.discoverdelay = 500
        if u.discover() == 0:
            return None, None
        u.selectigd()
        ext_ip = u.externalipaddress()
        local_ip = u.lanaddr
        if u.addportmapping(port, 'TCP', local_ip, port, 'MeshBay POC', ''):
            return ext_ip, port
    except Exception:
        pass
    return None, None


# ── Bidirectional UDP hole punch test ─────────────────────────────────────────

async def start_echo_server_on_meshbay() -> asyncio.subprocess.Process:
    """SSH to meshbay.org and start a one-shot UDP echo server."""
    echo_script = (
        f"python3 -c \""
        f"import socket;"
        f"s=socket.socket(socket.AF_INET,socket.SOCK_DGRAM);"
        f"s.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,1);"
        f"s.bind(('0.0.0.0',{MESHBAY_ECHO_PORT}));"
        f"s.settimeout(15);"
        f"print('ECHO_READY',flush=True);"
        f"data,addr=s.recvfrom(256);"
        f"print('RECEIVED from',addr,'data',data.decode(),flush=True);"
        f"s.sendto(b'ECHO:'+data,addr);"
        f"print('ECHOED to',addr,flush=True);"
        f"s.close()\""
    )
    proc = await asyncio.create_subprocess_exec(
        'ssh', '-o', 'StrictHostKeyChecking=no', '-o', 'ConnectTimeout=5',
        MESH_SERVER, echo_script,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    # Wait for ECHO_READY
    line = await asyncio.wait_for(proc.stdout.readline(), timeout=8)
    if b'ECHO_READY' not in line:
        proc.terminate()
        return None
    return proc


async def udp_hole_punch_test(local_port: int) -> tuple[bool, str|None, str|None]:
    """
    1. Start UDP listener locally
    2. Send probe to meshbay.org echo server (creates NAT mapping)
    3. Echo server responds to actual source addr/port it received
    4. Check if echo arrives locally
    Returns (success, seen_ext_addr, echo_content)
    """
    loop = asyncio.get_event_loop()
    result = {'addr': None, 'data': None}
    recv_event = asyncio.Event()

    class RxProtocol(asyncio.DatagramProtocol):
        def __init__(self, transport_holder):
            self._transport_holder = transport_holder
        def connection_made(self, transport):
            self._transport_holder.append(transport)
        def datagram_received(self, data, addr):
            result['addr'] = addr
            result['data'] = data
            recv_event.set()
        def error_received(self, exc):
            pass

    transport_holder = []
    transport, _ = await loop.create_datagram_endpoint(
        lambda: RxProtocol(transport_holder),
        local_addr=('0.0.0.0', local_port)
    )

    # Send probe — this is the hole punch
    transport.sendto(b'PING:MESHBAY:HOLEPUNCH', (MESHBAY_IP, MESHBAY_ECHO_PORT))

    try:
        await asyncio.wait_for(recv_event.wait(), timeout=8)
    except TimeoutError:
        pass
    finally:
        transport.close()

    if result['data']:
        return True, str(result['addr']), result['data'].decode()
    return False, None, None


# ── Main ──────────────────────────────────────────────────────────────────────

async def main():
    print("\n=== MeshBay Node POC — Spike 4: NAT Traversal ===\n")

    state = json.loads(STATE_FILE.read_text()) if STATE_FILE.exists() else {}

    endpoint_hint = None
    stun_ip1 = stun_ip2 = None
    stun_port1 = stun_port2 = None

    # ── Phase A: STUN discovery ────────────────────────────────────────────────
    print("[ A ] STUN external address discovery")
    stun_ip1, stun_port1 = stun_query(LOCAL_PORT, *STUN_SERVERS[0])
    if stun_ip1:
        print(f"  {PASS} via {STUN_SERVERS[0][0]}: {stun_ip1}:{stun_port1}")
    else:
        print(f"  {FAIL} {STUN_SERVERS[0][0]} unreachable")

    # ── Phase B: NAT type probe (cone vs symmetric) ────────────────────────────
    print("\n[ B ] NAT type detection")
    stun_ip2, stun_port2 = stun_query(LOCAL_PORT, *STUN_SERVERS[1])
    if stun_ip1 and stun_ip2:
        if stun_port1 == stun_port2:
            nat_type = "Cone NAT (same external port for both STUN servers)"
            print(f"  {PASS} {nat_type}")
            print(f"       Cloudflare STUN : {stun_ip1}:{stun_port1}")
            print(f"       Google STUN     : {stun_ip2}:{stun_port2}")
            endpoint_hint = f"{stun_ip1}:{stun_port1}"
        else:
            nat_type = "Symmetric NAT (different port per destination)"
            print(f"  {FAIL} {nat_type}")
            print(f"       Cloudflare STUN : {stun_ip1}:{stun_port1}")
            print(f"       Google STUN     : {stun_ip2}:{stun_port2}")
            print("       → Hole punching unreliable; TURN relay required")
    elif stun_ip1:
        nat_type = "Unknown (only one STUN server responded)"
        print(f"  {SKIP} {nat_type}")
        endpoint_hint = f"{stun_ip1}:{stun_port1}"
    else:
        nat_type = "Unknown (STUN unavailable)"
        print(f"  {FAIL} {nat_type}")

    # ── Phase C: Bidirectional UDP hole punching ───────────────────────────────
    print("\n[ C ] Bidirectional UDP hole punch (node → meshbay.org → echo back)")
    print(f"  Starting UDP echo server on meshbay.org:{MESHBAY_ECHO_PORT} ...")
    udp_ok = False
    seen_ext_addr = None
    try:
        echo_proc = await asyncio.wait_for(
            start_echo_server_on_meshbay(), timeout=10)
        if echo_proc:
            print(f"  Echo server ready. Sending probe from local:{LOCAL_PORT} ...")
            udp_ok, seen_ext_addr, echo_data = await udp_hole_punch_test(LOCAL_PORT)
            stdout, _ = await asyncio.wait_for(echo_proc.communicate(), timeout=5)
            server_log = stdout.decode().strip()

            if udp_ok:
                print(f"  {PASS} Echo received: '{echo_data}'")
                print(f"  {PASS} meshbay.org saw us as: {seen_ext_addr}")
                print(f"  {PASS} Server log: {server_log}")
                # seen_ext_addr is the actual external addr for meshbay.org dest
                # May differ from STUN if symmetric NAT
                actual_ext = seen_ext_addr.replace("('", "").replace("'", "").replace(", ", ":")
                if endpoint_hint and actual_ext != endpoint_hint:
                    print(f"  ⚠  STUN addr {endpoint_hint} ≠ actual {actual_ext} (symmetric NAT confirmed)")
                endpoint_hint = actual_ext
            else:
                print(f"  {FAIL} No echo received (timeout)")
                print(f"       Server log: {server_log}")
        else:
            print(f"  {FAIL} Could not start echo server on meshbay.org")
    except Exception as e:
        print(f"  {FAIL} Error: {e}")

    # ── Phase D: UPnP ─────────────────────────────────────────────────────────
    print("\n[ D ] UPnP port mapping")
    upnp_ip, upnp_port = try_upnp(LOCAL_PORT)
    if upnp_ip:
        print(f"  {PASS} Mapped {upnp_ip}:{upnp_port}")
        endpoint_hint = f"{upnp_ip}:{upnp_port}"
    else:
        print(f"  {FAIL} UPnP not available on this router")

    # ── Phase E: Update hub ────────────────────────────────────────────────────
    print("\n[ E ] Update endpoint_hint on hub")
    if endpoint_hint and state.get("username"):
        async with httpx.AsyncClient(timeout=10) as client:
            r = await client.post(f"{HUB_URL}/v1/users/login", json={
                "username": state["username"], "password": state["password"]})
            access_token = r.json()["access_token"]
            pk_b64 = state.get("sk_ed25519_b64", "")
            r = await client.post(f"{HUB_URL}/v1/nodes/announce",
                json={"pk_node": pk_b64, "endpoint_hint": endpoint_hint},
                headers={"Authorization": f"Bearer {access_token}"})
            if r.status_code == 201:
                state["node_id"]      = r.json()["node_id"]
                state["endpoint_hint"] = endpoint_hint
                print(f"  {PASS} endpoint_hint={endpoint_hint} on hub")
    else:
        print(f"  {SKIP} No endpoint to register")

    STATE_FILE.write_text(json.dumps(state, indent=2))

    print(f"\n{'='*55}")
    print("Spike 4 — NAT Traversal Summary")
    print(f"  External IP (STUN)  : {stun_ip1 or 'unknown'}")
    print(f"  NAT type            : {nat_type if stun_ip1 else 'unknown'}")
    print(f"  UDP bidirectional   : {PASS + ' works' if udp_ok else FAIL + ' blocked'}")
    print(f"  UPnP                : {PASS + ' works' if upnp_ip else FAIL + ' disabled'}")
    print(f"  endpoint_hint       : {endpoint_hint or 'none'}")
    if udp_ok:
        print("\n  P2P UDP is functional. QUIC transport will work.")
        print("  Spike 4 COMPLETE.")
    else:
        print("\n  P2P blocked — TURN relay needed for this configuration.")
        print("  Spike 4 COMPLETE (with finding: relay required).")

asyncio.run(main())