#!/usr/bin/env python
"""relay_stub_test.py -- Phase-1 verification for btconsole.py --relay mode.

Spawns the relay with a 3-pilot test egg, then drives stub clients through the
whole protocol: console egg delivery, game HELLO/registration, PEER_UP/DOWN,
unicast + broadcast routing (route rewrite), UDP HELLO/ACK, UDP forward,
UDP->TCP fallback, and the reject paths (bad magic, dup id, out-of-roster,
oversize).  Exits 0 with 'ALL PASS' or 1 with the first failure.
"""
import os
import socket
import struct
import subprocess
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
TOOLS = os.path.join(HERE, "..", "tools")
CONSOLE_PORT = 27500
GAME_PORT = CONSOLE_PORT + 1

ENV_TCP = "<iI"
ENV_UDP = "<iiI"
HELLO_MAGIC = 0x31525442
R_BCAST, R_HELLO, R_PEER_UP, R_PEER_DOWN, R_UDP_ACK = -1, -2, -3, -4, -5

fails = []


def check(cond, what):
    tag = "ok" if cond else "FAIL"
    print(f"  [{tag}] {what}")
    if not cond:
        fails.append(what)


def recv_exact(sock, n, timeout=5.0):
    sock.settimeout(timeout)
    buf = b""
    while len(buf) < n:
        d = sock.recv(n - len(buf))
        if not d:
            raise ConnectionError("peer closed")
        buf += d
    return buf


def recv_env(sock, timeout=5.0):
    route, length = struct.unpack(ENV_TCP, recv_exact(sock, 8, timeout))
    payload = recv_exact(sock, length, timeout) if length else b""
    return route, payload


def send_env(sock, route, payload):
    sock.sendall(struct.pack(ENV_TCP, route, len(payload)) + payload)


def hello(sock, host_id):
    send_env(sock, R_HELLO, struct.pack("<iI", HELLO_MAGIC, host_id))


def game_connect():
    s = socket.create_connection(("127.0.0.1", GAME_PORT), timeout=5)
    s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    return s


def main():
    egg = os.path.join(HERE, "relay_test.egg")
    with open(egg, "w", newline="\n") as f:
        f.write("[pilots]\npilot=10.99.0.1:1502\npilot=10.99.0.2:1502\n"
                "pilot=10.99.0.3:1502\n[10.99.0.1:1502]\nhostType=0\n"
                "[10.99.0.2:1502]\nhostType=0\n[10.99.0.3:1502]\nhostType=0\n")

    relay = subprocess.Popen(
        [sys.executable, os.path.join(TOOLS, "btconsole.py"),
         "--relay", str(CONSOLE_PORT), egg, "--bind", "127.0.0.1"],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    time.sleep(1.0)
    try:
        run_tests()
    finally:
        relay.terminate()
        try:
            out, _ = relay.communicate(timeout=5)
            print("---- relay log ----")
            print(out)
        except subprocess.TimeoutExpired:
            relay.kill()
    if fails:
        print(f"\n{len(fails)} FAILURE(S):")
        for f_ in fails:
            print(f"  - {f_}")
        return 1
    print("\nALL PASS")
    return 0


def run_tests():
    print("== console side: egg delivery ==")
    con = socket.create_connection(("127.0.0.1", CONSOLE_PORT), timeout=5)
    first = recv_exact(con, 1040)
    clid, gid, fh, ts = struct.unpack_from("<iiii", first, 0)
    mlen, mid, mfl, seq, total, this = struct.unpack_from("<IiIiii", first, 16)
    check(clid == 0 and fh == 1 and mid == 3 and seq == 0,
          f"egg chunk header (clid={clid} fh={fh} mid={mid} seq={seq})")
    got = this
    while got < total:
        pkt = recv_exact(con, 1040)
        got += struct.unpack_from("<IiIiii", pkt, 16)[5]
    check(got == total, f"egg fully delivered ({got}/{total} bytes)")

    print("== game side: HELLO/registration/PEER_UP ==")
    a = game_connect()
    hello(a, 2)
    b = game_connect()
    hello(b, 3)
    route, payload = recv_env(a)
    check(route == R_PEER_UP and struct.unpack("<i", payload)[0] == 3,
          "A got PEER_UP(3) when B joined")
    route, payload = recv_env(b)
    check(route == R_PEER_UP and struct.unpack("<i", payload)[0] == 2,
          "B (newcomer) got PEER_UP(2)")

    print("== routing: unicast + route rewrite ==")
    frame = b"\x00" * 16 + struct.pack("<IiI", 12, 99, 1)   # minimal fake frame
    send_env(a, 3, frame)                                    # A -> unicast to 3
    route, payload = recv_env(b)
    check(route == 2 and payload == frame,
          f"B received A's unicast with route rewritten to sender (route={route})")

    print("== routing: broadcast fan-out ==")
    c = game_connect()
    hello(c, 4)
    for s_, who in ((a, "A"), (b, "B")):
        route, payload = recv_env(s_)
        check(route == R_PEER_UP and struct.unpack("<i", payload)[0] == 4,
              f"{who} got PEER_UP(4)")
    for _ in range(2):                                       # C gets 2 PEER_UPs
        recv_env(c)
    send_env(a, R_BCAST, frame)                              # A broadcasts ONCE
    for s_, who in ((b, "B"), (c, "C")):
        route, payload = recv_env(s_)
        check(route == 2 and payload == frame, f"{who} received A's broadcast")
    a.settimeout(0.5)
    try:
        a.recv(1)
        check(False, "A must NOT receive its own broadcast")
    except socket.timeout:
        check(True, "A did not receive its own broadcast")

    print("== UDP: HELLO/ACK + forward ==")
    ua = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    ua.bind(("127.0.0.1", 0))
    ua.settimeout(5)
    ua.sendto(struct.pack(ENV_UDP, R_HELLO, 2, 0), ("127.0.0.1", GAME_PORT))
    ack, _ = ua.recvfrom(4096)
    check(struct.unpack_from(ENV_UDP, ack)[0] == R_UDP_ACK, "UDP HELLO acked")
    ub = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    ub.bind(("127.0.0.1", 0))
    ub.settimeout(5)
    ub.sendto(struct.pack(ENV_UDP, R_HELLO, 3, 0), ("127.0.0.1", GAME_PORT))
    ub.recvfrom(4096)
    dg = struct.pack(ENV_UDP, 3, 2, 7) + frame               # A -> unicast 3
    ua.sendto(dg, ("127.0.0.1", GAME_PORT))
    fwd, _ = ub.recvfrom(4096)
    check(fwd == dg, "UDP unicast forwarded verbatim")
    dg2 = struct.pack(ENV_UDP, R_BCAST, 2, 8) + frame        # A -> broadcast
    ua.sendto(dg2, ("127.0.0.1", GAME_PORT))
    fwd2, _ = ub.recvfrom(4096)
    check(fwd2 == dg2, "UDP broadcast reached B")
    # C never sent UDP HELLO -> its copy must fall back to TCP.
    route, payload = recv_env(c)
    check(route == 2 and payload == frame,
          "C (no UDP endpoint) got the broadcast via TCP fallback")

    print("== reject paths ==")
    bad = game_connect()
    send_env(bad, R_HELLO, struct.pack("<iI", 0xDEAD, 5))
    check(bad.recv(1) == b"", "bad HELLO magic -> connection dropped")
    dup = game_connect()
    hello(dup, 2)
    check(dup.recv(1) == b"", "duplicate hostID -> dropped")
    oor = game_connect()
    hello(oor, 99)
    check(oor.recv(1) == b"", "out-of-roster hostID -> dropped")
    big = game_connect()
    big.sendall(struct.pack(ENV_TCP, 3, 5000))
    check(big.recv(1) == b"", "oversize frame -> dropped")

    print("== PEER_DOWN ==")
    c.close()
    for s_, who in ((a, "A"), (b, "B")):
        s_.settimeout(5)
        route, payload = recv_env(s_)
        check(route == R_PEER_DOWN and struct.unpack("<i", payload)[0] == 4,
              f"{who} got PEER_DOWN(4)")


if __name__ == "__main__":
    sys.exit(main())