#!/usr/bin/env python3
"""Send a .dol to the running gc_bba_test over raw ethernet, for self-update.

Pairs with gc_bba_test/source/netboot.c.  Deliberately NOT IP: frames carry a
private ethertype and a 12-byte header, so there is no ARP, no IP stack and no
DHCP to go wrong -- it exercises exactly the BBA TX/RX paths the test suite
already proves on hardware.

Transfer is STOP-AND-WAIT: one block in flight, ACKed before the next.  That
is load-bearing, not laziness -- the gateware does not consume RRP, so the
SPRAM ring has no overflow protection, and one outstanding frame means it
cannot overrun.

Needs CAP_NET_RAW, so run with sudo.

    sudo ./scripts/send-dol.py enp3s0 gc_bba_test/gc_bba_test.dol
"""
import argparse
import os
import socket
import struct
import sys
import time

ETHERTYPE = 0x88B5
MAGIC = b"GCUP"
OP_DATA, OP_ACK, OP_DONE = 1, 2, 3
BLOCK = 1024
HDR = 26
# The GC's MAC comes from PAR0-5, which the gateware pre-populates with a
# Nintendo OUI.  Discovered from the first ACK, so this is only the initial
# destination for block 0.
BROADCAST = b"\xff" * 6


def frame(dst, src, op, seq, nblocks, payload=b""):
    f = bytearray()
    f += dst + src + struct.pack("!H", ETHERTYPE)
    f += MAGIC + bytes([op, 0]) + struct.pack("!HHH", seq, nblocks, len(payload))
    f += payload
    if len(f) < 60:                      # ethernet minimum
        f += b"\x00" * (60 - len(f))
    return bytes(f)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("iface", help="network interface facing the GameCube")
    ap.add_argument("dol", help=".dol to send")
    ap.add_argument("--timeout", type=float, default=0.5,
                    help="per-block ACK timeout in seconds (default 0.5)")
    ap.add_argument("--retries", type=int, default=20,
                    help="per-block retransmits before giving up (default 20)")
    args = ap.parse_args()

    img = open(args.dol, "rb").read()
    blocks = [img[i:i + BLOCK] for i in range(0, len(img), BLOCK)]
    if not blocks:
        sys.exit("empty image")
    if len(blocks) > 0xFFFF:
        sys.exit("image too large for a 16-bit block count")

    s = socket.socket(socket.AF_PACKET, socket.SOCK_RAW, socket.htons(ETHERTYPE))
    s.bind((args.iface, 0))
    src = s.getsockname()[4]
    s.settimeout(args.timeout)

    print(f"{args.dol}: {len(img)} bytes, {len(blocks)} blocks of {BLOCK}")
    print(f"sending from {src.hex(':')} on {args.iface}")

    dst = BROADCAST
    t0 = time.time()
    for seq, payload in enumerate(blocks):
        for attempt in range(args.retries):
            s.send(frame(dst, src, OP_DATA, seq, len(blocks), payload))
            try:
                while True:
                    pkt = s.recv(2048)
                    if len(pkt) < HDR or pkt[14:18] != MAGIC:
                        continue
                    if pkt[18] != OP_ACK:
                        continue
                    if struct.unpack("!H", pkt[20:22])[0] != seq:
                        continue        # stale ACK for an earlier block
                    if dst == BROADCAST:
                        dst = pkt[6:12]
                        print(f"GameCube is {dst.hex(':')}")
                    break
                break
            except socket.timeout:
                if attempt == 0:
                    print(f"  block {seq}: no ACK, retrying", file=sys.stderr)
        else:
            sys.exit(f"block {seq}: no ACK after {args.retries} tries")
        if seq % 32 == 0 or seq == len(blocks) - 1:
            pct = 100 * (seq + 1) // len(blocks)
            print(f"\r  {pct:3d}%  block {seq + 1}/{len(blocks)}", end="", flush=True)

    # DONE is not ACKed: the GC launches the new image immediately, so there is
    # nothing left running to answer.  Send it a few times in case one is lost.
    for _ in range(5):
        s.send(frame(dst, src, OP_DONE, 0, len(blocks)))
        time.sleep(0.02)

    dt = time.time() - t0
    print(f"\nsent {len(img)} bytes in {dt:.1f}s "
          f"({len(img) * 8 / dt / 1e6:.2f} Mbit/s); GC should relaunch now")


if __name__ == "__main__":
    if os.geteuid() != 0:
        print("warning: AF_PACKET needs root; try sudo", file=sys.stderr)
    main()