#!/bin/python

import time
import subprocess

from os import listdir
from sys import exit, argv
from time import sleep
from pathlib import Path


script_dir = Path(__file__).resolve().parent

output_csv_fn = lambda a: f"metrics_{a}.csv"
output_plot_fn = lambda a: f"plot_{a}.png"

INTERVAL     = 1.5
CSV_HEADER   = "ts,cpu_percent,ram_mb,open_fds"
GNUPLOT_FILE = script_dir / "plot_metrics.gp"
DIR_CSV      = Path("metrics")
DIR_PLOT     = Path("output")


def run(pid, fp):
    create_new_file_or_exit(fp);

    proc1, total1 = current_cpu(pid)
    p = Path(f"/proc/{pid}")

    num_cpus = open_file_or_exit("/proc/cpuinfo")
    num_cpus = num_cpus.count("processor\t:")

    print(CSV_HEADER)

    while p.exists():
        proc2, total2 = current_cpu(pid)
        total = total2 - total1

        if total == 0:
            proc1 = proc2
            total1 = total2
            continue

        cpu_percent = 100 * (proc2 - proc1) / total * num_cpus
        cpu_percent = round(cpu_percent, 1)

        proc1 = proc2
        total1 = total2

        ram = open_file_or_exit(f"/proc/{pid}/status")
        ram = ram.split("\n")

        for line in ram:
            if line.startswith("VmRSS"):
                ram = line.split()
                mb  = round(int(ram[1]) / 1024, 1)
                ram = f"{mb}"

        fds = Path(f"/proc/{pid}/fd")
        fds = len(listdir(fds))

        ts = int(time.monotonic())
        csv_line = f"{ts},{cpu_percent},{ram},{fds}"

        print(csv_line)
        line_append(fp, csv_line)
        sleep(INTERVAL)

    print("Terminated.")


def current_cpu(pid):
    fields = open_file_or_exit(f"/proc/{pid}/stat")
    fields = fields.split()

    utime = int(fields[13])
    stime = int(fields[14])

    cpu_fields = open_file_or_exit("/proc/stat")
    cpu_fields = cpu_fields.split("\n")[0].split()

    total = sum(int(x) for x in cpu_fields[1:])
    return utime + stime, total


def open_file_or_exit(fp):
    try:
        with open(fp, 'r', encoding="utf8") as f:
            return f.read()
    except Exception as e:
        print(e)
        exit(1)


def create_new_file_or_exit(fn):
    try:
        with open(fn, "w+", encoding="utf8") as f:
            f.write(CSV_HEADER + "\n")
    except Exception as e:
        print(e)
        exit(1)


def line_append(fn, line):
    try:
        with open(fn, "a+", encoding="utf8") as f:
            f.write(line + "\n")
    except Exception as e:
        print(e)
        exit(1)


def run_command(cmd):
    try:
        subprocess.run(cmd, capture_output=True, text=True, check=True)
    except subprocess.CalledProcessError as e:
        print(f"Command failed {e.returncode}: {e.stderr}")
        exit(1)


if __name__ == "__main__":
    args = argv

    if len(args) == 1 or len(args) > 4:
        print("Usage: record.py PID [-p, -plot]")
        exit(1)

    DIR_CSV.mkdir(exist_ok=True)
    DIR_PLOT.mkdir(exist_ok=True)

    if "-po" in args:
        if len(args) != 3:
            print("Need filepath to csv file")
            exit(1)

        fp = Path(args[2])
        out = fp.name
        out = Path(out).with_suffix(".png")
        gnuplot_cmd = [
            "gnuplot",
            "-e",
            f"file='{fp}'",
            "-e",
            f"set output '{out}'",
            GNUPLOT_FIL
        ]
        run_command(gnuplot_cmd)
        exit(0)

    pid = args[1]
    ts  = int(time.monotonic())
    fp  = DIR_CSV / Path(output_csv_fn(f"{pid}_{ts}"))
    out = DIR_PLOT / Path(output_plot_fn(f"{pid}_{ts}"))

    try:
        run(pid, fp)
    except KeyboardInterrupt:
        print("Terminated by user.")

    if "-p" in args:
        gnuplot_cmd = [
            "gnuplot",
            "-e",
            f"file='{fp}'",
            "-e",
            f"set output '{out}'",
            GNUPLOT_FILE
        ]
        run_command(gnuplot_cmd)