# vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4

from dataclasses import dataclass
import os
import platform
from time import monotonic

import plot_layout


@dataclass(frozen=True)
class Metric:
    axis: str
    cumulative: bool


metrics = {
    "peak_nonpaged_pool": Metric("x1y1", False),
    "nonpaged_pool": Metric("x1y1", False),
    "num_page_faults": Metric("x1y2", True),
    "paged_pool": Metric("x1y1", False),
    "peak_paged_pool": Metric("x1y1", False),
    "peak_pagefile": Metric("x1y1", False),
    "peak_wset": Metric("x1y1", False),
    "private": Metric("x1y1", False),
    "rss": Metric("x1y1", False),
    "uss": Metric("x1y1", False),
    "data": Metric("x1y1", False),
    "shared": Metric("x1y1", False),
    "text": Metric("x1y1", False),
    "dirty": Metric("x1y1", False),
    "lib": Metric("x1y1", False),
    "vms": Metric("x1y1", False),
    "other_bytes": Metric("x1y1", True),
    "other_count": Metric("x1y2", True),
    "read_bytes": Metric("x1y1", True),
    "read_chars": Metric("x1y1", True),
    "read_count": Metric("x1y2", True),
    "write_bytes": Metric("x1y1", True),
    "write_chars": Metric("x1y1", True),
    "write_count": Metric("x1y2", True),
    "pfaults": Metric("x1y2", True),
    "pageins": Metric("x1y2", True),
    "minor_faults": Metric("x1y2", True),
    "major_faults": Metric("x1y2", True),
    "pss": Metric("x1y1", False),
    "pss_anon": Metric("x1y1", False),
    "pss_file": Metric("x1y1", False),
    "pss_shmem": Metric("x1y1", False),
    "shared_clean": Metric("x1y1", False),
    "shared_dirty": Metric("x1y1", False),
    "private_clean": Metric("x1y1", False),
    "private_dirty": Metric("x1y1", False),
    "referenced": Metric("x1y1", False),
    "anonymous": Metric("x1y1", False),
    "lazyfree": Metric("x1y1", False),
    "anonhugepages": Metric("x1y1", False),
    "shmempmdmapped": Metric("x1y1", False),
    "filepmdmapped": Metric("x1y1", False),
    "shared_hugetlb": Metric("x1y1", False),
    "private_hugetlb": Metric("x1y1", False),
    "swap": Metric("x1y1", False),
    "swappss": Metric("x1y1", False),
    "locked": Metric("x1y1", False),
}


@dataclass(frozen=True)
class Plot:
    name: str
    title: str
    ylabel: str
    y2label: str
    lines: list[str]


plots = [
    Plot(
        "memory",
        "libtorrent memory usage",
        "Memory Size",
        "",
        [
            "pss",
            "pss_file",
            "pss_anon",
            "rss",
            "dirty",
            "private_dirty",
            "private_clean",
            "lazyfree",
            "anonymous",
            "vms",
            "private",
            "paged_pool",
        ],
    ),
    Plot(
        "vm",
        "libtorrent vm stats",
        "",
        "count",
        [
            "pfaults",
            "pageins",
            "num_page_faults",
            "major_faults",
            "minor_faults",
        ],
    ),
    Plot(
        "read",
        "libtorrent disk I/O (read)",
        "Size",
        "count",
        [
            "read_bytes",
            "read_chars",
            "read_count",
        ],
    ),
    Plot(
        "write",
        "libtorrent disk I/O (write)",
        "Size",
        "count",
        [
            "write_bytes",
            "write_chars",
            "write_count",
        ],
    ),
]

if platform.system() == "Linux":

    def capture_sample(
        pid: int, start_time: float, output: dict[str, list[float]]
    ) -> None:
        try:
            with open(f"/proc/{pid}/smaps_rollup") as f:
                sample = f.read()
            with open(f"/proc/{pid}/stat") as f:
                sample2 = f.read()
            with open(f"/proc/{pid}/io") as f:
                sample3 = f.read()
            timestamp = monotonic() - start_time
        except Exception:
            return

        if "time" not in output:
            time_delta = 0.0
            output["time"] = [timestamp]
        else:
            time_delta = timestamp - output["time"][-1]
            output["time"].append(timestamp)

        for line in sample.split("\n"):
            if "[rollup]" in line:
                continue
            if line.strip() == "":
                continue

            key, value = line.split(":")
            val = int(value.split()[0].strip())
            key = key.strip().lower()

            if key not in output:
                output[key] = [val * 1024]
            else:
                output[key].append(val * 1024)

        stats = sample2.split()

        def add_counter(key: str, val: float) -> None:
            m = metrics[key]
            if key not in output:
                if m.cumulative:
                    output[key + "-raw"] = [val]
                    # we only have a single value, nothing to compare it against
                    val = 0
                output[key] = [val]
            else:
                if m.cumulative:
                    raw_val = val
                    val = (val - output[key + "-raw"][-1]) / time_delta
                    output[key + "-raw"].append(raw_val)
                output[key].append(val)

        add_counter("minor_faults", float(stats[9]))
        add_counter("major_faults", float(stats[11]))

        for line in sample3.split("\n"):
            if line == "":
                continue
            key, val2 = line.split(": ")
            if key == "rchar":
                add_counter("read_chars", float(val2))
            elif key == "wchar":
                add_counter("write_chars", float(val2))
            elif key == "read_bytes":
                add_counter("read_bytes", float(val2))
            elif key == "write_bytes":
                add_counter("write_bytes", float(val2))
            elif key == "syscr":
                add_counter("read_count", float(val2))
            elif key == "syscw":
                add_counter("write_count", float(val2))


# example output:
# 8affffff000-7fffba926000 ---p 00000000 00:00 0  [rollup]
#    Rss:               76932 kB
#    Pss:               17508 kB
#    Pss_Anon:          11376 kB
#    Pss_File:           6101 kB
#    Pss_Shmem:            30 kB
#    Shared_Clean:      65380 kB
#    Shared_Dirty:         88 kB
#    Private_Clean:        80 kB
#    Private_Dirty:     11384 kB
#    Referenced:        76932 kB
#    Anonymous:         11376 kB
#    LazyFree:              0 kB
#    AnonHugePages:         0 kB
#    ShmemPmdMapped:        0 kB
#    FilePmdMapped:         0 kB
#    Shared_Hugetlb:        0 kB
#    Private_Hugetlb:       0 kB
#    Swap:                  0 kB
#    SwapPss:               0 kB
#    Locked:                0 kB

else:
    import psutil

    def capture_sample(
        pid: int, start_time: float, output: dict[str, list[float]]
    ) -> None:
        try:
            p = psutil.Process(pid)
            mem = p.memory_full_info()
            io_cnt = p.io_counters()
            timestamp = monotonic() - start_time
        except Exception:
            return

        if "time" not in output:
            time_delta = 0.0
            output["time"] = [timestamp]
        else:
            time_delta = timestamp - output["time"][-1]
            output["time"].append(timestamp)

        for key in dir(mem):
            if key not in metrics:
                if not key.startswith("_") and key not in [
                    "pagefile",
                    "wset",
                    "count",
                    "index",
                ]:
                    print(f"missing key: {key}")
                continue

            val = getattr(mem, key)

            m = metrics[key]
            if key not in output:
                if m.cumulative:
                    output[key + "-raw"] = [val]
                    # we only have a single value, nothing to compare it against
                    val = 0
                output[key] = [val]
            else:
                if m.cumulative:
                    raw_val = val
                    val = (val - output[key + "-raw"][-1]) / time_delta
                    output[key + "-raw"].append(raw_val)

                output[key].append(val)

        for key in dir(io_cnt):
            if key not in metrics:
                if not key.startswith("_") and key not in [
                    "pagefile",
                    "wset",
                    "count",
                    "index",
                ]:
                    print(f"missing key: {key}")
                continue

            m = metrics[key]
            if key not in output:
                if m.cumulative:
                    output[key + "-raw"] = [val]
                    # we only have a single value, nothing to compare it against
                    val = 0
                output[key] = [val]
            else:
                if m.cumulative:
                    raw_val = val
                    val = (val - output[key + "-raw"][-1]) / time_delta
                    output[key + "-raw"].append(raw_val)

                output[key].append(val)


def print_output_to_file(out: dict[str, list[float]], filename: str) -> list[str]:
    if out == {}:
        return []

    with open(filename, "w+") as stats_output:
        non_zero_keys: set[str] = set()
        non_zero_keys.add("time")
        keys = out.keys()
        for key in keys:
            stats_output.write(f"{key} ")
        stats_output.write("\n")
        idx = 0
        while len(out["time"]) > idx:
            for key in keys:
                stats_output.write(f"{out[key][idx]:.2f} ")
                if out[key][idx] != 0:
                    non_zero_keys.add(key)
            stats_output.write("\n")
            idx += 1
    return [k if k in non_zero_keys else "" for k in keys]


def plot_output(filename: str, keys: list[str]) -> None:
    if "time" not in keys:
        return

    try:
        import matplotlib

        matplotlib.use("Agg")  # headless backend; no display required
        import matplotlib.pyplot as plt
    except ImportError:
        return

    output_dir, in_file = os.path.split(filename)

    # read the stats table written by print_output_to_file: the first line is
    # the column names, the rest are rows of floats.
    with open(filename) as f:
        header = f.readline().split()
        cols: dict[str, list[float]] = {name: [] for name in header}
        for line in f:
            fields = line.split()
            if len(fields) != len(header):
                continue
            for name, value in zip(header, fields):
                cols[name].append(float(value))

    if not cols.get("time"):
        return
    t = cols["time"]

    # print_output_to_file blanks all-zero columns to "" in `keys`; plot only
    # the columns that actually saw a non-zero sample, matching the previous
    # gnuplot output.
    nonzero = {k for k in keys if k}

    mib = 1024 * 1024
    dpi = 100
    # draw the lines a single device pixel wide; matplotlib's ~1.5 pt default
    # is thick for the dense, spiky disk-I/O curves.
    line_width = 72.0 / dpi
    for plot in plots:
        names = [n for n in plot.lines if n in nonzero and n in cols]
        if not names:
            continue

        # count-style (x1y2) metrics go on the left (primary) axis; size
        # (x1y1) metrics go on the right axis, in MB. both axes are always
        # created so the layout is the same across plots, and an axis with no
        # data is left unlabeled. the right (MB) axis tick labels land in the
        # reserved right margin (see subplots_adjust below).
        has_count = any(metrics[n].axis == "x1y2" for n in names)
        has_size = any(metrics[n].axis == "x1y1" for n in names)

        fig, ax = plt.subplots(figsize=(12.0, 7.0))
        ax_mb = ax.twinx()

        # draw colors from a single cycle in plot-line order so the two axes
        # don't both restart at the first color (twinx gives each axis its own
        # cycle otherwise), matching gnuplot's single sequence.
        color_cycle = plt.rcParams["axes.prop_cycle"].by_key()["color"]
        handles = []
        for i, name in enumerate(names):
            m = metrics[name]
            label = name + ("/s" if m.cumulative else "")
            color = color_cycle[i % len(color_cycle)]
            if m.axis == "x1y1":
                # x1y1 metrics are byte counts; show them in MB
                target = ax_mb
                ys = [v / mib for v in cols[name]]
            else:
                target = ax
                ys = cols[name]
            (line,) = target.step(
                t, ys, where="post", label=label, color=color, linewidth=line_width
            )
            handles.append(line)

        ax.set_xlim(left=0)
        ax.set_ylim(bottom=0)
        ax_mb.set_ylim(bottom=0)
        ax.set_xlabel("time (s)")
        ax.set_title(plot.title)
        # label each y-axis only when it has data
        if has_count:
            ax.set_ylabel(plot.y2label)
        if has_size:
            ax_mb.set_ylabel(f"{plot.ylabel} (MB)")
        # vertical grid always; horizontal grid on whichever axis has data
        ax.grid(True, axis="x")
        (ax if has_count else ax_mb).grid(True, axis="y")
        # always top-left; matplotlib's auto-placement does a poor job here
        ax.legend(handles=handles, loc="upper left")

        # pin the plot box to the shared summary-page margins (see plot_layout)
        # so these line up with the piece-pass-order and latency plots.
        fig.subplots_adjust(
            left=plot_layout.BOX_LEFT,
            right=plot_layout.BOX_RIGHT,
            bottom=0.09,
            top=0.93,
        )
        fig.savefig(f"{output_dir}/{in_file}-{plot.name}.png", dpi=dpi)
        plt.close(fig)
