"""Measure the three protocols over a real lossy link. Given; do not change it.

Run it inside the course container (see README.md), as root:

    python3 experiments.py              # every experiment, about four minutes
    python3 experiments.py loss         # goodput against loss rate only
    python3 experiments.py window       # goodput against window size only

It builds the smallest possible network: two network namespaces, h1 and h2,
joined by one virtual cable (a veth pair), with netem on each end adding the
delay and the loss. That is exactly what Mininet builds for two hosts and a
link. h1 sends a file to h2 with transfer.py, and the run counts only if the
file arrives byte for byte. Loss is random, so every setting runs three times
and the tables show the median run. Results print as tables for FINDINGS.md
and go to results.json with every run kept.

netem's loss applies in both directions: at 5%, data packets and ACKs are
each lost 5% of the time.

References:
    https://man7.org/linux/man-pages/man8/ip-netns.8.html
    https://man7.org/linux/man-pages/man8/tc-netem.8.html
"""

import hashlib
import json
import os
import subprocess
import sys
import time

FILE_BYTES = 100_000
DELAY = "10ms"            # each way, so the round trip is about 20 ms
TIMEOUT = 0.1
TRIALS = 3
HERE = os.path.dirname(os.path.abspath(__file__))
ADDR = {"h1": "10.0.0.1", "h2": "10.0.0.2"}


def sh(*cmd):
    subprocess.run(cmd, check=True, capture_output=True, text=True)


def build(loss):
    """Two namespaces joined by a veth pair, with netem delay and loss on both ends."""
    teardown()
    for host in ADDR:
        sh("ip", "netns", "add", host)
    sh("ip", "link", "add", "h1-eth0", "netns", "h1", "type", "veth", "peer", "name", "h2-eth0", "netns", "h2")
    for host, addr in ADDR.items():
        dev = f"{host}-eth0"
        sh("ip", "-n", host, "addr", "add", f"{addr}/24", "dev", dev)
        sh("ip", "-n", host, "link", "set", dev, "up")
        sh("ip", "-n", host, "link", "set", "lo", "up")
        sh("ip", "netns", "exec", host, "tc", "qdisc", "add", "dev", dev, "root",
           "netem", "delay", DELAY, "loss", f"{loss}%")


def teardown():
    for host in ADDR:
        subprocess.run(["ip", "netns", "del", host], capture_output=True)


def run(proto, loss, window):
    build(loss)
    src, dst = "/tmp/p3_send.bin", "/tmp/p3_recv.bin"
    with open(src, "wb") as f:
        f.write(os.urandom(FILE_BYTES))
    if os.path.exists(dst):
        os.remove(dst)
    receiver = subprocess.Popen(["ip", "netns", "exec", "h2", sys.executable, "transfer.py", "recv",
                                 "--proto", proto, "--window", str(window), "--out", dst], cwd=HERE)
    time.sleep(0.3)                                   # let the receiver bind its port first
    try:
        sent = subprocess.run(["ip", "netns", "exec", "h1", sys.executable, "transfer.py", "send",
                               "--proto", proto, "--window", str(window), "--timeout", str(TIMEOUT),
                               ADDR["h2"], "5000", src], cwd=HERE, capture_output=True, text=True, timeout=300)
        out = sent.stdout + sent.stderr
    except subprocess.TimeoutExpired:
        out = "the sender ran for 300 s without finishing"
    try:
        receiver.wait(5)                              # 5 s to finish after the sender is done
    except subprocess.TimeoutExpired:
        receiver.kill()
        receiver.wait()
    teardown()
    try:
        stats = json.loads(out.strip().splitlines()[-1])
    except (ValueError, IndexError):
        return {"proto": proto, "loss": loss, "window": window, "ok": False, "error": out.strip()[-300:]}
    same = os.path.exists(dst) and _digest(src) == _digest(dst)
    stats.update(proto=proto, loss=loss, window=window, ok=same,
                 goodput_kbps=round(FILE_BYTES * 8 / 1000 / stats["seconds"], 1))
    return stats


def median_run(proto, loss, window):
    runs = [run(proto, loss, window) for _ in range(TRIALS)]
    good = sorted((r for r in runs if r["ok"]), key=lambda r: r["goodput_kbps"])
    best = dict(good[len(good) // 2]) if len(good) == TRIALS else dict(runs[0], ok=False)
    best["runs"] = runs
    return best


def _digest(path):
    with open(path, "rb") as f:
        return hashlib.sha256(f.read()).hexdigest()


def table(rows, key):
    print(f"\n| protocol | {key} | goodput (kbit/s) | transmissions | retransmissions | intact |")
    print("| :--- | --: | --: | --: | --: | :--- |")
    for r in rows:
        if not r["ok"]:
            print(f"| {r['proto']} | {r[key]} | failed | | | no |")
            continue
        print(f"| {r['proto']} | {r[key]} | {r['goodput_kbps']} | {r['transmissions']} "
              f"| {r['retransmissions']} | yes |")


def main():
    if os.geteuid() != 0:
        sys.exit("run this as root inside the course container: creating namespaces needs it")
    wanted = sys.argv[1:] or ["loss", "window"]
    results = {}
    if "loss" in wanted:
        rows = [median_run(p, loss, 8) for p in ("sw", "gbn", "sr") for loss in (0, 1, 5, 10, 20)]
        table(rows, "loss")
        results["loss"] = rows
    if "window" in wanted:
        rows = [median_run(p, 5, w) for p in ("gbn", "sr") for w in (1, 2, 4, 8, 16, 32)]
        table(rows, "window")
        results["window"] = rows
    with open(os.path.join(HERE, "results.json"), "w") as f:
        json.dump(results, f, indent=1)
    print("\nwrote results.json")


if __name__ == "__main__":
    main()