#!/usr/bin/env python3
"""Create thesis-ready plots from network_timing_benchmark.py results.

The script reads one or more benchmark result directories containing
`summary.json`, extracts the most relevant metrics, writes CSV files for
checking, and exports publication-friendly figures.
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path
from typing import Any


def import_plotting_stack() -> tuple[Any, Any, Any]:
    try:
        import matplotlib.pyplot as plt
        import pandas as pd
        import seaborn as sns
    except ModuleNotFoundError as exc:
        missing = exc.name or "plotting dependency"
        print(
            f"Missing Python package: {missing}\n\n"
            "Install the plotting dependencies with:\n"
            "  python3 -m pip install pandas matplotlib seaborn\n\n"
            "On Ubuntu with apt packages, this usually also works:\n"
            "  sudo apt install python3-pandas python3-matplotlib python3-seaborn",
            file=sys.stderr,
        )
        raise SystemExit(2) from exc
    return pd, plt, sns


def safe_name(value: str) -> str:
    return re.sub(r"[^A-Za-z0-9_.-]+", "_", value).strip("_") or "plot"


def find_summary_files(input_paths: list[Path]) -> list[Path]:
    files: list[Path] = []
    for path in input_paths:
        if path.is_file() and path.name == "summary.json":
            files.append(path)
        elif path.is_dir():
            files.extend(path.rglob("summary.json"))
    return sorted(set(files))


def infer_setup_name(summary_path: Path, payload: dict[str, Any]) -> str:
    label = str(payload.get("label") or "").strip()
    if label:
        return label

    folder = summary_path.parent.name
    match = re.match(r"^\d{8}-\d{6}-(?P<label>.+)$", folder)
    if match:
        return match.group("label")
    return folder


def setup_sort_key(setup: str) -> tuple[int, str]:
    normalized = setup.lower()
    order = {
        "direct": 0,
        "baseline": 0,
        "bridge": 1,
        "bridge-only": 1,
        "bridge-new": 1,
        "af_packet": 2,
        "af-packet": 2,
        "tc_ebpf": 3,
        "tc-ebpf": 3,
        "ebpf": 3,
        "nfqueue": 4,
        "nftables": 5,
    }
    for key, rank in order.items():
        if key in normalized:
            return rank, normalized
    return 99, normalized


def add_metric(
    rows: list[dict[str, Any]],
    *,
    setup: str,
    source: Path,
    category: str,
    metric: str,
    value: Any,
    unit: str,
    direction: str | None = None,
    payload_size_bytes: int | None = None,
    protocol: str | None = None,
    test: str | None = None,
) -> None:
    if value is None:
        return
    try:
        number = float(value)
    except (TypeError, ValueError):
        return
    rows.append(
        {
            "setup": setup,
            "category": category,
            "metric": metric,
            "value": number,
            "unit": unit,
            "direction": direction,
            "payload_size_bytes": payload_size_bytes,
            "protocol": protocol,
            "test": test,
            "source": str(source),
        }
    )


def extract_summary(summary_path: Path) -> list[dict[str, Any]]:
    payload = json.loads(summary_path.read_text(encoding="utf-8"))
    setup = infer_setup_name(summary_path, payload)
    rows: list[dict[str, Any]] = []

    for result in payload.get("tests", []):
        if result.get("skipped"):
            continue
        result_type = result.get("type")

        if result_type == "ping":
            size = result.get("payload_size_bytes")
            try:
                size = int(size) if size is not None else None
            except (TypeError, ValueError):
                size = None
            rtt = result.get("rtt_ms", {}) or {}
            jitter = result.get("jitter", {}) or {}
            for metric in ("mean", "median", "p95", "p99", "min", "max", "stdev", "mad", "iqr"):
                add_metric(
                    rows,
                    setup=setup,
                    source=summary_path,
                    category="ping",
                    metric=f"rtt_{metric}",
                    value=rtt.get(metric),
                    unit="ms",
                    payload_size_bytes=size,
                )
            add_metric(
                rows,
                setup=setup,
                source=summary_path,
                category="ping",
                metric="jitter_mean_abs_delta",
                value=jitter.get("mean_abs_delta_ms"),
                unit="ms",
                payload_size_bytes=size,
            )
            add_metric(
                rows,
                setup=setup,
                source=summary_path,
                category="ping",
                metric="loss",
                value=result.get("loss_percent"),
                unit="percent",
                payload_size_bytes=size,
            )

        elif result_type in {"iperf_tcp", "iperf_udp"}:
            category = result_type
            direction = result.get("direction", "forward")
            add_metric(
                rows,
                setup=setup,
                source=summary_path,
                category=category,
                metric="throughput",
                value=result.get("mbit_per_second"),
                unit="Mbit/s",
                direction=direction,
            )
            if result_type == "iperf_tcp":
                add_metric(
                    rows,
                    setup=setup,
                    source=summary_path,
                    category=category,
                    metric="retransmits",
                    value=result.get("retransmits"),
                    unit="count",
                    direction=direction,
                )
            else:
                add_metric(
                    rows,
                    setup=setup,
                    source=summary_path,
                    category=category,
                    metric="jitter",
                    value=result.get("jitter_ms"),
                    unit="ms",
                    direction=direction,
                )
                add_metric(
                    rows,
                    setup=setup,
                    source=summary_path,
                    category=category,
                    metric="loss",
                    value=result.get("lost_percent"),
                    unit="percent",
                    direction=direction,
                )

        elif result_type == "sockperf":
            protocol = result.get("protocol")
            extracted = result.get("extracted", {}) or {}
            for metric, unit in (
                ("avg_latency_usec", "us"),
                ("min_latency_usec", "us"),
                ("max_latency_usec", "us"),
            ):
                add_metric(
                    rows,
                    setup=setup,
                    source=summary_path,
                    category="sockperf",
                    metric=metric,
                    value=extracted.get(metric),
                    unit=unit,
                    protocol=protocol,
                )

        elif result_type == "arping":
            rtt = result.get("rtt_ms", {}) or {}
            for metric in ("mean", "median", "p95", "p99"):
                add_metric(
                    rows,
                    setup=setup,
                    source=summary_path,
                    category="arping",
                    metric=f"rtt_{metric}",
                    value=rtt.get(metric),
                    unit="ms",
                )

        elif result_type == "flent":
            add_metric(
                rows,
                setup=setup,
                source=summary_path,
                category="flent",
                metric="data_file_exists",
                value=1 if result.get("data_file_exists") else 0,
                unit="bool",
                test=result.get("test"),
            )

    return rows


def save_figure(fig: Any, output_dir: Path, name: str, formats: list[str], dpi: int) -> None:
    for fmt in formats:
        path = output_dir / f"{name}.{fmt}"
        fig.savefig(path, dpi=dpi, bbox_inches="tight")
        print(f"wrote {path}")


def maybe_warn_empty(df: Any, name: str) -> bool:
    if df.empty:
        print(f"skipping {name}: no data")
        return True
    return False


def plot_ping_percentiles(pd: Any, plt: Any, sns: Any, metrics: Any, output_dir: Path, formats: list[str], dpi: int) -> None:
    df = metrics[
        (metrics["category"] == "ping")
        & (metrics["metric"].isin(["rtt_median", "rtt_p95", "rtt_p99"]))
    ].copy()
    if maybe_warn_empty(df, "ping percentiles"):
        return
    df["payload"] = df["payload_size_bytes"].astype("Int64").astype(str) + " B"
    df["percentile"] = df["metric"].map(
        {
            "rtt_median": "median",
            "rtt_p95": "p95",
            "rtt_p99": "p99",
        }
    )

    grid = sns.catplot(
        data=df,
        kind="bar",
        x="setup",
        y="value",
        hue="percentile",
        col="payload",
        col_wrap=3,
        errorbar=None,
        height=3.2,
        aspect=1.2,
        palette="crest",
        order=sorted(df["setup"].unique(), key=setup_sort_key),
    )
    grid.set_axis_labels("", "RTT [ms]")
    grid.set_titles("{col_name}")
    for ax in grid.axes.flat:
        ax.tick_params(axis="x", rotation=25)
        ax.grid(axis="y", alpha=0.25)
    grid.figure.suptitle("Ping RTT Percentiles", y=1.04)
    save_figure(grid.figure, output_dir, "ping_rtt_percentiles", formats, dpi)
    plt.close(grid.figure)


def plot_added_delay(pd: Any, plt: Any, sns: Any, metrics: Any, output_dir: Path, formats: list[str], dpi: int, baseline: str) -> None:
    df = metrics[
        (metrics["category"] == "ping")
        & (metrics["metric"] == "rtt_median")
    ].copy()
    if maybe_warn_empty(df, "added one-way delay"):
        return

    baseline_candidates = df[df["setup"].str.lower() == baseline.lower()]
    if baseline_candidates.empty:
        baseline_candidates = df[df["setup"].str.lower().str.contains(baseline.lower(), regex=False)]
    if baseline_candidates.empty:
        print(f"skipping added delay: baseline setup '{baseline}' not found")
        return

    baseline_by_payload = baseline_candidates.groupby("payload_size_bytes")["value"].mean()
    df["baseline_rtt_ms"] = df["payload_size_bytes"].map(baseline_by_payload)
    df = df.dropna(subset=["baseline_rtt_ms"])
    df["added_one_way_ms"] = (df["value"] - df["baseline_rtt_ms"]) / 2.0
    df = df[df["setup"].str.lower() != baseline.lower()]
    if maybe_warn_empty(df, "added one-way delay"):
        return
    df["payload"] = df["payload_size_bytes"].astype("Int64").astype(str) + " B"

    fig, ax = plt.subplots(figsize=(8.2, 4.4))
    sns.barplot(
        data=df,
        x="setup",
        y="added_one_way_ms",
        hue="payload",
        errorbar=None,
        palette="flare",
        order=sorted(df["setup"].unique(), key=setup_sort_key),
        ax=ax,
    )
    ax.axhline(0, color="0.2", linewidth=0.8)
    ax.set_xlabel("")
    ax.set_ylabel("Added one-way delay [ms]")
    ax.set_title(f"Estimated Added One-Way Delay vs {baseline}")
    ax.tick_params(axis="x", rotation=25)
    ax.grid(axis="y", alpha=0.25)
    save_figure(fig, output_dir, "added_one_way_delay", formats, dpi)
    plt.close(fig)


def plot_throughput(pd: Any, plt: Any, sns: Any, metrics: Any, output_dir: Path, formats: list[str], dpi: int) -> None:
    df = metrics[
        (metrics["category"].isin(["iperf_tcp", "iperf_udp"]))
        & (metrics["metric"] == "throughput")
    ].copy()
    if maybe_warn_empty(df, "throughput"):
        return
    df["traffic"] = df["category"].map({"iperf_tcp": "TCP", "iperf_udp": "UDP"}) + " " + df["direction"].fillna("")

    fig, ax = plt.subplots(figsize=(9.2, 4.8))
    sns.barplot(
        data=df,
        x="setup",
        y="value",
        hue="traffic",
        errorbar=None,
        palette="mako",
        order=sorted(df["setup"].unique(), key=setup_sort_key),
        ax=ax,
    )
    ax.set_xlabel("")
    ax.set_ylabel("Throughput [Mbit/s]")
    ax.set_title("Throughput by Setup")
    ax.tick_params(axis="x", rotation=25)
    ax.grid(axis="y", alpha=0.25)
    save_figure(fig, output_dir, "throughput", formats, dpi)
    plt.close(fig)


def plot_udp_quality(pd: Any, plt: Any, sns: Any, metrics: Any, output_dir: Path, formats: list[str], dpi: int) -> None:
    df = metrics[
        (metrics["category"] == "iperf_udp")
        & (metrics["metric"].isin(["jitter", "loss"]))
    ].copy()
    if maybe_warn_empty(df, "UDP quality"):
        return
    df["direction"] = df["direction"].fillna("forward")
    df["metric_label"] = df["metric"].map({"jitter": "Jitter [ms]", "loss": "Loss [%]"})

    grid = sns.catplot(
        data=df,
        kind="bar",
        x="setup",
        y="value",
        hue="direction",
        col="metric_label",
        sharey=False,
        errorbar=None,
        height=3.8,
        aspect=1.25,
        palette="rocket",
        order=sorted(df["setup"].unique(), key=setup_sort_key),
    )
    grid.set_axis_labels("", "")
    grid.set_titles("{col_name}")
    for ax in grid.axes.flat:
        ax.tick_params(axis="x", rotation=25)
        ax.grid(axis="y", alpha=0.25)
    grid.figure.suptitle("UDP Jitter and Loss", y=1.04)
    save_figure(grid.figure, output_dir, "udp_jitter_loss", formats, dpi)
    plt.close(grid.figure)


def plot_sockperf(pd: Any, plt: Any, sns: Any, metrics: Any, output_dir: Path, formats: list[str], dpi: int) -> None:
    df = metrics[
        (metrics["category"] == "sockperf")
        & (metrics["metric"] == "avg_latency_usec")
    ].copy()
    if maybe_warn_empty(df, "sockperf"):
        return

    fig, ax = plt.subplots(figsize=(8.2, 4.2))
    sns.barplot(
        data=df,
        x="setup",
        y="value",
        hue="protocol",
        errorbar=None,
        palette="viridis",
        order=sorted(df["setup"].unique(), key=setup_sort_key),
        ax=ax,
    )
    ax.set_xlabel("")
    ax.set_ylabel("Average latency [us]")
    ax.set_title("Sockperf Application Latency")
    ax.tick_params(axis="x", rotation=25)
    ax.grid(axis="y", alpha=0.25)
    save_figure(fig, output_dir, "sockperf_latency", formats, dpi)
    plt.close(fig)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Plot benchmark summaries for thesis figures.")
    parser.add_argument(
        "inputs",
        nargs="*",
        type=Path,
        default=[Path("measurments"), Path("measurements")],
        help="Result directories or summary.json files. Defaults to ./measurments and ./measurements.",
    )
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=Path("documentation/thesis/figures/benchmark"),
        help="Directory for generated plots and CSV files.",
    )
    parser.add_argument(
        "--baseline",
        default="direct",
        help="Setup name used as baseline for added-delay plots.",
    )
    parser.add_argument(
        "--formats",
        default="pdf,svg,png",
        help="Comma-separated output formats, e.g. pdf,svg,png.",
    )
    parser.add_argument("--dpi", type=int, default=220, help="DPI for raster outputs such as PNG.")
    return parser


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    pd, plt, sns = import_plotting_stack()

    summary_files = find_summary_files(args.inputs)
    if not summary_files:
        parser.error("no summary.json files found in the provided inputs")

    formats = [item.strip().lower() for item in args.formats.split(",") if item.strip()]
    if not formats:
        parser.error("--formats must contain at least one format")

    args.out_dir.mkdir(parents=True, exist_ok=True)
    sns.set_theme(
        context="paper",
        style="whitegrid",
        font_scale=1.0,
        rc={
            "figure.dpi": args.dpi,
            "savefig.dpi": args.dpi,
            "axes.spines.right": False,
            "axes.spines.top": False,
        },
    )

    rows: list[dict[str, Any]] = []
    for path in summary_files:
        rows.extend(extract_summary(path))

    metrics = pd.DataFrame(rows)
    if metrics.empty:
        parser.error("summary files were found, but no plottable metrics were extracted")

    metrics["setup"] = metrics["setup"].astype(str)
    metrics = metrics.sort_values(
        by=["setup", "category", "metric", "payload_size_bytes", "direction", "protocol"],
        key=lambda col: col.map(lambda value: setup_sort_key(str(value)) if col.name == "setup" else value),
    )
    metrics_csv = args.out_dir / "benchmark_metrics.csv"
    metrics.to_csv(metrics_csv, index=False)
    print(f"wrote {metrics_csv}")

    summary_csv = args.out_dir / "benchmark_metric_summary.csv"
    grouped = (
        metrics.groupby(["setup", "category", "metric", "unit", "direction", "payload_size_bytes", "protocol"], dropna=False)[
            "value"
        ]
        .agg(["count", "mean", "median", "min", "max", "std"])
        .reset_index()
    )
    grouped.to_csv(summary_csv, index=False)
    print(f"wrote {summary_csv}")

    plot_ping_percentiles(pd, plt, sns, metrics, args.out_dir, formats, args.dpi)
    plot_added_delay(pd, plt, sns, metrics, args.out_dir, formats, args.dpi, args.baseline)
    plot_throughput(pd, plt, sns, metrics, args.out_dir, formats, args.dpi)
    plot_udp_quality(pd, plt, sns, metrics, args.out_dir, formats, args.dpi)
    plot_sockperf(pd, plt, sns, metrics, args.out_dir, formats, args.dpi)

    print(f"Processed {len(summary_files)} summary files.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())