"""Send or receive one file over real UDP with any of the three protocols. Given.

    python3 transfer.py recv --proto gbn --port 5000 --out got.bin
    python3 transfer.py send --proto gbn --window 8 --timeout 0.1 10.0.0.2 5000 file.bin

The receiver prints nothing until the transfer ends; the sender prints one
line of JSON with its stats and the elapsed time. experiments.py runs both
inside Mininet for you, so you rarely need this by hand.

Reference: https://docs.python.org/3/library/socket.html
"""

import argparse
import json
import socket
import time

import gbn
import sr
import stopwait

PROTOCOLS = {"sw": stopwait, "gbn": gbn, "sr": sr}


class UdpSocket:
    """A UDP socket whose timeout applies to receiving only.

    The protocols set short timeouts while they wait for ACKs. On a plain
    socket that timeout also covers sendto(), which then fails when the
    outgoing queue is briefly full. Here sendto() waits instead, the same way
    channel.py's simulated link behaves.
    """

    def __init__(self, sock):
        self._sock = sock
        self._timeout = None

    def settimeout(self, seconds):
        self._timeout = seconds

    def sendto(self, data, addr):
        self._sock.settimeout(None)
        return self._sock.sendto(data, addr)

    def recvfrom(self, bufsize):
        self._sock.settimeout(self._timeout)
        return self._sock.recvfrom(bufsize)


def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="role", required=True)
    r = sub.add_parser("recv")
    r.add_argument("--proto", choices=PROTOCOLS, required=True)
    r.add_argument("--port", type=int, default=5000)
    r.add_argument("--window", type=int, default=8)
    r.add_argument("--out", required=True)
    s = sub.add_parser("send")
    s.add_argument("--proto", choices=PROTOCOLS, required=True)
    s.add_argument("--window", type=int, default=8)
    s.add_argument("--timeout", type=float, default=0.1)
    s.add_argument("host")
    s.add_argument("port", type=int)
    s.add_argument("file")
    args = ap.parse_args()

    raw = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock = UdpSocket(raw)
    mod = PROTOCOLS[args.proto]
    if args.role == "recv":
        raw.bind(("0.0.0.0", args.port))
        receiver_cls = getattr(mod, "Receiver", stopwait.Receiver)
        kw = {"window": args.window} if receiver_cls is not stopwait.Receiver else {}
        data = receiver_cls(sock, **kw).receive()
        with open(args.out, "wb") as f:
            f.write(data)
    else:
        with open(args.file, "rb") as f:
            data = f.read()
        start = time.monotonic()
        stats = mod.Sender(sock, (args.host, args.port), window=args.window, timeout=args.timeout).send(data)
        stats["seconds"] = round(time.monotonic() - start, 3)
        print(json.dumps(stats))


if __name__ == "__main__":
    main()