#!/usr/bin/env python3
"""Mediciones de la semana 11 sobre la topología de dumbbell/.

    python3 bench.py link --rate 50 --delay 40 --loss 0 --buffer 250
    python3 bench.py flow --cc reno --seconds 20 --csv reno.csv
    python3 bench.py bloat --cc cubic --seconds 15
    python3 bench.py pair --cc reno --cc2 cubic --seconds 30
    python3 bench.py pair --cc cubic --parallel2 4 --seconds 30
    python3 bench.py queue

Todo corre dentro de los nodos con docker exec. Las cifras dependen de cada computadora.
"""
import argparse
import json
import re
import statistics
import time

from netlab import RECEIVER, sh, spawn

BASE_PORT = 5201
DATA_PORT = 40001  # puerto local del stream de datos, para no confundirlo con el de control
CWND = re.compile(r"\bcwnd:(\d+)")
SSTHRESH = re.compile(r"\bssthresh:(\d+)")
RTT = re.compile(r"\brtt:([\d.]+)/")


def set_link(rate=50, delay=40, loss=0.0, buffer_kb=250, extra_s1=0, extra_s2=0,
             queue="fifo", ecn=False):
    """Configura el enlace. rate en Mbit/s, delay y extra_* en ms, loss en %, buffer en KiB."""
    # 1: el emisor pide ECN al conectarse. 2: el valor por defecto; lo acepta si se lo piden.
    for node in ("s1", "s2"):
        sh(node, f"sysctl -qw net.ipv4.tcp_ecn={1 if ecn else 2}")
    if queue == "red-ecn":
        # RED marca paquetes con ECN entre min y max, antes de que la cola se llene.
        sh("r1", "tc qdisc del dev eth3 root 2>/dev/null; "
                 f"tc qdisc add dev eth3 root handle 1: tbf rate {rate}mbit burst 15kb "
                 f"limit {buffer_kb}kb && "
                 f"tc qdisc add dev eth3 parent 1:1 handle 10: red limit {buffer_kb}kb "
                 f"min 30kb max 90kb avpkt 1500 burst 30 bandwidth {rate}mbit "
                 "probability 0.1 ecn")
    else:
        sh("r1", "tc qdisc del dev eth3 root 2>/dev/null; "
                 f"tc qdisc add dev eth3 root tbf rate {rate}mbit burst 15kb "
                 f"limit {buffer_kb}kb")
    netem = f"netem delay {delay}ms limit 20000"
    if loss > 0:
        netem += f" loss {loss}%"
    sh("r2", f"tc qdisc replace dev eth2 root {netem}")
    # El retardo extra va en el sentido de los ACK: hacia cada emisor, nunca en su propia
    # salida, donde TCP lo notaría antes de que el paquete salga de la máquina.
    for device, extra in (("eth1", extra_s1), ("eth2", extra_s2)):
        if extra > 0:
            sh("r1", f"tc qdisc replace dev {device} root netem delay {extra}ms limit 20000")
        else:
            sh("r1", f"tc qdisc del dev {device} root 2>/dev/null; true")


def start_servers():
    sh("rx", "pkill iperf3; true", check=False)
    time.sleep(0.2)
    for port in (BASE_PORT, BASE_PORT + 1):
        sh("rx", f"iperf3 -s -D -p {port}")
    time.sleep(0.5)


def stop_servers():
    sh("rx", "pkill iperf3; true", check=False)


def iperf_command(cc, seconds, port=BASE_PORT, cport=DATA_PORT, parallel=1):
    return (f"iperf3 -c {RECEIVER} -p {port} --cport {cport} -t {seconds} -C {cc} "
            f"-P {parallel} -i 1 -J")


def sampler_command(cport, seconds, interval=0.1):
    """Un ciclo dentro del nodo: es la única forma de muestrear cada 100 ms."""
    return (
        'read start rest </proc/uptime; limit=${start%.*}; limit=$((limit+' f"{seconds + 2}" '));'
        ' while :; do read now rest </proc/uptime; [ "${now%.*}" -ge "$limit" ] && break;'
        ' echo "@ $now"; ss -tinH "sport = :' f"{cport}" '";'
        f' sleep {interval}; done')


def parse_ss(text):
    """Convierte la salida del muestreo en filas (t, cwnd, ssthresh o None, rtt en ms)."""
    rows, now, first = [], None, None
    for line in text.splitlines():
        if line.startswith("@ "):
            now = float(line[2:])
            continue
        cwnd, ssthresh, rtt = CWND.search(line), SSTHRESH.search(line), RTT.search(line)
        if cwnd and rtt and now is not None:
            if first is None:
                first = now
            rows.append((round(now - first, 2), int(cwnd.group(1)),
                         int(ssthresh.group(1)) if ssthresh else None, float(rtt.group(1))))
    return rows


def write_csv(path, rows):
    with open(path, "w", encoding="utf-8") as file:
        file.write("t_s,cwnd_segments,ssthresh_segments,rtt_ms\n")
        for time_s, cwnd, ssthresh, rtt in rows:
            file.write(f"{time_s},{cwnd},{'' if ssthresh is None else ssthresh},{rtt}\n")


def summary(report):
    """Saca del JSON de iperf3 lo que nos importa."""
    end = report["end"]
    return {
        "cc": end.get("sender_tcp_congestion", "?"),
        "mbps": end["sum_received"]["bits_per_second"] / 1e6,
        "retrans": end["sum_sent"]["retransmits"],
    }


def run_flow(cc, seconds, node="s1", cport=DATA_PORT, sample=True):
    """Un flujo de iperf3 con su muestreo de ss. Devuelve (resumen, filas)."""
    sampler = spawn(node, sampler_command(cport, seconds)) if sample else None
    time.sleep(0.05)
    output = sh(node, iperf_command(cc, seconds, cport=cport), timeout=seconds + 40)
    rows = []
    if sampler:
        text, _ = sampler.communicate(timeout=30)
        rows = parse_ss(text)
    return summary(json.loads(output)), rows


def run_bloat(cc, seconds):
    """Mide el RTT de un ping antes y durante una transferencia que cruza la misma cola."""
    idle = sh("s2", f"ping -c 10 -i 0.2 {RECEIVER}", timeout=30)
    base = [float(x) for x in re.findall(r"time=([\d.]+)", idle)]
    pinger = spawn("s2", f"sleep 4; ping -c {(seconds - 5) * 5} -i 0.2 {RECEIVER}")
    result, rows = run_flow(cc, seconds)
    text, _ = pinger.communicate(timeout=30)
    during = [float(x) for x in re.findall(r"time=([\d.]+)", text)]
    return {
        "mbps": result["mbps"],
        "retrans": result["retrans"],
        "ping_idle_ms": statistics.median(base),
        "ping_load_ms": statistics.median(during) if during else float("nan"),
        "ping_max_ms": max(during) if during else float("nan"),
        "ping_sent": (seconds - 5) * 5,
        "ping_received": len(during),
        "rows": rows,
    }


def queue_stats():
    """Los contadores de la cola del cuello de botella."""
    return sh("r1", "tc -s qdisc show dev eth3")


def run_pair(cc1, cc2, seconds, parallel2=1):
    """Dos emisores a la vez, s1 y s2. Devuelve los dos resúmenes."""
    first = spawn("s1", iperf_command(cc1, seconds, BASE_PORT, DATA_PORT))
    second = spawn("s2", iperf_command(cc2, seconds, BASE_PORT + 1, DATA_PORT, parallel2))
    out1, _ = first.communicate(timeout=seconds + 40)
    out2, _ = second.communicate(timeout=seconds + 40)
    return summary(json.loads(out1)), summary(json.loads(out2))


def jain(values):
    """Índice de equidad de Jain: (suma x)^2 / (n * suma x^2)."""
    return sum(values) ** 2 / (len(values) * sum(x * x for x in values))


def main():
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("what", choices=["link", "flow", "bloat", "pair", "queue"])
    parser.add_argument("--rate", type=int, default=50, help="Mbit/s del cuello de botella")
    parser.add_argument("--delay", type=int, default=40, help="ms de retardo en un sentido")
    parser.add_argument("--loss", type=float, default=0.0, help="%% de pérdida aleatoria")
    parser.add_argument("--buffer", type=int, default=250, help="KiB de cola en el cuello")
    parser.add_argument("--extra-s1", type=int, default=0, help="ms extra de RTT para s1")
    parser.add_argument("--extra-s2", type=int, default=0, help="ms extra de RTT para s2")
    parser.add_argument("--queue", choices=["fifo", "red-ecn"], default="fifo")
    parser.add_argument("--ecn", action="store_true", help="los emisores piden ECN")
    parser.add_argument("--cc", default="cubic", help="algoritmo: reno o cubic")
    parser.add_argument("--cc2", default=None, help="algoritmo del segundo flujo en pair")
    parser.add_argument("--parallel2", type=int, default=1,
                        help="conexiones paralelas del segundo emisor en pair")
    parser.add_argument("--seconds", type=int, default=20)
    parser.add_argument("--csv", default=None, help="archivo para guardar cwnd en el tiempo")
    args = parser.parse_args()

    if args.what == "link":
        set_link(args.rate, args.delay, args.loss, args.buffer, args.extra_s1, args.extra_s2,
                 args.queue, args.ecn)
        print(sh("r1", "tc qdisc show dev eth3"), end="")
        print(sh("r2", "tc qdisc show dev eth2"), end="")
        return
    if args.what == "queue":
        print(queue_stats(), end="")
        return

    start_servers()
    try:
        if args.what == "flow":
            result, rows = run_flow(args.cc, args.seconds)
            print(f"{result['cc']}: {result['mbps']:.1f} Mbit/s, "
                  f"{result['retrans']} retransmisiones")
            if args.csv:
                write_csv(args.csv, rows)
                print(f"{len(rows)} muestras en {args.csv}")
        elif args.what == "bloat":
            found = run_bloat(args.cc, args.seconds)
            print(f"{args.cc}: {found['mbps']:.1f} Mbit/s, {found['retrans']} retransmisiones")
            print(f"ping sin carga: {found['ping_idle_ms']:.1f} ms; con carga: mediana "
                  f"{found['ping_load_ms']:.1f} ms, máximo {found['ping_max_ms']:.1f} ms; "
                  f"recibidos {found['ping_received']} de {found['ping_sent']}")
        else:
            one, two = run_pair(args.cc, args.cc2 or args.cc, args.seconds, args.parallel2)
            print(f"s1 {one['cc']}: {one['mbps']:.1f} Mbit/s, {one['retrans']} retransmisiones")
            print(f"s2 {two['cc']}: {two['mbps']:.1f} Mbit/s, {two['retrans']} retransmisiones")
            print(f"Jain: {jain([one['mbps'], two['mbps']]):.3f}")
    finally:
        stop_servers()


if __name__ == "__main__":
    main()