#!/usr/bin/env python3
"""Lo que probamos en la red leaf-spine, uno por uno.

    python3 scenarios.py <escenario>     corre un escenario y muestra lo que mide
    python3 scenarios.py reset           vuelve al estado inicial
    python3 scenarios.py                 lista los escenarios

Cada escenario imprime los comandos que corre en cada nodo. Léanlos: son lo mismo que harían
a mano. Un escenario no deshace el anterior; para eso está reset.
"""
import json
import os
import re
import statistics
import subprocess
import sys
import threading
import time

import netlab
from netlab import ANSWERS, HERE, Lab, wait_until

LAB = Lab("fabric", "w14-fabric", "redes-w14-fabric")

SPINES = ("s1", "s2")
LEAVES = ("l1", "l2", "l3")
HOSTS = ("h1", "h2", "h3", "h4", "h5", "h6")
TENANTS = ("ta1", "tb1", "ta2", "tb2")
H1, H3 = "10.214.1.10", "10.214.2.10"
S1_FROM_L1, S2_FROM_L1 = "10.214.11.0", "10.214.21.0"
# Por qué interfaz sale cada spine hacia cada leaf.
TOWARD = {"l1": "eth1", "l2": "eth2", "l3": "eth3"}
# Un flujo cuyo hueco pasa de esto se cortó: los datagramas van cada 20 ms.
CUT_MS = 150
RING = (("h1", "10.214.2.10"), ("h3", "10.214.1.20"), ("h2", "10.214.2.20"), ("h4", "10.214.1.10"))


# --- Medir --------------------------------------------------------------------------------

def spine_tx(leaf="l2"):
    """Bytes que cada spine lleva enviados hacia un leaf."""
    return {spine: int(LAB.out(spine, "cat", f"/sys/class/net/{TOWARD[leaf]}/statistics/tx_bytes"))
            for spine in SPINES}


def next_hop(sport, dport=5201, source=H1, dest=H3, protocol="tcp"):
    """El next hop que elige l1 para un paquete de esta 5-tupla que entra por su LAN."""
    text = LAB.out("l1", "ip", "route", "get", dest, "from", source, "iif", "br0",
                   "ipproto", protocol, "sport", str(sport), "dport", str(dport))
    found = re.search(r"via (\S+)", text)
    return found.group(1) if found else None


def set_policy(value, nodes=LEAVES):
    """0: direcciones IP. 1: 5-tupla o el hash que ya traiga el paquete. 3: los campos que pidamos."""
    for leaf in nodes:
        if value == 3:
            # 0x37: origen, destino, protocolo, puerto de origen y puerto de destino.
            LAB.sh(leaf, "sysctl", "-qw", "net.ipv4.fib_multipath_hash_fields=0x37")
        LAB.sh(leaf, "sysctl", "-qw", f"net.ipv4.fib_multipath_hash_policy={value}")


def start_iperf_servers(node, *ports):
    for port in ports:
        LAB.sh(node, "iperf3", "-s", "-D", "-p", str(port))
    time.sleep(0.5)


def stop_programs(nodes):
    for node in nodes:
        LAB.sh(node, "sh", "-c", "pkill iperf3; pkill flow_probe; pkill ring_reduce; pkill tcpdump; true")


def iperf_mbits(node, dest, *options, port=5201, seconds=4):
    """Corre iperf3 y devuelve (Mbit/s en el receptor, retransmisiones del emisor)."""
    done = LAB.sh(node, "iperf3", "-c", dest, "-p", str(port), "-t", str(seconds), "-J", *options,
                  timeout=seconds + 60)
    try:
        end = json.loads(done.stdout)["end"]
        return end["sum_received"]["bits_per_second"] / 1e6, end["sum_sent"].get("retransmits", 0)
    except (ValueError, KeyError):
        return 0.0, 0


def nstat(node, *names):
    """Contadores del kernel de un nodo, desde que arrancó su espacio de red."""
    text = LAB.out(node, "nstat", "-az", *names)
    return {found[0]: int(found[1]) for found in re.findall(r"^(\S+)\s+(\d+)", text, re.MULTILINE)}


def flows_on_failure(fail, flows=16, seconds=30, fail_after=4):
    """Corre flow_probe, aplica una falla y devuelve [(flujo, recibidos, perdidos, hueco_ms)]."""
    LAB.sh("h3", "sh", "-c", "pkill flow_probe; true")
    receiver = subprocess.Popen(
        ["docker", "exec", LAB.prefix + "h3", "/w14/flow_probe", "recv", "6000", str(seconds + 3)],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    time.sleep(0.5)
    sender = subprocess.Popen(
        ["docker", "exec", LAB.prefix + "h1", "/w14/flow_probe", "send", H3, "6000", str(flows),
         "20", str(seconds)],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    time.sleep(fail_after)
    fail()
    sender.communicate(timeout=seconds + 30)
    text, _ = receiver.communicate(timeout=seconds + 30)
    rows = re.findall(r"^\s*(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s*$", text, re.MULTILINE)
    # El receptor solo mide huecos entre dos datagramas. Si un flujo nunca vuelve, lo que falta
    # se ve al contar: se esperaba uno cada 20 ms.
    expected = seconds * 1000 // 20
    return [(int(r[0]), int(r[2]), expected - int(r[2]), max(int(r[4]), (expected - int(r[2])) * 20))
            for r in rows]


def summarize_failure(rows):
    hit = [row for row in rows if row[3] >= CUT_MS]
    print(f"  flujos: {len(rows)}; con corte: {len(hit)}; sin corte: {len(rows) - len(hit)}")
    if hit:
        print(f"  hueco en los flujos con corte: {min(r[3] for r in hit)} a "
              f"{max(r[3] for r in hit)} ms")
    calm = [r[3] for r in rows if r[3] < CUT_MS]
    if calm:
        print(f"  hueco máximo en los flujos sin corte: {max(calm)} ms")


# --- Hash y reparto ------------------------------------------------------------------------

def policy0():
    """Hash solo con las direcciones IP (el valor de fábrica) y 16 flujos de h1 a h3."""
    set_policy(0)
    measure_flows()


def policy1():
    """Hash con la 5-tupla y los mismos 16 flujos."""
    set_policy(1)
    measure_flows()


def measure_flows(flows=16, seconds=4):
    """Reparte flujos de h1 a h3 y devuelve qué fracción cruzó cada spine."""
    start_iperf_servers("h3", 5201)
    before = spine_tx()
    rate, _ = iperf_mbits("h1", H3, "-P", str(flows), seconds=seconds)
    after = spine_tx()
    delta = {spine: after[spine] - before[spine] for spine in SPINES}
    total = sum(delta.values()) or 1
    for spine in SPINES:
        print(f"  {spine} hacia l2: {delta[spine] / 1e6:10.1f} MB  {100 * delta[spine] / total:5.1f} %")
    print(f"  total: {rate:.0f} Mbit/s (depende de la computadora)")
    return {spine: delta[spine] / total for spine in SPINES}


def routeget():
    """Pregunta al kernel qué next hop usaría para 64 puertos de origen, con cada política."""
    results = {}
    for policy in (0, 1, 3):
        set_policy(policy)
        counts = {}
        for sport in range(40000, 40064):
            hop = next_hop(sport)
            counts[hop] = counts.get(hop, 0) + 1
        results[policy] = counts
        print(f"  política {policy}: {counts}")
    return results


def predict(flows=12):
    """¿Acierta ip route get con el camino real de un flujo TCP? Con la política 1 y con la 3."""
    start_iperf_servers("h3", 5201)
    results = {}
    # Cada corrida usa puertos nuevos: un puerto recién usado se queda en TIME_WAIT.
    first = 20000 + int(time.time()) % 20000
    for policy in (1, 3):
        set_policy(policy)
        hits = 0
        for offset in range(flows):
            sport = first + offset
            hop = next_hop(sport)
            before = spine_tx()
            iperf_mbits("h1", H3, "--cport", str(sport), seconds=1)
            after = spine_tx()
            used = S1_FROM_L1 if after["s1"] - before["s1"] > after["s2"] - before["s2"] else S2_FROM_L1
            hits += used == hop
        first += flows
        results[policy] = hits
        print(f"  política {policy}: ip route get acertó {hits} de {flows} flujos")
    return results


def shape(rate="100mbit"):
    """Limita a rate los dos uplinks de l1: así el spine se vuelve el cuello de botella."""
    for dev in ("eth1", "eth2"):
        LAB.sh("l1", "tc", "qdisc", "replace", "dev", dev, "root", "tbf", "rate", rate,
               "burst", "64kb", "latency", "20ms")


def port_on(hop, dport, first):
    """Busca un puerto de origen para que la 5-tupla caiga en ese next hop, según el kernel."""
    for sport in range(first, first + 200):
        if next_hop(sport, dport) == hop:
            return sport
    raise RuntimeError(f"ningún puerto de origen cae en {hop}")


def two_elephants(cports):
    """Dos iperf3 a la vez, cada uno con su puerto de origen. Devuelve los Mbit/s de cada uno."""
    start_iperf_servers("h3", 5201, 5202)
    jobs = [subprocess.Popen(
        ["docker", "exec", LAB.prefix + "h1", "iperf3", "-c", H3, "-p", str(port), "-t", "6", "-J",
         "--cport", str(cport)], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True)
        for port, cport in zip((5201, 5202), cports)]
    rates = []
    for job in jobs:
        out, _ = job.communicate(timeout=60)
        try:
            rates.append(json.loads(out)["end"]["sum_received"]["bits_per_second"] / 1e6)
        except (ValueError, KeyError):
            rates.append(0.0)
    return rates


def elephants():
    """Dos flujos grandes con uplinks de 100 Mbit/s: chocando en un spine y separados."""
    set_policy(3)
    shape()
    together = two_elephants((port_on(S1_FROM_L1, 5201, 41000), port_on(S1_FROM_L1, 5202, 41200)))
    apart = two_elephants((port_on(S1_FROM_L1, 5201, 42000), port_on(S2_FROM_L1, 5202, 42200)))
    print(f"  los dos por s1:     {together[0]:6.1f} + {together[1]:6.1f} = {sum(together):6.1f} Mbit/s")
    print(f"  uno por cada spine: {apart[0]:6.1f} + {apart[1]:6.1f} = {sum(apart):6.1f} Mbit/s")
    return together, apart


# --- Por qué no por paquete --------------------------------------------------------------

def steer_tables():
    """Tablas 101 y 102 de l1: una sola salida en cada una. Las usan perpacket y sdn_ctl.py."""
    LAB.sh("l1", "ip", "route", "replace", "10.214.2.0/24", "via", S1_FROM_L1, "table", "101")
    LAB.sh("l1", "ip", "route", "replace", "10.214.2.0/24", "via", S2_FROM_L1, "table", "102")
    for mark, table in ((1, 101), (2, 102)):
        LAB.sh("l1", "ip", "rule", "add", "fwmark", str(mark), "lookup", str(table))


def unsteer():
    LAB.sh("l1", "nft", "delete", "table", "ip", "w14steer")
    for mark, table in ((1, 101), (2, 102)):
        LAB.sh("l1", "ip", "rule", "del", "fwmark", str(mark), "lookup", str(table))
        LAB.sh("l1", "ip", "route", "flush", "table", str(table))


def two_paths(rate="100mbit", delay="2ms"):
    """Dos uplinks iguales para BGP, pero el de s2 con 2 ms más de latencia."""
    LAB.sh("l1", "tc", "qdisc", "replace", "dev", "eth1", "root", "netem", "rate", rate, "limit", "2000")
    LAB.sh("l1", "tc", "qdisc", "replace", "dev", "eth2", "root", "netem", "rate", rate, "delay", delay,
           "limit", "2000")


def nft_install(rules):
    done = LAB.sh("l1", "nft", rules)
    if done.returncode != 0:
        raise RuntimeError(done.stderr)


def perpacket_on():
    steer_tables()
    rules = ("table ip w14steer { chain pre { type filter hook prerouting priority mangle; "
             "ip saddr 10.214.1.10 ip daddr 10.214.2.10 meta mark set numgen inc mod 2 offset 1; }; }")
    nft_install(rules)


def reorder_run(seconds=5):
    """Un flujo TCP de h1 a h3. Devuelve Mbit/s, retransmisiones y segmentos fuera de orden."""
    names = ("TcpExtTCPOFOQueue",)
    before = nstat("h3", *names).get(names[0], 0)
    rate, retransmits = iperf_mbits("h1", H3, seconds=seconds)
    after = nstat("h3", *names).get(names[0], 0)
    print(f"  {rate:8.0f} Mbit/s  {retransmits:6d} retransmisiones  {after - before:8d} fuera de orden en h3")
    return rate, retransmits, after - before


def perpacket():
    """Un flujo TCP por dos caminos con 2 ms de diferencia: por flujo y por paquete."""
    set_policy(3)
    start_iperf_servers("h3", 5201)
    two_paths()
    print("  por flujo:")
    per_flow = reorder_run()
    perpacket_on()
    print("  por paquete:")
    per_packet = reorder_run()
    unsteer()
    return per_flow, per_packet


# --- Fallas --------------------------------------------------------------------------------

def spine_links(spine):
    """(nodo, interfaz) de los dos extremos de cada enlace de un spine."""
    number = spine[1]
    ends = [(spine, f"eth{i}") for i in (1, 2, 3)]
    ends += [(leaf, f"eth{number}") for leaf in LEAVES]
    return ends


def cut_silently(spine="s1"):
    for node, dev in spine_links(spine):
        LAB.sh(node, "tc", "qdisc", "replace", "dev", dev, "root", "netem", "loss", "100%")


def cut_with_carrier(spine="s1"):
    for dev in ("eth1", "eth2", "eth3"):
        LAB.sh(spine, "ip", "link", "set", dev, "down")


def set_timers(keepalive, hold):
    """Primero se configura en todos y después se reinicia: el hold se negocia al mínimo de los dos."""
    for node in (*SPINES, "l1", "l2"):
        LAB.sh(node, "vtysh", "-c", "configure terminal", "-c", f"router bgp {ASN[node]}",
               "-c", f"timers bgp {keepalive} {hold}", "-c", "end")
    for node in (*SPINES, "l1", "l2"):
        LAB.sh(node, "vtysh", "-c", "clear bgp *")


def bfd_config(enable):
    """Prende o apaga BFD en cada sesión entre los spines y l1, l2."""
    no = "" if enable else "no "
    for spine, leaf, address in NEIGHBORS:
        for node, peer in ((spine, address), (leaf, spine_side(address))):
            LAB.sh(node, "vtysh", "-c", "configure terminal", "-c", f"router bgp {ASN[node]}",
                   "-c", f"{no}neighbor {peer} bfd", "-c", "end")


def spine_side(address):
    return address[:-1] + "0"


ASN = {"s1": 65140, "s2": 65140, "l1": 65141, "l2": 65142, "l3": 65143}
NEIGHBORS = (("s1", "l1", "10.214.11.1"), ("s1", "l2", "10.214.12.1"),
             ("s2", "l1", "10.214.21.1"), ("s2", "l2", "10.214.22.1"))


def ecmp_ready():
    """¿l1 tiene dos next hop hacia la LAN de l2 y hay paso de h1 a h3?"""
    route = LAB.out("l1", "ip", "route", "show", "10.214.2.0/24")
    return route.count("nexthop") == 2 and "linkdown" not in route and ", 0% packet loss" in LAB.out(
        "h1", "ping", "-c", "1", "-W", "1", H3)


def carrier():
    """Falla con aviso: se apagan las interfaces de s1. Los flujos que iban por él se cortan."""
    set_policy(1)
    summarize = flows_on_failure(cut_with_carrier, seconds=12)
    summarize_failure(summarize)
    return summarize


def silent():
    """Falla silenciosa de s1 con temporizadores 3/9 (keepalive/hold)."""
    set_policy(1)
    rows = flows_on_failure(cut_silently, seconds=18)
    summarize_failure(rows)
    return rows


def silent_short():
    """Falla silenciosa de s1 con temporizadores 1/3."""
    set_policy(1)
    set_timers(1, 3)
    wait_until(ecmp_ready, seconds=40)
    rows = flows_on_failure(cut_silently, seconds=12)
    summarize_failure(rows)
    return rows


def watch_l1(seconds, cut_at, log):
    """Anota cuándo cambian, en l1, la entrada ARP de s1 y los next hop hacia l2."""
    start = time.monotonic()
    last = None
    while time.monotonic() - start < seconds:
        neighbor = LAB.out("l1", "ip", "neigh", "show", S1_FROM_L1).split()
        state = (neighbor[-1] if neighbor else "sin entrada",
                 LAB.out("l1", "ip", "route", "show", "10.214.2.0/24").count("via"))
        if state != last:
            log.append((time.monotonic() - start - cut_at, *state))
            last = state
        time.sleep(0.5)


def silent_default():
    """Falla silenciosa de s1 con los temporizadores de fábrica, 60/180. Tarda unos cuatro minutos."""
    set_policy(1)
    set_timers(60, 180)
    wait_until(ecmp_ready, seconds=40)
    log = []
    watcher = threading.Thread(target=watch_l1, args=(205, 4, log))
    watcher.start()
    rows = flows_on_failure(cut_silently, seconds=210)
    watcher.join()
    summarize_failure(rows)
    for moment, neighbor, hops in log:
        print(f"  a los {moment:6.1f} s del corte: ARP de s1 {neighbor}; next hop hacia l2 en la tabla: {hops}")
    return rows


def silent_bfd():
    """Falla silenciosa de s1 con BFD: 100 ms y multiplicador 3."""
    set_policy(1)
    for spine, leaf, address in NEIGHBORS:
        for node, peer in ((spine, address), (leaf, spine_side(address))):
            LAB.sh(node, "vtysh", "-c", "configure terminal", "-c", "bfd", "-c", f"peer {peer}",
                   "-c", "receive-interval 100", "-c", "transmit-interval 100", "-c", "end")
    bfd_config(True)
    wait_until(lambda: LAB.out("l1", "vtysh", "-c", "show bfd peers brief").count("up") >= 2, seconds=30)
    print(LAB.out("l1", "vtysh", "-c", "show bfd peers brief"))
    rows = flows_on_failure(cut_silently, seconds=12)
    summarize_failure(rows)
    return rows


def partial():
    """Se cae solo el enlace s1-l2: l1 deja de usar s1 para l2, y lo sigue usando para lo demás."""
    LAB.sh("s1", "ip", "link", "set", "eth2", "down")
    time.sleep(2)
    print(LAB.out("l1", "ip", "route", "show", "10.214.2.0/24"))
    print(LAB.out("s1", "vtysh", "-c", "show bgp ipv4 unicast 10.214.2.0/24"))


# --- VXLAN ---------------------------------------------------------------------------------

VTEP = {"l1": ("10.214.0.1", "10.214.0.2"), "l2": ("10.214.0.2", "10.214.0.1")}
SEGMENTS = ((100, "eth5"), (200, "eth6"))


def overlay(vxlan_mtu=1450):
    """Dos clientes con las mismas direcciones: VNI 100 y VNI 200, con VTEP en l1 y l2."""
    for leaf, (local, remote) in VTEP.items():
        for vni, port in SEGMENTS:
            LAB.sh(leaf, "ip", "link", "add", f"vx{vni}", "type", "vxlan", "id", str(vni),
                   "local", local, "dstport", "4789", "srcport", "49152", "65535")
            LAB.sh(leaf, "ip", "link", "set", f"vx{vni}", "mtu", str(vxlan_mtu))
            # Lista de VTEP a los que se copia lo que no se sabe a quién mandar.
            LAB.sh(leaf, "bridge", "fdb", "append", "00:00:00:00:00:00", "dev", f"vx{vni}", "dst", remote)
            LAB.sh(leaf, "ip", "link", "add", f"br{vni}", "type", "bridge")
            LAB.sh(leaf, "ip", "link", "set", f"vx{vni}", "master", f"br{vni}")
            LAB.sh(leaf, "ip", "link", "set", port, "master", f"br{vni}")
            LAB.sh(leaf, "ip", "link", "set", f"vx{vni}", "up")
            LAB.sh(leaf, "ip", "link", "set", f"br{vni}", "up")
    if netlab.ECHO:
        print(LAB.out("ta1", "ping", "-c", "2", "10.99.0.2"))
        print(LAB.out("l1", "bridge", "fdb", "show", "dev", "vx100"))


def overlay_down():
    for leaf in VTEP:
        for vni, port in SEGMENTS:
            LAB.sh(leaf, "ip", "link", "set", port, "nomaster")
            LAB.sh(leaf, "ip", "link", "del", f"vx{vni}")
            LAB.sh(leaf, "ip", "link", "del", f"br{vni}")
    for tenant in TENANTS:
        LAB.sh(tenant, "ip", "neigh", "flush", "all")
        LAB.sh(tenant, "ip", "link", "set", "eth1", "mtu", "1500")
    for node in (*SPINES, "l1", "l2"):
        for dev in ("eth1", "eth2"):
            LAB.sh(node, "ip", "link", "set", dev, "mtu", "1500")


def host_tunnel_down():
    for host in ("h1", "h3"):
        LAB.sh(host, "ip", "link", "del", "vx77")


def host_tunnel(srcport=("49152", "65535")):
    """VTEP en los hosts h1 y h3, para que l1 reparta el tráfico ya encapsulado."""
    for host, remote, inner in (("h1", H3, "10.77.0.1/24"), ("h3", H1, "10.77.0.2/24")):
        LAB.sh(host, "ip", "link", "add", "vx77", "type", "vxlan", "id", "77", "remote", remote,
               "dstport", "4789", "dev", "eth1", "srcport", *srcport)
        LAB.sh(host, "ip", "addr", "add", inner, "dev", "vx77")
        LAB.sh(host, "ip", "link", "set", "vx77", "up")


def tunnel_spread(srcport=("49152", "65535")):
    """16 flujos TCP dentro del túnel de hosts: ¿qué fracción cruza cada spine?"""
    set_policy(3)
    host_tunnel(srcport)
    start_iperf_servers("h3", 5201)
    wait_until(lambda: ", 0% packet loss" in LAB.out("h1", "ping", "-c", "1", "-W", "1", "10.77.0.2"), seconds=10)
    before = spine_tx()
    iperf_mbits("h1", "10.77.0.2", "-P", "16", seconds=3)
    after = spine_tx()
    host_tunnel_down()
    delta = {spine: after[spine] - before[spine] for spine in SPINES}
    total = sum(delta.values()) or 1
    for spine in SPINES:
        print(f"  {spine} hacia l2: {delta[spine] / 1e6:9.1f} MB  {100 * delta[spine] / total:5.1f} %")
    return {spine: delta[spine] / total for spine in SPINES}


def entropy():
    """El túnel entre hosts con todos los puertos de origen externos, y con uno solo."""
    print("  puerto de origen externo según el flujo interno (49152 a 65535):")
    wide = tunnel_spread()
    print("  puerto de origen externo fijo (50000):")
    narrow = tunnel_spread(("50000", "50001"))
    return wide, narrow


def capture_vxlan(filename, action, wait=1.0):
    """Captura VXLAN en los dos spines mientras corre action. Devuelve las rutas de los pcap."""
    os.makedirs(os.path.join(HERE, "captures"), exist_ok=True)
    paths = []
    for spine in SPINES:
        path = f"/captures/{filename}-{spine}.pcap"
        paths.append(os.path.join(HERE, "captures", f"{filename}-{spine}.pcap"))
        LAB.daemon(spine, f"tcpdump -i eth1 -n -s 200 -U -Z root -w {path} udp port 4789")
    time.sleep(1.5)
    action()
    time.sleep(wait)
    for spine in SPINES:
        LAB.sh(spine, "pkill", "tcpdump")
    time.sleep(0.5)
    return paths


def read_capture(paths, *options):
    out = ""
    for path in paths:
        done = subprocess.run(["tshark", "-r", path, *options], capture_output=True, text=True, timeout=60)
        out += done.stdout
    return out


def tenants():
    """Lo que ve cada cliente, y un ping de ta1 capturado en el underlay."""
    print(LAB.out("ta1", "ping", "-c", "1", "10.99.0.2"))
    print(LAB.out("tb1", "ping", "-c", "1", "10.99.0.2"))
    for node in ("ta1", "tb1", "ta2", "tb2"):
        print(f"  {node}: {LAB.out(node, 'cat', '/sys/class/net/eth1/address').strip()}  "
              f"{LAB.out(node, 'ip', 'neigh', 'show', 'dev', 'eth1').strip()}")
    paths = capture_vxlan("vxlan", lambda: LAB.sh("ta1", "ping", "-c", "3", "-i", "0.2", "10.99.0.2"))
    print(read_capture(paths, "-Y", "icmp", "-T", "fields", "-e", "ip.src", "-e", "ip.dst", "-e",
                       "udp.srcport", "-e", "udp.dstport", "-e", "vxlan.vni", "-e", "eth.src", "-e",
                       "eth.dst", "-e", "icmp.type"))


def inner_ports():
    """Ocho flujos TCP dentro de un cliente: ¿cuántos puertos de origen externos se ven?"""
    set_policy(1)
    LAB.sh("ta2", "iperf3", "-s", "-D", "-p", "5201")
    time.sleep(0.5)
    before = spine_tx("l2")
    paths = capture_vxlan("entropy", lambda: LAB.sh(
        "ta1", "iperf3", "-c", "10.99.0.2", "-P", "8", "-t", "2", "-M", "1300", timeout=60))
    after = spine_tx("l2")
    # El capturador corre en el eth1 de cada spine; lo que importa es qué spine llevó qué.
    text = read_capture(paths, "-T", "fields", "-e", "udp.srcport")
    ports = {line.strip() for line in text.splitlines() if line.strip()}
    print(f"  puertos de origen externos distintos: {len(ports)}")
    for spine in SPINES:
        print(f"  {spine} hacia l2: {(after[spine] - before[spine]) / 1e6:.1f} MB")
    LAB.sh("ta2", "pkill", "iperf3")
    return ports, {spine: after[spine] - before[spine] for spine in SPINES}


def tenant_ping(size):
    """Un ping con DF de ta1 a ta2. El paquete mide size + 28 bytes."""
    return LAB.sh("ta1", "/w14/df_ping", "10.99.0.2", str(size)).returncode == 0


def tenant_tcp():
    """Mbit/s de TCP entre ta1 y ta2; 0 si la conexión no logra pasar datos."""
    LAB.sh("ta2", "iperf3", "-s", "-D", "-p", "5201")
    time.sleep(0.5)
    rate, _ = iperf_mbits("ta1", "10.99.0.2", "--connect-timeout", "3000", seconds=3)
    LAB.sh("ta2", "pkill", "iperf3")
    return rate


def mtu_table():
    """Pings con DF desde un cliente, justo en el borde y en el tamaño completo, más un flujo TCP."""
    sizes = {size: tenant_ping(size) for size in (1422, 1423, 1472)}
    for size, passed in sizes.items():
        print(f"  df_ping {size} (paquete de {size + 28} bytes): {'pasa' if passed else 'no pasa'}")
    rate = tenant_tcp()
    print(f"  TCP de ta1 a ta2: {rate:.0f} Mbit/s")
    return sizes, rate


def tenant_mtu():
    """Una salida: decirle a cada cliente que su MTU es 1450."""
    for tenant in TENANTS:
        LAB.sh(tenant, "ip", "link", "set", "eth1", "mtu", "1450")


def jumbo_underlay(mtu=1550):
    """La otra salida, la de los data centers: darle más MTU al underlay."""
    for node in (*SPINES, "l1", "l2"):
        for dev in ("eth1", "eth2"):
            LAB.sh(node, "ip", "link", "set", dev, "mtu", str(mtu))
    for leaf in VTEP:
        for vni, _ in SEGMENTS:
            LAB.sh(leaf, "ip", "link", "set", f"vx{vni}", "mtu", "1500")


def warn():
    """VXLAN con MTU 1500 sobre un underlay de 1500: Linux avisa al cliente con ICMP."""
    overlay(vxlan_mtu=1500)
    LAB.sh("ta1", "ping", "-c", "1", "10.99.0.2")
    first = LAB.out("ta1", "/w14/df_ping", "10.99.0.2", "1472").strip()
    print(f"  df_ping 1472, primera vez: {first}")
    learned = " ".join(LAB.out("ta1", "ip", "route", "get", "10.99.0.2").split())
    print(f"  ta1 ahora cree: {learned}")
    print(f"  df_ping 1472, otra vez: {LAB.out('ta1', '/w14/df_ping', '10.99.0.2', '1472').strip()}")
    rate = tenant_tcp()
    print(f"  TCP de ta1 a ta2: {rate:.0f} Mbit/s")
    return learned, rate


# --- Trabajo colectivo -----------------------------------------------------------------

def ring_rounds(bytes_count=1000000, rounds=10):
    """Un all-reduce en anillo con h1, h3, h2 y h4. Devuelve los ms de cada ronda en cada nodo."""
    jobs = []
    for rank, (node, next_ip) in enumerate(RING):
        jobs.append(subprocess.Popen(
            ["docker", "exec", LAB.prefix + node, "/w14/ring_reduce", str(rank), str(len(RING)), "6100",
             next_ip, str(bytes_count), str(rounds)],
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True))
    results = {}
    for (node, _), job in zip(RING, jobs):
        out, _ = job.communicate(timeout=240)
        results[node] = [float(t) for t in re.findall(r"ronda \d+ ([\d.]+) ms", out)]
    return results


def collective():
    """Un all-reduce entre cuatro hosts, sin y con un host atrasado."""
    start = ring_rounds()
    LAB.sh("h4", "tc", "qdisc", "replace", "dev", "eth1", "root", "tbf", "rate", "100mbit",
           "burst", "32kb", "latency", "20ms")
    slow = ring_rounds()
    LAB.sh("h4", "tc", "qdisc", "del", "dev", "eth1", "root")
    for node in ("h1", "h3", "h2", "h4"):
        print(f"  {node}: mediana sin atraso {statistics.median(start[node]):7.1f} ms; "
              f"con h4 a 100 Mbit/s {statistics.median(slow[node]):7.1f} ms")
    return start, slow


# --- Un enlace con pérdida que no es congestión -------------------------------------------

def lossy_link(*netem):
    LAB.sh("l2", "tc", "qdisc", "replace", "dev", "eth3", "root", "netem", "rate", "100mbit",
           "delay", "20ms", "limit", "2000", *netem)


def wireless():
    """TCP sobre un último enlace de 100 Mbit/s y 40 ms de RTT, con y sin pérdida."""
    start_iperf_servers("h3", 5201)
    results = {}
    for name, netem in (("sin pérdida", ()), ("0,1 % de pérdida", ("loss", "0.1%")),
                        ("1 % de pérdida", ("loss", "1%")),
                        ("sin pérdida, con 4 ms de jitter", ("delay", "20ms", "4ms"))):
        LAB.sh("l2", "tc", "qdisc", "del", "dev", "eth3", "root")
        lossy_link(*netem)
        rate, retransmits = iperf_mbits("h1", H3, seconds=8)
        results[name] = rate
        print(f"  {name:32} {rate:7.1f} Mbit/s  {retransmits:6d} retransmisiones")
    LAB.sh("l2", "tc", "qdisc", "del", "dev", "eth3", "root")
    return results


# --- Laboratorio --------------------------------------------------------------------------

def load_l3(path=None):
    """Carga en l3 la configuración de BGP de lab/l3-bgp.conf."""
    path = path or os.path.join(ANSWERS, "l3-bgp.conf")
    with open(path, encoding="utf-8") as file:
        lines = [line.strip() for line in file if line.strip() and not line.strip().startswith("!")]
    args = ["vtysh", "-c", "configure terminal"]
    for line in lines:
        args += ["-c", line]
    print(LAB.out("l3", *args))


def load_vxlan(path=None):
    """Corre lab/vxlan.sh, que crea un túnel entre h2 y h6."""
    path = path or os.path.join(ANSWERS, "vxlan.sh")
    with open(path, encoding="utf-8") as file:
        script = file.read()
    done = subprocess.run(["sh", "-c", script], capture_output=True, text=True, timeout=120,
                          env={**os.environ, "P": LAB.prefix})
    print(done.stdout + done.stderr, end="")
    return done


# --- Volver al inicio ---------------------------------------------------------------------

def reset():
    """Vuelve al estado inicial."""
    stop_programs((*HOSTS, *SPINES, "l1", "l2", "l3", "ta1", "ta2"))
    for node in (*SPINES, *LEAVES):
        for dev in ("eth1", "eth2", "eth3"):
            LAB.sh(node, "tc", "qdisc", "del", "dev", dev, "root")
            if node in SPINES:
                LAB.sh(node, "ip", "link", "set", dev, "up")
    for host in HOSTS:
        LAB.sh(host, "tc", "qdisc", "del", "dev", "eth1", "root")
    LAB.sh("l2", "tc", "qdisc", "del", "dev", "eth3", "root")
    set_policy(0)
    unsteer()
    overlay_down()
    host_tunnel_down()
    bfd_config(False)
    for spine, leaf, address in NEIGHBORS:
        for node, peer in ((spine, address), (leaf, spine_side(address))):
            LAB.sh(node, "vtysh", "-c", "configure terminal", "-c", "bfd", "-c", f"no peer {peer}", "-c", "end")
    if "timers bgp 3 9" not in LAB.out("l1", "vtysh", "-c", "show running-config"):
        set_timers(3, 9)
    wait_until(ecmp_ready, seconds=40)


SCENARIOS = {
    "routeget": routeget,
    "predict": predict,
    "policy0": policy0,
    "policy1": policy1,
    "elephants": elephants,
    "perpacket": perpacket,
    "partial": partial,
    "carrier": carrier,
    "silent": silent,
    "silent_short": silent_short,
    "silent_default": silent_default,
    "silent_bfd": silent_bfd,
    "overlay": overlay,
    "tenants": tenants,
    "inner_ports": inner_ports,
    "mtu": mtu_table,
    "tenant_mtu": tenant_mtu,
    "jumbo": jumbo_underlay,
    "warn": warn,
    "entropy": entropy,
    "collective": collective,
    "wireless": wireless,
    "l3": load_l3,
    "vxlan": load_vxlan,
    "reset": reset,
}


def main():
    if len(sys.argv) != 2 or sys.argv[1] not in SCENARIOS:
        for name, function in SCENARIOS.items():
            print(f"{name:14} {(function.__doc__ or 'vuelve al estado inicial').strip().splitlines()[0]}")
        return 1
    netlab.ECHO = True
    SCENARIOS[sys.argv[1]]()
    return 0


if __name__ == "__main__":
    sys.exit(main())