#!/usr/bin/python3
# usage: $0: sensors_bin_path
# Run sensors under strace to discover where sensor files
# live on the current system and generate a draft file.
import warnings
import sys
import os.path
import tempfile
import subprocess
import argparse
import os

def fix_label(s, nodash=False):
    if s[-2:] == "\\n":
        s = s[:-2]
    s = s.replace(" ", "_")
    if nodash:
        s = s.replace("-", "_")
    return s

IGNORE_TYPES = [
    "crit",
    "alarm",
    "hyst",
]

def get_freq_file(temp):
    met_file1 = temp.replace("max_fr", "cur_fr")
    if os.path.isfile(met_file1):
        return met_file1
    met_file2 = temp.replace("cpuinfo_max_freq", "scaling_cur_freq")
    if os.path.isfile(met_file2):
        return met_file2

def main():
    parser = argparse.ArgumentParser(
        description="Generate filesingle metric list")
    parser.add_argument("--sensors", metavar="SENSORS-PATH",
                        default="/usr/bin/sensors",
                        help="Full path to lm_sensor 'sensors' program"
                       )
    parser.add_argument("--lscpu", metavar="LSCPU-PATH", default="/usr/bin/lscpu",
                        help="Full path to 'lscpu' program"
                       )
    parser.add_argument("--nodash", default=False, action='store_true',
                        help="enable replacing - with _ in names."
                       )
    parser.add_argument("--test-lscpu", metavar="TEST-INPUT", required=False,
                        help="path to strace result file for parser testing"
                       )
    parser.add_argument("--test-sensors", metavar="TEST-INPUT", required=False,
                        help="path to strace result file for parser testing"
                       )

    args = parser.parse_args()

    if args.test_sensors:
        if not os.path.isfile(args.test_sensors):
            print(f"Missing {args.test_sensors}. Specify a trace log.")
            exit(1)
    else:
        if not os.path.isfile(args.sensors):
            print(f"Missing {args.sensors}. Need --sensors=<progpath>")
            exit(1)

    if args.test_lscpu:
        if not os.path.isfile(args.test_lscpu):
            print(f"Missing {args.test_lscpu}. Specify a trace log.")
            exit(1)
    else:
        if not os.path.isfile(args.lscpu):
            print(f"Missing {args.lscpu}. Need --lscpu=<progpath>")
            exit(1)

    sensors_bin = args.sensors
    lscpu_bin = args.lscpu
    nodash = args.nodash

    # Build and parse sensors trace
    trace_file = tempfile.mktemp(dir="/tmp", suffix=".trace", prefix="ldms-config-sensors.")
    if args.test_sensors:
        trace_file = args.test_sensors
    else:
        cmd = f"script -c 'strace -e trace=openat,open,read {sensors_bin} -u' {trace_file} > /dev/null"
        trval = subprocess.call(cmd, shell=True)
        if trval != 0:
            print("Unable to collect sensors config output")
            exit(1)

    devices = { }
    with open(trace_file) as file:
        lines = file.readlines()

    if not args.test_sensors:
        os.unlink(trace_file)

#review    lastopen = None
#review    lastread = None
    last_file = None
    last_file_type = None
    last_group = None

    for line in lines:
        line = line.strip()
        if len(line) < 5:
            continue
        if line.startswith("openat("):
            temp = line.split(',')[1].strip().strip('"')
            if temp.startswith('/sys/class/'):
                last_file = temp
                if temp.endswith("/name"):
                    last_device = os.path.dirname(temp)
                    devices[last_device] = {"items": []}
                    last_file_type = "name"
                else:
                    suffix = temp.split('_')[-1]
                    last_file_type = suffix
            else:
                last_file = None
                last_file_type = None
                last_device = None
            continue
        if line.startswith("open("):
            temp = line.split(',')[0].split("(")[1].strip().strip('"')
            if temp.startswith('/sys/class/'):
                last_file = temp
                if temp.endswith("/name"):
                    last_device = os.path.dirname(temp)
                    devices[last_device] = {"items": []}
                    last_file_type = "name"
                else:
                    suffix = temp.split('_')[-1]
                    last_file_type = suffix
            else:
                last_file = None
                last_file_type = None
                last_device = None
            continue
        if line.startswith("read("):
            temp = line.split(',')[1].strip().strip('"')
            if not temp:
                continue
            if last_file_type is None:
                continue
            if last_file_type == "name":
                devices[last_device]['name'] = fix_label(temp)
                continue
            if last_file_type == "label":
                last_device = os.path.dirname(last_file)
                base_name = os.path.basename(last_file)
                split_index = base_name.rfind("_")
                item = base_name[:split_index]
                if not "group" in devices[last_device]:
                    devices[last_device]["group"] = last_group
                if not item in devices[last_device]["items"]:
                    devices[last_device]["items"].append(item)
                    devices[last_device][item] = {}
                    devices[last_device][item]["label"] = fix_label(temp)
                continue
            if last_file_type in ["input", "average"]:
                last_device = os.path.dirname(last_file)
                base_name = os.path.basename(last_file)
                split_index = base_name.rfind("_")
                item = base_name[:split_index]
                if not "group" in devices[last_device]:
                    devices[last_device]["group"] = last_group
                if not item in devices[last_device]["items"]:
                    devices[last_device]["items"].append(item)
                    devices[last_device]["group"] = last_group
                    devices[last_device][item] = {}
                    devices[last_device][item]["label"] = fix_label(item)
                devices[last_device][item]["inputfile"] = last_file
                continue
            if last_file_type in IGNORE_TYPES:
                continue
            continue
        # exclude   lines with "xxxxx:" and "sdf: bar" but allow : in mlx device names.
        if ": " not in line and line.strip()[-1] != ':':
            last_group = line

    for device, data in devices.items():
        if data["items"]:
            for item in data["items"]:
                xg = data["group"] if data["group"] else "NOGROUP"
                item_label = fix_label(".".join([xg, data[item]["label"] ]), nodash)
                print(f"{item_label} {data[item]['inputfile']} S64 -1")

    devices = {}
    # Build and parse lscpu trace
    if args.test_lscpu:
        lscpu_trace_file = args.test_lscpu
    else:
        trace_file = tempfile.mktemp(dir="/tmp", suffix=".trace", prefix="ldms-config-sensors.")
        lscpu_trace_file = trace_file + "a"
        cmd = f"script -c 'strace -e trace=open,openat {lscpu_bin}' {trace_file} > /dev/null; grep cpuinfo_max_freq {trace_file} > {lscpu_trace_file}"
        trval = subprocess.call(cmd, shell=True)
        if trval != 0:
            print("#Unable to collect lscpu config output")
            exit(1)
    with open(lscpu_trace_file) as file:
        lines = file.readlines()

    if not args.test_lscpu:
        os.unlink(trace_file)
        os.unlink(lscpu_trace_file)

    for line in lines:
        line = line.strip()
        if len(line) < 5:
            continue
        if line.startswith("openat("):
            temp = line.split('"')[1]
            met_file = get_freq_file(temp)
            parts = temp.split("/")
            device = parts[-3]
            item = "cur_freq"
            if device not in devices:
                devices[device] = {"items": [], "group": device}
            devices[device]["items"].append(item)
            devices[device][item] = {"label": "cur_freq", "inputfile": met_file}
        elif line.startswith("open("):
            temp = line.split('"')[1]
            met_file = get_freq_file(temp)
            parts = temp.split("/")
            device = parts[-3]
            item = "cur_freq"
            if device not in devices:
                devices[device] = {"items": [], "group": device}
            devices[device]["items"].append(item)
            devices[device][item] = {"label": "cur_freq", "inputfile": met_file}
        else:
            print(f"Unexpected line from lscpu trace: {line}")

    # CPU temps reported in millicentigrade
    # Intel power averages in microwatts
    # CPU freq reported in millihertz
    output_lines = []
    for device, data in devices.items():
        if data["items"]:
            for item in data["items"]:
                item_label = fix_label(".".join([data["group"], data[item]["label"] ]), nodash)
                output_line = f"{item_label} {data[item]['inputfile']} S64 -1"
                output_lines.append(output_line)

    for line in sorted(output_lines):
        print(line)

if __name__ == "__main__":
    main()
