#!/usr/bin/env python3
"""pat-perf — Prisma Access Toolbox command-line tunnel test.

Runs TCP and UDP tests with iperf3 against the toolbox's iperf3 server and
attaches the result to the newest VPN Protocol Performance Testing run with the same label
(e.g. MASQUE or IPsec), so the browser report compares everything together.

Requires Python 3.8+ (standard library only) and iperf3 in PATH
(macOS: brew install iperf3 · Debian/Ubuntu: apt install iperf3 · Windows:
download an iperf3 build and pass --iperf3 C:\\path\\iperf3.exe).

  python3 pat-perf.py --label MASQUE
  python3 pat-perf.py --label IPsec --duration 20 --udp-rate 10M
  python3 pat-perf.py --compare masque.json ipsec.json
"""
import argparse
import json
import platform
import shutil
import socket
import statistics
import subprocess
import sys
import time
import urllib.request

# Filled in by the toolbox when the script is downloaded from the tool page.
DEFAULT_SERVER = "masque-test.demo-lab.at"
DEFAULT_TOOLBOX = "https://masque-test.demo-lab.at/tunnel-perf"


def unset(v):
    return not v or v.startswith("__TP_")


# --- scoring (same simplified ITU-T G.107 E-model as lib/score.js) ---------------
def mos(rtt_ms, jitter_ms, loss_pct):
    if None in (rtt_ms, jitter_ms, loss_pct):
        return None
    eff = rtt_ms / 2 + 2 * jitter_ms + 10
    r = 93.2 - eff / 40 if eff < 160 else 93.2 - (eff - 120) / 10
    r = max(0.0, min(100.0, r - 2.5 * loss_pct))
    m = 1 + 0.035 * r + 0.000007 * r * (r - 60) * (100 - r)
    return round(max(1.0, min(4.5, m)), 2)


# --- iperf3 ----------------------------------------------------------------------
def run_iperf(args, extra, label):
    cmd = [args.iperf3, "-c", args.server, "-p", str(args.port), "-J", "-t", str(args.duration)] + extra
    for attempt in range(1, 6):
        print(f"  {label}...", end="", flush=True)
        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=args.duration + 60)
        try:
            data = json.loads(proc.stdout)
        except ValueError:
            data = {"error": (proc.stderr or proc.stdout or "no output").strip()}
        err = data.get("error")
        if err and "busy" in err.lower() and attempt < 5:
            print(" server busy, retrying in 3 s")
            time.sleep(3)
            continue
        print(" failed: " + err if err else " done")
        return data
    return data


def tcp_result(data):
    if data.get("error"):
        return {"error": data["error"]}
    end = data.get("end", {})
    rtts = [s.get("sender", {}).get("mean_rtt") for s in end.get("streams", [])]
    rtts = [r / 1000 for r in rtts if r]  # µs -> ms, only when the OS reports it
    return {
        "mbps": end.get("sum_received", {}).get("bits_per_second", 0) / 1e6,
        "retransmits": end.get("sum_sent", {}).get("retransmits"),
        "rttMs": statistics.mean(rtts) if rtts else None,
    }


def udp_result(data):
    if data.get("error"):
        return {"error": data["error"]}
    end = data.get("end", {})
    s = end.get("sum_received") or end.get("sum") or {}
    return {"mbps": s.get("bits_per_second", 0) / 1e6, "jitterMs": s.get("jitter_ms"), "lossPct": s.get("lost_percent")}


def connect_rtt(server, port):
    """Median TCP connect time to the iperf3 port: an idle-path RTT estimate."""
    samples = []
    for _ in range(5):
        t = time.perf_counter()
        try:
            with socket.create_connection((server, port), timeout=3):
                samples.append((time.perf_counter() - t) * 1000)
        except OSError:
            pass
        time.sleep(0.2)
    return statistics.median(samples) if samples else None


# --- toolbox upload --------------------------------------------------------------
def http_json(url, body=None):
    req = urllib.request.Request(url, data=json.dumps(body).encode() if body is not None else None,
                                 headers={"Content-Type": "application/json"} if body is not None else {})
    with urllib.request.urlopen(req, timeout=20) as r:
        return json.loads(r.read().decode())


def egress():
    try:
        p = http_json("https://ipinfo.io/json")
        return {k: p.get(k) for k in ("ip", "city", "region", "country", "org")}
    except Exception:
        return {}


def upload(args, result, eg):
    base = args.toolbox.rstrip("/")
    r = http_json(f"{base}/api/cli/upload", {"label": args.label, "newRun": args.new_run, "egress": eg,
                                             "userAgent": f"pat-perf ({platform.platform()})", "iperf": result})
    if r.get("attached"):
        return f"attached to run \"{r['label']}\" from {r['createdAt'][:16].replace('T', ' ')} UTC"
    return f"saved as new run \"{args.label}\""


# --- output ----------------------------------------------------------------------
def fmt(v, unit="", digits=1):
    return "n/a" if v is None else f"{v:.{digits}f}{(' ' + unit) if unit else ''}"


ROWS = [
    ("TCP upload", "tcpUp", "mbps", "Mbit/s", False),
    ("TCP download", "tcpDown", "mbps", "Mbit/s", False),
    ("TCP upload retransmits", "tcpUp", "retransmits", "", True),
    ("TCP download retransmits", "tcpDown", "retransmits", "", True),
    ("Idle RTT (TCP connect)", None, "connectRttMs", "ms", True),
    ("UDP voice jitter", "udpVoice", "jitterMs", "ms", True),
    ("UDP voice loss", "udpVoice", "lossPct", "%", True),
    ("UDP voice MOS", "udpVoice", "mos", "", False),
    ("UDP rate throughput", "udpRate", "mbps", "Mbit/s", False),
    ("UDP rate jitter", "udpRate", "jitterMs", "ms", True),
    ("UDP rate loss", "udpRate", "lossPct", "%", True),
]


def digits(key, unit):
    return 0 if key == "retransmits" else 2 if unit in ("%", "") else 1


def value(result, section, key):
    src = result if section is None else (result.get(section) or {})
    v = src.get(key)
    return v if isinstance(v, (int, float)) else None


def print_summary(result):
    print()
    for name, section, key, unit, _ in ROWS:
        print(f"  {name:<28} {fmt(value(result, section, key), unit, digits(key, unit)):>14}")
    for section in ("tcpUp", "tcpDown", "udpVoice", "udpRate"):
        if (result.get(section) or {}).get("error"):
            print(f"  {section}: {result[section]['error']}")


def compare(path_a, path_b):
    a, b = (json.load(open(p)) for p in (path_a, path_b))
    la, lb = a.get("label", "A"), b.get("label", "B")
    print(f"\n  {'Metric':<28} {la:>14} {lb:>14} {'Δ':>8}  Better")
    for name, section, key, unit, lower in ROWS:
        va, vb = value(a, section, key), value(b, section, key)
        delta = better = ""
        if va is not None and vb is not None:
            if va:
                delta = f"{(vb - va) / abs(va) * 100:+.0f} %"
            # Same rule as the report (lib/score.js): < 5 % or a tiny absolute gap is a tie.
            tie = abs(vb - va) <= max(abs(va), abs(vb)) * 0.05 or abs(vb - va) < {"%": 0.1, "": 0.05}.get(unit, 1)
            better = "equivalent" if tie else (lb if (vb < va) == lower else la)
        print(f"  {name:<28} {fmt(va, unit, digits(key, unit)):>14} {fmt(vb, unit, digits(key, unit)):>14} {delta:>8}  {better}")


def main():
    ap = argparse.ArgumentParser(description="Prisma Access Toolbox tunnel test (iperf3 TCP/UDP).")
    ap.add_argument("--label", help="run label, e.g. MASQUE or IPsec")
    ap.add_argument("--server", default=None if unset(DEFAULT_SERVER) else DEFAULT_SERVER, help="iperf3 server host")
    ap.add_argument("--port", type=int, default=5201)
    ap.add_argument("--duration", type=int, default=10, help="seconds per test (default 10)")
    ap.add_argument("--streams", type=int, default=4, help="parallel TCP streams (default 4)")
    ap.add_argument("--udp-rate", default="5M", help="UDP rate test bitrate, iperf3 syntax (default 5M)")
    ap.add_argument("--toolbox", default=None if unset(DEFAULT_TOOLBOX) else DEFAULT_TOOLBOX, help="VPN Protocol Performance Testing URL, e.g. https://host/tunnel-perf")
    ap.add_argument("--new-run", action="store_true", help="always create a new run instead of attaching")
    ap.add_argument("--no-upload", action="store_true", help="don't send the result to the toolbox")
    ap.add_argument("--json-out", help="also write the result to this file (input for --compare)")
    ap.add_argument("--iperf3", default=shutil.which("iperf3") or "iperf3", help="path to the iperf3 binary")
    ap.add_argument("--compare", nargs=2, metavar=("A.json", "B.json"), help="compare two --json-out files and exit")
    args = ap.parse_args()

    if args.compare:
        return compare(*args.compare)
    if not args.label:
        ap.error("--label is required (e.g. MASQUE or IPsec)")
    if not args.server:
        ap.error("--server is required")
    if not args.no_upload and not args.toolbox:
        ap.error("--toolbox is required (or use --no-upload)")
    try:
        version = subprocess.run([args.iperf3, "--version"], capture_output=True, text=True).stdout.split("\n")[0]
    except OSError:
        sys.exit("iperf3 not found. Install it (brew install iperf3 / apt install iperf3) or pass --iperf3 PATH.")

    print(f"pat-perf: {args.label} -> {args.server}:{args.port} ({version})")
    eg = egress()
    if eg.get("ip"):
        print(f"  egress {eg['ip']} ({eg.get('city') or '?'}, {eg.get('country') or '?'}, {eg.get('org') or '?'})")
    rtt = connect_rtt(args.server, args.port)
    d = args.duration
    result = {
        "label": args.label,
        "server": args.server, "port": args.port, "durationS": d, "streams": args.streams,
        "client": {"host": socket.gethostname(), "os": platform.platform(), "iperf3": version},
        "egress": eg,
        "connectRttMs": rtt,
        "tcpUp": tcp_result(run_iperf(args, ["-P", str(args.streams)], f"TCP upload ({args.streams} streams, {d} s)")),
        "tcpDown": tcp_result(run_iperf(args, ["-P", str(args.streams), "-R"], f"TCP download ({args.streams} streams, {d} s)")),
        "udpVoice": udp_result(run_iperf(args, ["-u", "-b", "64K", "-l", "160"], f"UDP voice (64 kbit/s, {d} s)")),
        "udpRate": udp_result(run_iperf(args, ["-u", "-b", args.udp_rate, "-l", "1200"], f"UDP rate ({args.udp_rate}bit/s, {d} s)")),
        "recordedAt": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    }
    result["udpRate"]["target"] = args.udp_rate
    v = result["udpVoice"]
    if not v.get("error"):
        v["mos"] = mos(rtt, v.get("jitterMs"), v.get("lossPct"))
    print_summary(result)

    if args.json_out:
        with open(args.json_out, "w") as f:
            json.dump(result, f, indent=2)
        print(f"\n  written to {args.json_out}")
    if not args.no_upload:
        try:
            print(f"\n  toolbox: {upload(args, result, eg)}")
        except Exception as e:
            print(f"\n  toolbox upload failed: {e}")
            return 1
    return 0


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