#!/usr/bin/env python3
"""
Test de velocidad — Ciberkit
==================================

Mide descarga, subida, latencia, jitter y bufferbloat (latencia con la red
cargada) contra la red de Cloudflare, igual que la versión web, pero desde
la terminal: ideal para servidores, sedes remotas o para programarlo y tener
un historial con el que reclamar a tu proveedor.

Sin dependencias: Python 3.8+ en Windows, Linux o macOS.

Ejemplos
--------
  python test_velocidad.py                          # test completo
  python test_velocidad.py --ligero                 # ahorra datos (~80 MB)
  python test_velocidad.py --csv velocidad.csv      # añade una fila al historial
  python test_velocidad.py --json resultado.json
  python test_velocidad.py --minimo-bajada 250      # código 1 si baja de 250 Mbps (alertas)

Programarlo cada hora (Linux, crontab -e):
  0 * * * * /usr/bin/python3 /opt/test_velocidad.py --ligero --csv /var/log/velocidad.csv

Programarlo en Windows (cada hora):
  schtasks /Create /SC HOURLY /TN "TestVelocidad" /TR "python C:\\scripts\\test_velocidad.py --ligero --csv C:\\scripts\\velocidad.csv"

Ten en cuenta que cada test completo puede gastar cientos de MB.
"""

import argparse
import csv
import http.client
import json
import os
import re
import ssl
import statistics
import sys
import threading
import time

HOST = "speed.cloudflare.com"
UA = "suite-redes-test-velocidad/1.0"
PERFILES = {
    "completo": {"segundos": 10, "flujos": 4, "flujos_subida": 3, "max_bytes": float("inf"), "primer_bloque": 100_000},
    "ligero": {"segundos": 7, "flujos": 3, "flujos_subida": 2, "max_bytes": 40e6, "primer_bloque": 1_000_000},
}
NOTAS = [("A+", 5), ("A", 30), ("B", 60), ("C", 200), ("D", 400), ("F", float("inf"))]
CAMPOS_CSV = ["fecha", "descarga_mbps", "subida_mbps", "latencia_ms", "jitter_ms", "latencia_carga_ms", "bufferbloat", "servidor", "perfil", "fiabilidad", "rechazos"]


# ───────────────────────── Cálculos (mismos que la web: core/speed.js) ─────────────────────────

def percentil(xs, p):
    a = sorted(x for x in xs if x is not None)
    if not a:
        return None
    i = (len(a) - 1) * p / 100
    lo, hi = int(i), min(int(i) + 1, len(a) - 1)
    return a[lo] + (a[hi] - a[lo]) * (i - lo)


def jitter(xs):
    a = [x for x in xs if x is not None]
    if len(a) < 2:
        return None
    return sum(abs(a[i] - a[i - 1]) for i in range(1, len(a))) / (len(a) - 1)


def velocidad_final(muestras, arranque=0.2):
    xs = [x for x in muestras if x is not None and x >= 0]
    if not xs:
        return None
    estables = xs[int(len(xs) * arranque):]
    return percentil(estables or xs, 80)


def ms_servidor(cabecera):
    """Suma todas las métricas dur= de Server-Timing (Cloudflare las reparte en varias)."""
    return sum(float(x) for x in re.findall(r"dur=([\d.]+)", cabecera or ""))


def nota_bufferbloat(reposo, carga):
    if reposo is None or carga is None:
        return None
    aumento = max(carga - reposo, 0)
    return next(n for n, maximo in NOTAS if aumento < maximo)


# ───────────────────────── Medición ─────────────────────────

def fiabilidad(rechazos, muestras_bajada, muestras_subida, pings):
    """Igual que la web: una medición con peticiones rechazadas nunca se da por buena."""
    if rechazos > 0:
        return "baja"
    if min(muestras_bajada, muestras_subida) < 10 or pings < 15:
        return "media"
    return "alta"


def conexion():
    return http.client.HTTPSConnection(HOST, timeout=15, context=ssl.create_default_context())


def ping(conn):
    """Latencia de una petición vacía (hasta recibir las cabeceras), sin el tiempo del servidor."""
    t0 = time.perf_counter()
    conn.request("GET", f"/__down?bytes=0&r={time.time_ns()}", headers={"User-Agent": UA})
    r = conn.getresponse()
    t1 = time.perf_counter()
    r.read()
    if r.status != 200:
        raise OSError(f"HTTP {r.status}")
    return max((t1 - t0) * 1000 - ms_servidor(r.getheader("Server-Timing")), 0.1)


def info_servidor():
    """Servidor que responde. /meta solo responde a navegadores, así que el código de
    aeropuerto (IATA) se toma de la cabecera pública cf-ray («8f1c…-BOG»). No se usa la IP."""
    try:
        c = conexion()
        c.request("GET", f"/__down?bytes=0&r={time.time_ns()}", headers={"User-Agent": UA})
        r = c.getresponse()
        r.read()
        m = re.search(r"-([A-Z]{3})$", r.getheader("cf-ray") or "")
        return {"iata": m.group(1)} if m else {}
    except (OSError, http.client.HTTPException):
        return {}


class Pinger(threading.Thread):
    """Mide la latencia cada 400 ms mientras la red está cargada."""

    def __init__(self):
        super().__init__(daemon=True)
        self.valores = []
        self.parar = threading.Event()

    def run(self):
        conn = conexion()
        while not self.parar.is_set():
            try:
                self.valores.append(ping(conn))
            except (OSError, http.client.HTTPException):
                conn = conexion()
            self.parar.wait(0.4)


def caudal(direccion, perfil, al_muestrear=None):
    total = [0]
    rechazos = [0]
    lock = threading.Lock()
    t0 = time.perf_counter()
    limite = t0 + perfil["segundos"]
    fin = threading.Event()

    def terminado():
        return fin.is_set() or time.perf_counter() >= limite or total[0] >= perfil["max_bytes"]

    def sumar(n):
        with lock:
            total[0] += n

    cuerpo = os.urandom(1 << 20)

    def flujo():
        tam, techo, fallos = perfil["primer_bloque"], (100e6 if direccion == "down" else 32e6), 0
        conn = conexion()
        while not terminado():
            inicio = time.perf_counter()
            try:
                if direccion == "down":
                    conn.request("GET", f"/__down?bytes={int(tam)}&r={time.time_ns()}", headers={"User-Agent": UA})
                    r = conn.getresponse()
                    if r.status != 200:
                        raise OSError(f"HTTP {r.status}")
                    while True:
                        bloque = r.read(65536)
                        if not bloque:
                            break
                        sumar(len(bloque))
                        if terminado():
                            conn.close()
                            return
                else:
                    tam = int(tam)
                    conn.putrequest("POST", f"/__up?r={time.time_ns()}")
                    conn.putheader("User-Agent", UA)
                    conn.putheader("Content-Type", "application/octet-stream")
                    conn.putheader("Content-Length", str(tam))
                    conn.endheaders()
                    enviados = 0
                    while enviados < tam:
                        trozo = cuerpo[: min(65536, tam - enviados)]
                        conn.send(trozo)
                        enviados += len(trozo)
                        sumar(len(trozo))
                        if terminado():
                            conn.close()
                            return
                    r = conn.getresponse()
                    r.read()
                    if r.status != 200:
                        raise OSError(f"HTTP {r.status}")
            except (OSError, http.client.HTTPException):
                fallos += 1
                with lock:
                    rechazos[0] += 1
                if fallos > 3:
                    return
                techo = max(tam / 4, 100_000)
                tam = techo
                conn.close()
                conn = conexion()
                time.sleep(0.3 * fallos)
                continue
            if time.perf_counter() - inicio < 1.5:
                tam = min(tam * 4, techo)

    n = perfil["flujos"] if direccion == "down" else perfil["flujos_subida"]
    hilos = [threading.Thread(target=flujo, daemon=True) for _ in range(n)]
    for h in hilos:
        h.start()
    muestras, previo, t_prev = [], 0, t0
    while not terminado() and any(h.is_alive() for h in hilos):
        time.sleep(0.2)
        ahora = time.perf_counter()
        with lock:
            actual = total[0]
        mbps = (actual - previo) * 8 / ((ahora - t_prev) * 1e6)
        muestras.append(mbps)
        previo, t_prev = actual, ahora
        if al_muestrear:
            al_muestrear(mbps)
    fin.set()
    if total[0] == 0:
        raise OSError(f"no se pudo medir la {'descarga' if direccion == 'down' else 'subida'} (sin conexión o el servidor limitó las peticiones)")
    return velocidad_final(muestras), total[0], rechazos[0], len(muestras)


def medir(perfil, progreso=print):
    r = {}
    conn = conexion()
    try:
        ping(conn)  # calentar conexión y TLS
    except (OSError, http.client.HTTPException):
        conn = conexion()
    pings = []
    for _ in range(20):
        try:
            pings.append(ping(conn))
        except (OSError, http.client.HTTPException):
            conn = conexion()
    if len(pings) < 5:
        raise OSError("no se pudo medir la latencia (¿sin conexión?)")
    r["latencia_ms"], r["jitter_ms"] = statistics.median(pings), jitter(pings)
    progreso(f"  Latencia: {r['latencia_ms']:.1f} ms (jitter {r['jitter_ms']:.1f} ms)")

    for direccion, nombre in (("down", "descarga"), ("up", "subida")):
        p = Pinger()
        p.start()
        mbps, usados, rech, n_muestras = caudal(direccion, perfil)
        p.parar.set()
        p.join(2)
        r[f"{nombre}_mbps"] = mbps
        r[f"carga_{nombre}_ms"] = statistics.median(p.valores) if p.valores else None
        r[f"bytes_{nombre}"] = usados
        r[f"rechazos_{nombre}"] = rech
        r[f"muestras_{nombre}"] = n_muestras
        progreso(f"  {nombre.capitalize()}: {mbps:.1f} Mbps")

    cargas = [x for x in (r["carga_descarga_ms"], r["carga_subida_ms"]) if x is not None]
    r["latencia_carga_ms"] = max(cargas) if cargas else None
    r["bufferbloat"] = nota_bufferbloat(r["latencia_ms"], r["latencia_carga_ms"])
    r["rechazos"] = r["rechazos_descarga"] + r["rechazos_subida"]
    r["fiabilidad"] = fiabilidad(r["rechazos"], r["muestras_descarga"], r["muestras_subida"], len(pings))
    return r


# ───────────────────────── Salida ─────────────────────────

def fila_csv(r, servidor, perfil, fecha=None):
    f = lambda v, d=2: "" if v is None else f"{v:.{d}f}"  # noqa: E731
    return {
        "fecha": fecha or time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        "descarga_mbps": f(r.get("descarga_mbps")),
        "subida_mbps": f(r.get("subida_mbps")),
        "latencia_ms": f(r.get("latencia_ms"), 1),
        "jitter_ms": f(r.get("jitter_ms"), 1),
        "latencia_carga_ms": f(r.get("latencia_carga_ms"), 1),
        "bufferbloat": r.get("bufferbloat") or "",
        "servidor": servidor.get("iata", ""),
        "perfil": perfil,
        "fiabilidad": r.get("fiabilidad", ""),
        "rechazos": r.get("rechazos", ""),
    }


def guardar_csv(ruta, fila):
    nuevo = not os.path.exists(ruta) or os.path.getsize(ruta) == 0
    with open(ruta, "a", newline="", encoding="utf-8") as fh:
        w = csv.DictWriter(fh, fieldnames=CAMPOS_CSV)
        if nuevo:
            w.writeheader()
        w.writerow(fila)


def preparar_salida():
    for flujo in (sys.stdout, sys.stderr):
        try:
            if not flujo.isatty():
                flujo.reconfigure(encoding="utf-8", errors="replace")
            else:
                flujo.reconfigure(errors="replace")
        except (AttributeError, ValueError):
            pass


def main(argv=None):
    preparar_salida()
    ap = argparse.ArgumentParser(description="Test de velocidad (descarga, subida, latencia y bufferbloat) contra Cloudflare.")
    ap.add_argument("--ligero", action="store_true", help="ahorro de datos: ~80 MB como máximo")
    ap.add_argument("--csv", metavar="FICHERO", help="añade el resultado a un CSV (historial)")
    ap.add_argument("--json", metavar="FICHERO", help="guarda el resultado en JSON")
    ap.add_argument("--minimo-bajada", type=float, metavar="MBPS", help="código de salida 1 si la descarga queda por debajo")
    ap.add_argument("--silencioso", action="store_true", help="solo muestra el resumen final")
    args = ap.parse_args(argv)

    perfil = "ligero" if args.ligero else "completo"
    servidor = info_servidor()
    if not args.silencioso:
        donde = f"el servidor de Cloudflare {servidor['iata']}" if servidor.get("iata") else "Cloudflare"
        print(f"Midiendo contra {donde} — perfil {perfil}…")
    try:
        r = medir(PERFILES[perfil], progreso=(lambda *_: None) if args.silencioso else print)
    except OSError as e:
        print(f"Error: {e}", file=sys.stderr)
        return 2

    print(f"\nDescarga {r['descarga_mbps']:.1f} Mbps · Subida {r['subida_mbps']:.1f} Mbps · "
          f"Latencia {r['latencia_ms']:.0f} ms · Con carga {r['latencia_carga_ms'] or 0:.0f} ms (bufferbloat {r['bufferbloat'] or '?'})")
    print(f"Fiabilidad: {r['fiabilidad']}" + (f" — el servidor rechazó {r['rechazos']} peticiones: la velocidad real puede ser mayor. Repite más tarde." if r["rechazos"] else ""))
    if args.csv:
        guardar_csv(args.csv, fila_csv(r, servidor, perfil))
        print(f"Añadido a {args.csv}")
    if args.json:
        with open(args.json, "w", encoding="utf-8") as fh:
            json.dump({"formato": "suite-redes.velocidad/1", "servidor": servidor, "perfil": perfil, **r}, fh, ensure_ascii=False, indent=2)
        print(f"Guardado en {args.json}")
    if args.minimo_bajada and r["descarga_mbps"] < args.minimo_bajada:
        print(f"ALERTA: la descarga ({r['descarga_mbps']:.1f} Mbps) está por debajo de {args.minimo_bajada} Mbps", file=sys.stderr)
        return 1
    return 0


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