#!/usr/bin/env python3
#
# -*- coding: utf-8 -*-
# vim: ts=4 sw=4 tw=100 et ai si
#
# Copyright (C) 2022 Intel Corporation
# SPDX-License-Identifier: BSD-3-Clause
#
# Author: Antti Laakso <antti.laakso@linux.intel.com>

"""Test data generator, for collecting and creating data used for testing."""

import os
import re
import stat
import sys
import logging
from pathlib import Path
from pepclibs.msr import EnergyPerfBias, FSBFreq, HWPRequest, HWPRequestPkg, MiscFeatureControl
from pepclibs.msr import PlatformInfo, PowerCtl, PCStateConfigCtl, PMEnable, TurboRatioLimit
from pepclibs.msr import TurboRatioLimit1, PackagePowerLimit, PackagePowerInfo, RaplPowerUnit
from pepclibs.msr import SwLTROvrd, PMLogicalId, HWPCapabilities
from pepclibs.helperlibs import ArgParse, Logging, ProcessManager, YAML
from pepclibs.helperlibs.Exceptions import Error

OWN_NAME = "tdgen"
VERSION = "0.1"
LOG = logging.getLogger()
Logging.setup_logger(prefix=OWN_NAME)

# pylint: disable=line-too-long
CPUInfoData = {
    "commands" : [
        {"command": "lscpu",
         "dirname": "lscpu"}],
    "files" : [
        {"path": "/proc/cpuinfo",
         "readonly": True}],
    "inlinefiles" : [
        # Command for reading die info.
        {"command": r"grep -H '.*' "
                    r"/sys/devices/system/cpu/cpu[0-9]*/topology/die_id "
                    r"/sys/devices/system/cpu/cpu[0-9]*/topology/die_cpus_list",
         "separator": ":",
         "readonly": True,
         "dirname": "die-info",
         "filename": "die.txt"},
        # Command for reading module info.
        {"command": r"grep -H '.*' "
                    r"/sys/devices/system/cpu/cpu[0-9]*/cache/index[0-9]/id "
                    r"/sys/devices/system/cpu/cpu[0-9]*/cache/index[0-9]/shared_cpu_list",
         "separator": ":",
         "readonly": True,
         "dirname": "module-info",
         "filename": "module.txt"},
        # Command for reading online CPUs and all CPUs.
        {"command": r"grep -H '.*' "
                    r"/sys/devices/system/cpu/online "
                    r"/sys/devices/system/cpu/present",
         "separator": ":",
         "readonly": True,
         "dirname": "cpu-info",
         "filename": "cpu.txt"},
        # Command for reading node info.
        {"command": r"grep -H '.*' "
                    r"/sys/devices/system/node/online "
                    r"/sys/devices/system/node/node[0-9]/cpulist",
         "separator": ":",
         "readonly": True,
         "dirname": "node-info",
         "filename": "node.txt"},
        # Command for reading hybrid topology info.
        {"command": r"grep -H '.*' "
                    r"/sys/devices/cpu_atom/cpus "
                    r"/sys/devices/cpu_core/cpus",
         "separator": ":",
         "readonly": True,
         "dirname": "hybrid-info",
         "filename": "hybrid.txt"}],
    "msrs" :
        {"addresses": [
            EnergyPerfBias.MSR_ENERGY_PERF_BIAS,
            FSBFreq.MSR_FSB_FREQ,
            HWPRequest.MSR_HWP_REQUEST,
            HWPRequestPkg.MSR_HWP_REQUEST_PKG,
            HWPCapabilities.MSR_HWP_CAPABILITIES,
            MiscFeatureControl.MSR_MISC_FEATURE_CONTROL,
            PlatformInfo.MSR_PLATFORM_INFO,
            PowerCtl.MSR_POWER_CTL,
            PCStateConfigCtl.MSR_PKG_CST_CONFIG_CONTROL,
            PMEnable.MSR_PM_ENABLE,
            TurboRatioLimit.MSR_TURBO_RATIO_LIMIT,
            TurboRatioLimit1.MSR_TURBO_RATIO_LIMIT1,
            TurboRatioLimit1.MSR_TURBO_GROUP_CORECNT,
            TurboRatioLimit1.MSR_TURBO_RATIO_LIMIT_CORES,
            PackagePowerLimit.MSR_PKG_POWER_LIMIT,
            PackagePowerInfo.MSR_PKG_POWER_INFO,
            RaplPowerUnit.RAPL_POWER_UNIT,
            SwLTROvrd.MSR_SW_LTR_OVRD,
            PMLogicalId.MSR_PM_LOGICAL_ID],
         "separator1": ":",
         "separator2": "|",
         "dirname": "msr",
         "filename": "msr.txt"}
}

ASPMData = {
    "inlinefiles" : [
        {"command": r"grep -H '.*' "
                    r"/sys/module/pcie_aspm/parameters/policy",
         "separator": ":",
         "readonly": False,
         "dirname": "aspm-info",
         "filename": "aspm.txt"}]
}

CPUOnlineData = {
    "inlinefiles" : [
        {"command": r"grep -H '.*' "
                    r"/sys/devices/system/cpu/cpu[0-9]*/online",
         "separator": ":",
         "readonly": False,
         "dirname": "cpuonline-info",
         "filename": "cpuonline.txt"}],
    "inlinedirs" : [
        {"command": r"find /sys/devices/system/cpu -type d -regextype posix-extended -regex " \
                    r"'.*cpu([[:digit:]]+)'",
         "dirname": "cpuonline-info",
         "filename": "cpuonline-dirs.txt"}],
}

CStatesData = {
    "files" : [
        {"path": "/proc/cmdline",
         "readonly": True}],
    "inlinefiles" : [
        {"command": r"grep -H --directories=skip '.*' "
                    r"/sys/devices/system/cpu/cpu[0-9]*/cpuidle/state[0-9]/* "
                    r"/sys/devices/system/cpu/cpuidle/*",
         "separator": ":",
         "readonly": False,
         "dirname": "cstates",
         "filename": "cstates.txt"}]
}

PStatesData = {
    "prepare-cmds" : [
        "modprobe intel_uncore_frequency",
        "modprobe msr",
        "pepc pstates config --min-freq min --max-freq max --cpus all",
        "pepc pstates config --min-uncore-freq min --max-uncore-freq max --cpus all",
        "pepc cstates config --enable all --cpus all"],
    "inlinefiles" : [
        {"command": r"grep -H --directories=skip '.*' "
                    r"/sys/devices/system/cpu/cpufreq/policy[0-9]*/*",
         "separator": ":",
         "readonly": False,
         "dirname": "pstates",
         "filename": "pstates.txt"},
        {"command" : r"grep -H --directories=skip '.*' "
                     r"/sys/devices/system/cpu/intel_pstate/*",
         "separator": ":",
         "readonly": False,
         "dirname": "pstates",
         "filename": "intel_pstates.txt"},
        {"command" : r"grep -H --directories=skip '.*' "
                     r"/sys/devices/system/cpu/intel_uncore_frequency/*/*",
         "separator": ":",
         "readonly": False,
         "dirname": "pstates",
         "filename": "uncore.txt"},
        {"command" : r"grep -H '.*' "
                     r"/sys/devices/system/cpu/cpu[0-9]*/power/energy_perf_bias",
         "separator": ":",
         "readonly": False,
         "dirname": "pstates",
         "filename": "epb.txt"},
        {"command" : r"grep -H '.*' "
                     r"/sys/devices/system/cpu/cpu[0-9]*/acpi_cppc/* --exclude '*_ctrs'",
         "separator": ":",
         "readonly": False,
         "dirname": "pstates",
         "filename": "cppc.txt"},
        {"command" : r"grep -H '.*' "
                     r"/sys/devices/system/cpu/cpufreq/policy[0-9]*/energy_performance_preference",
         "separator": ":",
         "readonly": False,
         "dirname": "pstates",
         "filename": "epp.txt"},
        {"command" : r"grep -H '.*' "
                     r"/sys/devices/system/cpu/cpufreq/policy[0-9]*/energy_performance_available_preferences",
         "separator": ":",
         "readonly": True,
         "dirname": "pstates",
         "filename": "epp_policies.txt"},
        {"command" : r"grep -H '.*' "
                     r"/sys/devices/system/cpu/cpu[0-9]*/cpufreq/bios_limit",
         "separator": ":",
         "readonly": True,
         "dirname": "pstates",
         "filename": "bios_limit.txt"}]
}

SystemctlData = {
    "commands" : [
        {"command": "systemctl is-active -- 'tuned'",
         "dirname": "systemctl-tuned-active",
         "ignore_exitcode": True}],
}

PowerData = {
}

TpmiData = {
    "recursive_copy" : [
        {"path": "/sys/kernel/debug/tpmi-.*"}],
}
# pylint: disable=line-too-long

MODULE_TESTDATA = {
    "CPUInfo" : CPUInfoData,
    "ASPM" : ASPMData,
    "CPUOnline" : CPUOnlineData,
    "CStates" : CStatesData,
    "PStates" : PStatesData,
    "Systemctl" : SystemctlData,
    "Power" : PowerData,
    "TPMI" : TpmiData,
    }

def build_arguments_parser():
    """A helper function which parses the input arguments."""

    text = f"{OWN_NAME} - Test data generator, for collecting and creating test data."
    parser = ArgParse.SSHOptsAwareArgsParser(description=text, prog=OWN_NAME, ver=VERSION)

    ArgParse.add_ssh_options(parser)

    text = """Path to the directory to store the output of the commands at. Default value is the
              name of the host the command is run on. See the '-H' option."""
    parser.add_argument("-o", "--outdir", type=Path, help=text)

    return parser

def parse_arguments():
    """Parse input arguments."""

    parser = build_arguments_parser()
    args = parser.parse_args()

    return args

def collect_cmd_output(cmdinfo, pman, outdir):
    """
    Run the command defined in 'cmdinfo' and save the output to a file. The arguments are as
    follows.
      * cmdinfo - a dictionary of a command to run, see 'MODULE_TESTDATA'.
      * pman - the process manager object that defines the remote host to run the 'cmdinfo' on.
      * outdir - the directory to save the command output to.
    """

    datapath = outdir / cmdinfo["dirname"]
    os.makedirs(datapath, exist_ok=True)

    res = pman.run(cmdinfo["command"])
    if res.exitcode != 0 and not cmdinfo.get("ignore_exitcode"):
        LOG.error("running command '%s' failed and returned '%s'", cmdinfo["command"], res.exitcode)

    for fname, data in ("stdout", res.stdout), ("stderr", res.stderr):
        path = datapath / f"{fname}.txt"

        with open(path, "w", encoding="utf-8") as fobj:
            fobj.write(data)

def collect_files(cmdinfo, pman, outdir):
    """
    Read the content of files by running the command defined in 'cmdinfo' and save the output to a
    file. The arguments are as follows.
      * cmdinfo - a dictionary of a command to run to collect the file contents, see
                  'MODULE_TESTDATA'.
      * pman - the process manager object that defines the remote host to read the files from.
      * outdir - the directory to save the output to.
    """

    cmdpath = outdir / cmdinfo["dirname"]
    os.makedirs(cmdpath, exist_ok=True)

    res = pman.run(cmdinfo["command"])
    if res.exitcode != 0:
        LOG.notice("running command '%s' failed and returned '%s'",
                   cmdinfo["command"], res.exitcode)

    path = cmdpath / cmdinfo["filename"]
    with open(path, "w", encoding="utf-8") as fobj:
        fobj.write(res.stdout)

def collect_msrs(msrinfo, pman, outdir):
    """
    Read the values of the MSR registers defined in 'msrinfo' and save the output to a file. The
    arguments are as follows.
      * msrinfo - a dictionary of a MSR registers to read, see 'MODULE_TESTDATA'.
      * pman - the process manager object that defines the remote host to read the MSR values from.
      * outdir - the directory to save the command output to.
    """

    lines, _ = pman.run_verify("lscpu -p=cpu", join=False)

    cpus = []
    for line in lines:
        if line.startswith("#"):
            continue

        cpu = int(line.strip())
        cpus.append(cpu)

    cmdpath = outdir / msrinfo["dirname"]
    os.makedirs(cmdpath, exist_ok=True)

    path = cmdpath / msrinfo["filename"]
    with open(path, "w+", encoding="utf-8") as fobj:
        for cpu in cpus:
            line = f"/dev/cpu/{cpu}/msr{msrinfo['separator1']}"

            for addr in msrinfo["addresses"]:
                result = pman.run(f"rdmsr {addr} -p {cpu}")
                if result.exitcode != 0:
                    continue

                value = result.stdout.strip()
                line += f"{addr}{msrinfo['separator2']}{value} "

            fobj.write(line + "\n")

def copy_file(pman, src, outdir):
    """
    Copy file contents. The arguments are as follows.
      * pman - the process manager object that defines the remote host to read the files from.
      * src - path to the source file.
      * outdir - the destination to save 'src'.
    """

    dst = Path(outdir / src.lstrip("/"))
    os.makedirs(dst.parent, exist_ok=True)

    # In some cases '/proc/cpuinfo' is not fully copied when using 'scp' or 'rsync'.
    res = pman.run(f"cat {src}")
    if res.exitcode != 0:
        LOG.notice("running command 'cat %s' failed and returned '%s'", src, res.exitcode)

    with open(dst, "w", encoding="utf-8") as fobj:
        fobj.write(res.stdout)

def copy_dir(pman, src, outdir):
    """
    Copy directory contents recursively. The arguments are as follows.
      * pman - the process manager object that defines the remote host to copy the directory from.
      * src - path to the source directory. The last part of the 'src' can be a regular expression,
              in which case all the matching directories will be copied.
      * outdir - path to the destination directory.

    Return a list of dictionaries for the files that were copied.
    """

    src = Path(src)
    files = []

    for name, path, mode in pman.lsdir(src.parent):
        if not re.fullmatch(src.name, name):
            continue

        is_dir = stat.S_ISDIR(mode)
        is_readonly = not mode & stat.S_IWUSR

        if is_dir:
            files += copy_dir(pman, path / ".*", outdir)
        else:
            copy_file(pman, str(path), outdir)
            files += [{"path": path, "readonly": is_readonly}]

    if not files:
        files = [{"path": src.parent}]

    return files

def generate_config_file(modname, testdata, outdir):
    """Generate configuration file for python module name 'modname' from testdata 'testdata'."""

    with open(outdir / f"{modname}.yaml", "w", encoding="utf-8") as fobj:
        fobj.write(f"# This file was generated by the '{OWN_NAME}' tool.\n")
        YAML.dump(testdata, fobj)

def main():
    """Script entry point."""

    try:
        args = parse_arguments()

        # pylint: disable=no-member
        if args.hostname == "localhost":
            args.username = args.privkey = args.timeout = None

        with ProcessManager.get_pman(args.hostname, username=args.username,
                                     privkeypath=args.privkey, timeout=args.timeout) as pman:
            outdir = args.outdir
            if not outdir:
                outdir = Path(pman.hostname)

            pman.run_verify("pepc cpu-hotplug online --cpus all")

            for modname, testdata in MODULE_TESTDATA.items():
                datapath = outdir / modname

                if "prepare-cmds" in testdata:
                    for command in testdata["prepare-cmds"]:
                        pman.run(command)
                    del testdata["prepare-cmds"]

                if "commands" in testdata:
                    for cmdinfo in testdata["commands"]:
                        collect_cmd_output(cmdinfo, pman, datapath)
                        cmdinfo["dirname"] = f"{modname}/{cmdinfo['dirname']}"

                if "recursive_copy" in testdata:
                    files = []
                    for directory in testdata["recursive_copy"]:
                        files += copy_dir(pman, directory["path"], datapath)

                    del testdata["recursive_copy"]
                    testdata["recursive_copy"] = files

                if "files" in testdata:
                    for file in testdata["files"]:
                        copy_file(pman, file["path"], datapath)

                for section in ("inlinedirs", "inlinefiles"):
                    if section not in testdata:
                        continue

                    for cmdinfo in testdata[section]:
                        collect_files(cmdinfo, pman, datapath)
                        cmdinfo["dirname"] = f"{modname}/{cmdinfo['dirname']}"
                        # We do not need command used to collect file contents, remove it.
                        del cmdinfo["command"]

                if "msrs" in testdata:
                    collect_msrs(testdata["msrs"], pman, datapath)
                    testdata["msrs"]["dirname"] = f"{modname}/{testdata['msrs']['dirname']}"

                generate_config_file(modname, testdata, outdir)

    except KeyboardInterrupt:
        LOG.info("\nInterrupted, exiting")
        return -1
    except Error as err:
        LOG.error_out(err)

    return 0

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