#!/usr/bin/python3
#
# -*- coding: utf-8 -*-
# vim: ts=4 sw=4 tw=100 et ai si
#
# Copyright (C) 2023-2024 Intel Corporation
# SPDX-License-Identifier: BSD-3-Clause
#
# Author: Tero Kristo <tero.kristo@linux.intel.com>

"""wult-freq-helper - a helper for scaling CPU/uncore frequencies in a loop."""

import argparse
import contextlib
import sys
import time
from pathlib import Path
from pepclibs import CPUInfo, _UncoreFreq, _CPUFreq, _SysfsIO
from pepclibs.helperlibs import Logging, ArgParse, Trivial, LocalProcessManager
from pepclibs.helperlibs.Exceptions import Error
from wultlibs.helperlibs import Human

TOOLNAME = "wult-freq-helper"
VERSION = "1.0"

CPU_SYSFS_BASE = Path("/sys/devices/system/cpu")
UNCORE_FREQ_SYSFS_BASE = CPU_SYSFS_BASE / "intel_uncore_frequency"

_LOG = Logging.getLogger(Logging.MAIN_LOGGER_NAME).configure(prefix=TOOLNAME)

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

    text = f"""{TOOLNAME} - a tool for repeatedly setting CPU/uncore frequency between two values,
              in order to induce frequency CPU/uncore frequency and voltage scaling."""
    parser = ArgParse.ArgsParser(description=text, prog=TOOLNAME, ver=VERSION)

    parser.add_argument("-n", "--nowait", help="Don't wait for frequency to reach target value.",
                        action="store_true")

    text = """Specify parameters for a single frequency modification operation. 'SPEC' is specified
              as 'TYPE:ID:MIN:MAX', where:
              TYPE should be 'cpu' or 'uncore', specifies whether CPU or uncore frequency should be
              modified;
              ID is either CPU number or uncore domain ID to modify the frequency for (e.g. 'cpu:12:..' would target CPU12);
              MIN is the minimum CPU/uncore frequency value;
              MAX is the maximum CPU/uncore frequency value.
              For example, '-s cpu:2:800MHz:1GHz' would set the frequency for CPU2 repeatedly
              between 800MHz and 1GHz. To modify frequency of multiple CPUs or uncore domains,
              specify '--specs' multiple times."""
    parser.add_argument("-s", "--spec", action="append", help=text, dest="specs")

    text = """Sleep time between each frequency modification loop, in microseconds. By default,
              50000us"""
    parser.add_argument("--sleep", type=int, help=text, default=50000)

    parser.add_argument("--print-module-paths", action="store_true", help=argparse.SUPPRESS)

    args = parser.parse_args()
    return args

def sysfs_read(fname):
    """Read and return the contents of a sysfs file 'fname' as an integer."""

    with open(fname, "r", encoding="utf-8") as fobj:
        try:
            return Trivial.str_to_int(fobj.read())
        except OSError as err:
            raise Error("failed to read sysfs file '{fname}'") from err

def sysfs_write(fname, val):
    """Write 'val' to a sysfs file named 'fname'."""

    _LOG.debug("writing '%s' to '%s'", val, fname)
    with open(fname, "w", encoding="utf-8") as fobj:
        fobj.write(str(val))

def parse_freq(freq_type, f, fmin, fmax):
    """
    Parse a frequency string and return it as an integer. Arguments are as follows.
      * freq_type - type of frequency, 'cpu' or 'uncore'.
      * f - frequency string to be parsed.
      * fmin - minimum acceptable value for the frequency.
      * fmax - maximum acceptable value for the frequency.
    """

    if f == "min":
        f = fmin
    elif f == "max":
        f = fmax
    else:
        f = int(Human.parse_human(f, unit="Hz", integer=True) / 1000)

    if f < fmin or f > fmax:
        raise Error(f"frequency {f} out of range [{fmin}-{fmax}] for {freq_type}")

    return f

def set_freq(args, spec, cfg):
    """
    Set either uncore or CPU frequency. Arguments are as follows.
      * args - command line arguments.
      * spec - frequency specification dict.
      * cfg - name of the config, either 'min' or 'max'.
    """

    f = spec[cfg]
    obj = spec["obj"]

    if cfg == "min":
        obj.set_min_freq(f, spec["id"])
        obj.set_max_freq(f, spec["id"])
    else:
        obj.set_max_freq(f, spec["id"])
        obj.set_min_freq(f, spec["id"])

    if args.nowait:
        return

    count = 0

    while True:
        curf = next(obj.get_cur_freq(spec["id"]))[-1]
        if f == curf:
            break
        count += 1
        if count > 100000 and (count % 100000) == 0:
            _LOG.warning("waiting for %s%d to reach %d, current %d, loops=%d", spec["type"],
                         spec["num"], f, curf, count)
            if count > 500000:
                raise Error(f"unable to change frequency of {spec['type']}{spec['num']} to {f}")

def parse_freq_specs(args, cpuinfo, ucfobj, cpufobj):
    """
    Parse raw frequency specification data, and return an array of parsed frequency specs. Arguments
    are as follows.
      * args - raw frequency specification data.
      * cpuinfo - 'CPUInfo' object.
      * ucfobj - '_UncoreFreqSysfs' object.
      * cpufobj - '_CPUFreq' object.
    """

    parsed_specs = []

    for spec_raw in args.specs:
        tokens = spec_raw.split(":")
        if len(tokens) != 4:
            raise Error(f"bad frequency spec '{spec_raw}', expected four elements separated with "
                        "':'")

        freq_type = tokens[0]
        if freq_type not in ("cpu", "uncore"):
            raise Error(f"bad frequency type '{freq_type}'. Only 'cpu', 'uncore' are supported.")

        num = Trivial.str_to_int(tokens[1], what="frequency ID")
        siblings = None

        if freq_type == "cpu":
            obj = cpufobj
            spec_id = (num,)
            fmin = next(obj.get_min_freq_limit(spec_id))[1]
            fmax = next(obj.get_max_freq_limit(spec_id))[1]
            siblings = cpuinfo.get_cpu_siblings(num, sname="core")
        else:
            obj = ucfobj
            dies = obj.get_dies_info()
            dies_map = []
            for pkg in dies:
                for die in dies[pkg]:
                    dies_map += [{pkg: [die]}]
            if num >= len(dies_map):
                raise Error(f"bad ID {num} for uncore: valid range is [0-{len(dies_map)-1}]")
            spec_id = dies_map[num]

            fmin = next(obj.get_min_freq_limit_dies(spec_id))[2]
            fmax = next(obj.get_max_freq_limit_dies(spec_id))[2]

        fmin = parse_freq(freq_type, tokens[2], fmin, fmax)
        fmax = parse_freq(freq_type, tokens[3], fmin, fmax)

        if fmin > fmax or fmin == fmax:
            raise Error(f"bad frequency range: [{fmin}-{fmax}]")

        parsed_specs.append({"min": fmin, "max": fmax, "type": freq_type, "id": spec_id,
                             "num": num, "obj": obj, "siblings": siblings})

    return parsed_specs

def init_freq(spec):
    """
    Initialize a frequency specified in 'spec' to a known state. We force the frequency to minimum
    limit provided by the frequency specification.
    """

    f = spec["min"]
    obj = spec["obj"]

    obj.set_min_freq(f, spec["id"])
    obj.set_max_freq(f, spec["id"])
    obj.set_min_freq(f, spec["id"])

    if spec["type"] == "cpu":
        for sibling in spec["siblings"]:
            obj.set_min_freq(f, (sibling,))
            obj.set_max_freq(f, (sibling,))
            obj.set_min_freq(f, (sibling,))

def print_module_paths():
    """
    Print paths to all modules other than standard.
    """

    for mobj in sys.modules.values():
        path = getattr(mobj, "__file__", None)
        if not path:
            continue

        if not path.endswith(".py"):
            continue

        if "pepclibs/" in path or "statscollectlibs/" in path:
            print(path)

def main():
    """
    Implement the base functionality for the tool. Parses arguments and modifies the frequencies
    in a loop by the given command line options.
    """

    args = parse_arguments()
    if args.print_module_paths:
        print_module_paths()
        return 0

    with contextlib.ExitStack() as stack:
        pman = LocalProcessManager.LocalProcessManager()
        stack.enter_context(pman)

        cpuinfo = CPUInfo.CPUInfo(pman=pman)
        stack.enter_context(cpuinfo)

        sysfs_io = _SysfsIO.SysfsIO(pman, enable_cache=False)
        stack.enter_context(sysfs_io)

        ucfobj = _UncoreFreq.UncoreFreqSysfs(cpuinfo, pman=pman, sysfs_io=sysfs_io, enable_cache=False)
        stack.enter_context(ucfobj)

        cpufobj = _CPUFreq.CPUFreqSysfs(cpuinfo=cpuinfo, pman=pman, sysfs_io=sysfs_io,
                                        enable_cache=False, verify=False)
        stack.enter_context(cpufobj)

        ucfobj.set_min_freq = ucfobj.set_min_freq_dies
        ucfobj.set_max_freq = ucfobj.set_max_freq_dies
        ucfobj.get_cur_freq = ucfobj.get_cur_freq_dies

        try:
            if not args.specs:
                raise Error("no frequency specs defined, nothing to do")

            parsed_specs = parse_freq_specs(args, cpuinfo, ucfobj, cpufobj)
            if not parsed_specs:
                return 0

            for spec in parsed_specs:
                init_freq(spec)

            loops = 0

            if args.sleep:
                args.sleep /= 1000000

            _LOG.info("Looping...")
            while True:
                for spec in parsed_specs:
                    set_freq(args, spec, "max")
                if args.sleep:
                    time.sleep(args.sleep)
                for spec in parsed_specs:
                    set_freq(args, spec, "min")
                loops += 1
                if args.sleep:
                    time.sleep(args.sleep)
        except KeyboardInterrupt:
            _LOG.info("Loop terminated after %d loops.", loops)
        except Error as err:
            _LOG.error_out(err)

    return 0

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