#!/usr/bin/env python3
"""Mediciones de la semana 1 sobre la topología topo/path.clab.yml.

    make up
    python3 experiments.py cable        # traceroute con un enlace largo
    python3 experiments.py size         # RTT según el tamaño del paquete
    python3 experiments.py bottleneck   # throughput según el enlace más lento
    python3 experiments.py queue        # cola y pérdida según la intensidad
    python3 experiments.py web          # qué cambia en un paquete al cruzar routers
    make down

Se corre dentro del contenedor del curso. Las cifras dependen de la máquina y
de lo que esté corriendo en ella. Las pruebas de test_path.py usan estas mismas
funciones, con márgenes amplios.
"""
import json
import re
import subprocess
import sys
import threading
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
PREFIX = "clab-w01-path-"
H2_ADDRESS = "10.2.0.2"

# Cada enlace, con los dos extremos que hay que configurar: tc controla lo que
# sale por una interfaz, así que un enlace necesita una regla en cada sentido.
LINKS = {
    "h1-r1": (("h1", "eth1"), ("r1", "eth1")),
    "r1-r2": (("r1", "eth2"), ("r2", "eth1")),
    "r2-r3": (("r2", "eth2"), ("r3", "eth1")),
    "r3-h2": (("r3", "eth2"), ("h2", "eth1")),
}
CABLE = "r2-r3"
CABLE_DELAY_MS = 21  # el retardo de propagación teórico de Lima a Miami


def sh(node, command, check=True):
    """Corre un comando en un nodo y devuelve su salida."""
    result = subprocess.run(
        ["docker", "exec", PREFIX + node, "sh", "-c", command],
        capture_output=True, text=True, check=False,
    )
    if check and result.returncode != 0:
        raise RuntimeError(f"{node}: {command}: {result.stderr.strip()}")
    return result.stdout


def start_detached(node, command):
    """Deja un comando corriendo en un nodo, en segundo plano."""
    subprocess.run(["docker", "exec", "-d", PREFIX + node, "sh", "-c", command], check=True)


def set_link(link, rate_mbit=None, delay_ms=None, limit=None):
    """Fija la velocidad, el retardo y el tamaño del buffer de un enlace."""
    options = ""
    if delay_ms is not None:
        options += f" delay {delay_ms}ms"
    if rate_mbit is not None:
        options += f" rate {rate_mbit}mbit"
    if limit is not None:
        options += f" limit {limit}"
    for node, device in LINKS[link]:
        sh(node, f"tc qdisc replace dev {device} root netem{options}")


def clear_links():
    """Quita todas las reglas: los enlaces vuelven a ser lo más rápidos posible."""
    for ends in LINKS.values():
        for node, device in ends:
            sh(node, f"tc qdisc del dev {device} root 2>/dev/null || true")


def wait_until_ready(timeout=60):
    """Espera a que FRRouting instale las rutas y h1 alcance a h2."""
    deadline = time.time() + timeout
    while time.time() < deadline:
        if sh("h1", f"ping -c 1 -W 1 {H2_ADDRESS} >/dev/null 2>&1 && echo ok", check=False).strip():
            return
        time.sleep(1)
    raise RuntimeError("h1 no alcanza a h2: revisen las rutas con ip route")


def traceroute():
    """Devuelve la salida de traceroute y, por salto, (dirección, RTT mínimo en ms)."""
    output = sh("h1", f"traceroute -n -w 2 {H2_ADDRESS}")
    hops = []
    for line in output.splitlines()[1:]:
        fields = line.split()
        times = [float(t) for t in re.findall(r"([\d.]+) ms", line)]
        hops.append((fields[1], min(times) if times else None))
    return output, hops


def min_rtt(size, count=40, interval=0.05):
    """RTT mínimo de varios ping: el mínimo descarta el ruido de la máquina."""
    output = sh("h1", f"ping -q -c {count} -i {interval} -s {size} {H2_ADDRESS}")
    return float(re.search(r"= ([\d.]+)/", output).group(1))


def fit_line(points):
    """Recta por mínimos cuadrados. Devuelve (pendiente, intercepto)."""
    n = len(points)
    mean_x = sum(x for x, _ in points) / n
    mean_y = sum(y for _, y in points) / n
    slope = sum((x - mean_x) * (y - mean_y) for x, y in points) / sum((x - mean_x) ** 2 for x, _ in points)
    return slope, mean_y - slope * mean_x


SIZES = [64, 400, 800, 1200, 1472]


def rtt_by_size(slow_links, rate_mbit=5, cable=False):
    """RTT mínimo para cada tamaño, con algunos enlaces lentos y quizá el cable."""
    clear_links()
    for link in slow_links:
        set_link(link, rate_mbit=rate_mbit)
    if cable:
        set_link(CABLE, delay_ms=CABLE_DELAY_MS, rate_mbit=rate_mbit if CABLE in slow_links else None)
    return {size: min_rtt(size) for size in SIZES}


def estimated_rate_mbit(rtts, links_crossed):
    """Velocidad que explica la pendiente del RTT: pendiente = 2 * enlaces * 8 / R."""
    slope_ms_per_byte, _ = fit_line(list(rtts.items()))
    return 2 * links_crossed * 8 / slope_ms_per_byte / 1000


def tcp_throughput_mbit(seconds=5):
    """Throughput de TCP de h1 a h2, medido por el receptor con iperf3."""
    output = sh("h1", f"iperf3 -c {H2_ADDRESS} -t {seconds} -J")
    return json.loads(output)["end"]["sum_received"]["bits_per_second"] / 1e6


def start_iperf_server():
    """Deja un servidor iperf3 en h2, si no hay uno."""
    sh("h2", "pgrep iperf3 >/dev/null || iperf3 -s -D")
    time.sleep(1)


FRAME_OVERHEAD = 8 + 20 + 14  # encabezados UDP, IP y Ethernet de cada datagrama
UDP_PAYLOAD = 1400
QUEUE_RATE_MBIT = 50
QUEUE_LIMIT = 200


def udp_load(rho, burst=None, seconds=7):
    """Envía UDP con intensidad rho por el enlace r1-r2 y mide el RTT mientras tanto.

    La intensidad se calcula con los bytes que de verdad cruzan el enlace,
    encabezados incluidos. Devuelve el RTT (promedio y máximo) y la pérdida de UDP.
    """
    payload_mbit = rho * QUEUE_RATE_MBIT * UDP_PAYLOAD / (UDP_PAYLOAD + FRAME_OVERHEAD)
    target = f"{payload_mbit:.2f}M" + (f"/{burst}" if burst else "")
    result = {}

    def send():
        output = sh("h1", f"iperf3 -c {H2_ADDRESS} -u -b {target} -l {UDP_PAYLOAD} -t {seconds} -J")
        result["lost"] = json.loads(output)["end"]["sum"]["lost_percent"]

    sender = threading.Thread(target=send)
    sender.start()
    time.sleep(2)  # que la cola llegue a su estado normal antes de medir
    output = sh("h1", f"ping -q -c 20 -i 0.2 -W 2 {H2_ADDRESS}", check=False)
    sender.join()
    summary = re.search(r"= ([\d.]+)/([\d.]+)/([\d.]+)", output)
    result["avg"] = float(summary.group(2)) if summary else None
    result["max"] = float(summary.group(3)) if summary else None
    return result


def capture_web():
    """Captura una petición web junto a h1 y junto a h2. Devuelve las rutas de las capturas."""
    clear_links()
    directory = HERE / "captures"
    directory.mkdir(exist_ok=True)
    # Los corchetes evitan que pgrep -f se encuentre a sí mismo.
    if not sh("h2", "pgrep -f '[h]ttp.server'", check=False).strip():
        start_detached("h2", "python3 -m http.server 80 --directory /srv/web")
    time.sleep(1)
    points = {"near-h1": ("r1", "eth1"), "near-h2": ("r3", "eth2")}
    for node, device in points.values():
        sh(node, "rm -f /tmp/web.pcap")
        start_detached(node, f"tcpdump -U -i {device} -w /tmp/web.pcap tcp port 80")
    time.sleep(2)
    sh("h1", f"curl -sS -o /dev/null http://{H2_ADDRESS}/")
    time.sleep(1)
    paths = {}
    for name, (node, _) in points.items():
        sh(node, "pkill -INT tcpdump")
        time.sleep(1)
        paths[name] = directory / f"{name}.pcap"
        subprocess.run(["docker", "cp", "-q", f"{PREFIX}{node}:/tmp/web.pcap", str(paths[name])], check=True)
    return paths


def packet_fields(path):
    """Campos de cada paquete de una captura, leídos con tshark."""
    names = ["eth.src", "eth.dst", "ip.ttl", "ip.src", "ip.dst", "ip.checksum", "tcp.seq_raw", "tcp.len"]
    command = ["tshark", "-r", str(path), "-T", "fields"]
    for name in names:
        command += ["-e", name]
    output = subprocess.run(command, capture_output=True, text=True, check=True).stdout
    return [dict(zip(names, line.split("\t"))) for line in output.splitlines()]


def show_cable():
    clear_links()
    set_link(CABLE, delay_ms=CABLE_DELAY_MS)
    print(f"Enlace {CABLE} con {CABLE_DELAY_MS} ms por sentido\n")
    print(traceroute()[0])


def show_size():
    runs = [
        ("un enlace a 5 Mbps", rtt_by_size(["r1-r2"]), 1),
        ("tres enlaces a 5 Mbps", rtt_by_size(["r1-r2", "h1-r1", "r3-h2"]), 3),
        ("un enlace a 5 Mbps y cable de 21 ms", rtt_by_size(["r1-r2"], cable=True), 1),
    ]
    print("RTT mínimo en ms, según los bytes de datos de ping\n")
    print(f"{'caso':<37}" + "".join(f"{size:>8}" for size in SIZES) + f"{'R estimada':>13}")
    for name, rtts, links in runs:
        row = "".join(f"{rtts[size]:8.2f}" for size in SIZES)
        print(f"{name:<37}{row}{estimated_rate_mbit(rtts, links):9.1f} Mbps")


def show_bottleneck():
    start_iperf_server()
    print("TCP de h1 a h2, 5 s, throughput en el receptor\n")
    print(f"{'r1-r2':>6}{'r2-r3':>7}{'r3-h2':>7}{'medido':>13}")
    for rates in [(50, 200, 200), (200, 50, 200), (200, 200, 50), (200, 200, 200), (20, 200, 200)]:
        clear_links()
        for link, rate in zip(["r1-r2", "r2-r3", "r3-h2"], rates):
            set_link(link, rate_mbit=rate)
        print(f"{rates[0]:>6}{rates[1]:>7}{rates[2]:>7}{tcp_throughput_mbit():>8.1f} Mbps")


def show_queue():
    start_iperf_server()
    clear_links()
    set_link("r1-r2", rate_mbit=QUEUE_RATE_MBIT, limit=QUEUE_LIMIT)
    print(f"Enlace r1-r2 a {QUEUE_RATE_MBIT} Mbps, buffer de {QUEUE_LIMIT} paquetes, UDP de h1 a h2\n")
    print(f"{'carga':<22}{'rho':>5}{'RTT prom.':>12}{'RTT max.':>11}{'pérdida UDP':>14}")
    for name, rho, burst in [("sin carga", 0, None), ("regular", 0.5, None), ("regular", 0.8, None),
                             ("regular", 0.95, None), ("regular", 1.2, None),
                             ("ráfagas de 100", 0.5, 100)]:
        if rho == 0:
            output = sh("h1", f"ping -q -c 10 -i 0.2 {H2_ADDRESS}")
            values = re.search(r"= ([\d.]+)/([\d.]+)/([\d.]+)", output)
            print(f"{name:<22}{rho:>5.2f}{float(values.group(2)):>9.2f} ms{float(values.group(3)):>8.2f} ms{'-':>14}")
            continue
        r = udp_load(rho, burst)
        print(f"{name:<22}{rho:>5.2f}{r['avg']:>9.2f} ms{r['max']:>8.2f} ms{r['lost']:>13.1f}%")


def show_web():
    paths = capture_web()
    near_h1, near_h2 = packet_fields(paths["near-h1"]), packet_fields(paths["near-h2"])
    print("El primer paquete de h1 a h2, visto a cada lado de los tres routers\n")
    print(f"{'campo':<14}{'junto a h1':<20}{'junto a h2':<20}")
    for field in near_h1[0]:
        changed = "  cambia" if near_h1[0][field] != near_h2[0][field] else ""
        print(f"{field:<14}{near_h1[0][field]:<20}{near_h2[0][field]:<20}{changed}")
    print(f"\nCapturas en {paths['near-h1'].parent.name}/; ábranlas en Wireshark o con ./layers")


COMMANDS = {
    "cable": show_cable,
    "size": show_size,
    "bottleneck": show_bottleneck,
    "queue": show_queue,
    "web": show_web,
}


def main():
    names = sys.argv[1:] or list(COMMANDS)
    unknown = [name for name in names if name not in COMMANDS]
    if unknown:
        sys.exit(f"uso: {sys.argv[0]} [{'|'.join(COMMANDS)}]")
    wait_until_ready()
    for name in names:
        COMMANDS[name]()
        print()
    clear_links()


if __name__ == "__main__":
    main()