#!/usr/bin/env python3
#
# -*- coding: utf-8 -*-
# vim: ts=4 sw=4 tw=100 et ai si
#
# Copyright (C) 2023 Intel Corporation
# SPDX-License-Identifier: BSD-3-Clause
#
# Author: Niklas Neronin <niklas.neronin@intel.com>

"""
A tool for finding out the I/O scope of writable MSRs. I/O scope is defined by the observability MSR
value changes, not by the functional impact of the MSR.
"""

import sys
import logging
import contextlib

try:
    import argcomplete
except ImportError:
    # We can live without argcomplete, we only lose tab completions.
    argcomplete = None

from pepctool._Pepc import ArgParse
from pepclibs.helperlibs.Exceptions import Error, ErrorNotSupported
from pepclibs.helperlibs import Logging, ProcessManager, Human
from pepclibs.msr import MSR
from pepclibs import CPUInfo

TOOLNAME = "msrscope"
VERSION = "0.2"
_LOG = logging.getLogger()
Logging.setup_logger(prefix=TOOLNAME)

def _get_possible_ioscopes(cpuinfo, cpu, siblings):
    """Find all scopes shared by all the sibling CPUs."""

    ioscopes = []
    if len(siblings) == 1:
        ioscopes.append("CPU")

    topology = cpuinfo.get_topology()
    tinfo = cpuinfo.get_cpu_levels(cpu)
    for ioscope in ("core", "module", "die", "package"):
        for tline in topology:
            if tline["CPU"] in siblings:
                if tline[ioscope] != tinfo[ioscope]:
                    break
            elif tline[ioscope] == tinfo[ioscope]:
                # The 'core' and 'die' numbers may be relative to the package on some systems.
                if ioscope in {"core", "die"} and tline["package"] != tinfo["package"]:
                    continue
                break
        else:
            ioscopes.append(ioscope)

    return ioscopes

def _find_msr_ioscope(pman, cpuinfo, msr, regaddr, bits, values, cpu):
    """
    Find out MSR I/O scope. The arguments are as follows.
      * pman - the process manager object that defines the host to run the tests on.
      * cpuinfo - CPU information object generated by 'CPUInfo.CPUInfo()'.
      * msr - the 'MSR.MSR()' object to use for writing to the MSR register.
      * regaddr - address of the MSR to read/write the bits.
      * bits - the MSR bits range. A list of 2 decimals: (msb, lsb), where 'msb' is the more
               significant bit, and 'lsb' is a less significant bit.
      * cpu - the CPU which MSR is modified.
    """

    info = {}
    # Collect initial values for all online CPUs.
    for _cpu, val in msr.read_bits(regaddr, bits):
        info[_cpu] = val

    cpus = set()
    # Set CPU 'args.cpu' MSR to a different value and record all CPUs which values have changed.
    for value in values:
        value = int(value)
        if value == info[cpu]:
            continue

        msr.write_bits(regaddr, bits, value, cpus=(cpu, ))
        if value != msr.read_cpu_bits(regaddr, bits, cpu):
            raise Error(f"failed to write value '{value}' to MSR 0x{regaddr:x} bits "
                        f"{':'.join(bits)}")

        for _cpu, val in msr.read_bits(regaddr, bits):
            if info[_cpu] != val:
                cpus.add(_cpu)
        break

    scopes = _get_possible_ioscopes(cpuinfo, cpu, cpus)

    # Restore the MSR to the initial value. It is enough to restore only from 'args.cpu', because
    # other CPUs share the same I/O scope.
    msr.write_bits(regaddr, bits, info[cpu], cpus=(cpu, ))
    val = msr.read_cpu_bits(regaddr, bits, cpu)
    _LOG.debug("initial value was '%d' and value after test is '%d'", info[cpu], val)

    if len(scopes) == 1:
        _LOG.info("I/O scope of MSR %#x bits %s is \"%s\".", regaddr, ":".join(bits), scopes[0])
    elif scopes:
        _LOG.info("Possible I/O scopes of MSR %#x bits %s are: %s.\nThese scopes are identical%s.",
                  regaddr, ":".join(bits), ", ".join(scopes), pman.hostmsg)
    else:
        _LOG.info("Could not determine I/O scope of MSR %#x bits %s. Modified the following "
                  "CPUs:\n%s", regaddr, ":".join(bits), Human.rangify(cpus))

def _test_specific_msr(pman, cpuinfo, msr, args):
    """Test a specific MSR register."""

    # Argument 'regaddr' can either be an decimal or hex.
    try:
        args.regaddr = int(args.regaddr)
    except (ValueError, TypeError):
        try:
            args.regaddr = int(args.regaddr, 16)
        except (ValueError, TypeError):
            raise Error("please supply register addres in decimal or hex") from None

    args.bits = args.bits.split(":")
    if len(args.bits) != 2:
        raise Error("please supply bits in the following format \"h:l\"")

    if args.values[0] == args.values[1]:
        raise Error("please supply 2 unique values")

    cpuinfo.normalize_cpu(args.cpu)
    _find_msr_ioscope(pman, cpuinfo, msr, args.regaddr, args.bits, args.values, args.cpu)

def _find_feature_ioscope(cpuinfo, fmsr, fname, values, cpu):
    """
    Find I/O scope of an MSR feature. The arguments are as follows.
      * cpuinfo - CPU information object generated by 'CPUInfo.CPUInfo()'.
      * fmsr - the featured MSR object, e.g, 'PowerCtl', 'PCStateConfigCtl', 'HWPRequest'
      * fname - the feature name, e.g, 'epp', 'c1e_autopromote'
      * values - tuple of values to set the MSR to.
      * cpu - the CPU which is modified.
    """

    if not fmsr.is_feature_supported(fname, cpus=(cpu, )):
        return

    regaddr = fmsr.regaddr
    bits = fmsr.features[fname]["bits"]
    iosname = fmsr.features[fname]["iosname"]

    info = {}
    for _cpu, val in fmsr.read_feature(fname):
        info[_cpu] = val

    cpus = set()
    # Set CPU 'args.cpu' feature to a different value and record all CPUs which values have changed.
    for value in values:
        if value == info[cpu]:
            continue

        fmsr.write_feature(fname, value, cpus=(cpu, ))
        rval = fmsr.read_cpu_feature(fname, cpu)
        if rval != value and (fname != "pkg_cstate_limit" or value != rval["pkg_cstate_limit"]):
            _LOG.warning("failed to write '%s' value '%s' to MSR %#x %d:%d", fname, str(value),
                         fmsr.regaddr, fmsr.features[fname]["bits"][0],
                         fmsr.features[fname]["bits"][1])
            return

        for _cpu, val in fmsr.read_feature(fname):
            if info[_cpu] != val:
                cpus.add(_cpu)
        break

    scopes = _get_possible_ioscopes(cpuinfo, cpu, cpus)

    # Restore the MSR feature to the initial value. It is enough to restore only from 'args.cpu',
    # because other CPUs share the same I/O scope.
    if fname == "pkg_cstate_limit":
        fmsr.write_feature(fname, info[cpu]["pkg_cstate_limit"], cpus=(cpu, ))
        val = fmsr.read_cpu_feature(fname, cpu)["pkg_cstate_limit"]
    else:
        fmsr.write_feature(fname, info[cpu], cpus=(cpu, ))
        val = fmsr.read_cpu_feature(fname, cpu)
    _LOG.debug("initial value was '%d' and value after test is '%d'", info[cpu], val)

    if not scopes:
        _LOG.info("'%s' (%#x %d:%d) I/O scope couldn't be determined. Modified the following "
                  "CPUs:\n%s", fname, regaddr, bits[0], bits[1], Human.rangify(cpus))
    elif iosname in scopes:
        _LOG.info("'%s' (%#x %d:%d) %s I/O scope is correct.",
                  fname, regaddr, bits[0], bits[1], iosname)
    else:
        _LOG.info("'%s' (%#x %d:%d) %s I/O scope is incorrect, correct I/O scopes are: %s.",
                  fname, regaddr, bits[0], bits[1], iosname, ", ".join(scopes))

def _test_all_msrs(pman, cpuinfo, msr, args):
    """Test MSRs which are used by the 'pepc' tool."""

    # pylint: disable=import-outside-toplevel
    from pepclibs.msr import EnergyPerfBias, HWPRequest, HWPRequestPkg, PowerCtl, PCStateConfigCtl
    from pepclibs.msr import PackagePowerLimit, MiscFeatureControl, PMEnable, UncoreRatioLimit
    from pepclibs.msr import SwLTROvrd

    cpu = cpuinfo.normalize_cpu(args.cpu)

    try:
        fmsr = EnergyPerfBias.EnergyPerfBias(pman=pman, cpuinfo=cpuinfo, msr=msr)
    except ErrorNotSupported as err:
        _LOG.notice(err)
    else:
        _find_feature_ioscope(cpuinfo, fmsr, "epb", (4, 8), cpu)
    finally:
        fmsr.close()

    try:
        fmsr = HWPRequest.HWPRequest(pman=pman, cpuinfo=cpuinfo, msr=msr)
    except ErrorNotSupported as err:
        _LOG.notice(err)
    else:
        _find_feature_ioscope(cpuinfo, fmsr, "epp", (64, 128), cpu)

        with contextlib.suppress(ErrorNotSupported):
            perf = fmsr.read_cpu_feature("min_perf", cpu)
            perf1, perf2 = (perf - 1, perf) if perf > 0 else (perf, perf + 1)
            _find_feature_ioscope(cpuinfo, fmsr, "min_perf", (perf1, perf2), cpu)

        with contextlib.suppress(ErrorNotSupported):
            perf = fmsr.read_cpu_feature("max_perf", cpu)
            perf1, perf2 = (perf - 1, perf) if perf > 254 else (perf, perf + 1)
            _find_feature_ioscope(cpuinfo, fmsr, "max_perf", (perf1, perf2), cpu)
    finally:
        fmsr.close()

    try:
        fmsr = HWPRequestPkg.HWPRequestPkg(pman=pman, cpuinfo=cpuinfo, msr=msr)
    except ErrorNotSupported as err:
        _LOG.notice(err)
    else:
        _find_feature_ioscope(cpuinfo, fmsr, "epp", (64, 128), cpu)

        with contextlib.suppress(ErrorNotSupported):
            perf = fmsr.read_cpu_feature("min_perf", cpu)
            perf1, perf2 = (perf - 1, perf) if perf > 0 else (perf, perf + 1)
            _find_feature_ioscope(cpuinfo, fmsr, "min_perf", (perf1, perf2), cpu)

        with contextlib.suppress(ErrorNotSupported):
            perf = fmsr.read_cpu_feature("max_perf", cpu)
            perf1, perf2 = (perf - 1, perf) if perf > 254 else (perf, perf + 1)
            _find_feature_ioscope(cpuinfo, fmsr, "max_perf", (perf1, perf2), cpu)
    finally:
        fmsr.close()

    try:
        fmsr = PowerCtl.PowerCtl(pman=pman, cpuinfo=cpuinfo, msr=msr)
    except ErrorNotSupported as err:
        _LOG.notice(err)
    else:
        _find_feature_ioscope(cpuinfo, fmsr, "c1e_autopromote", ("off", "on"), cpu)
        _find_feature_ioscope(cpuinfo, fmsr, "cstate_prewake", ("off", "on"), cpu)
        _find_feature_ioscope(cpuinfo, fmsr, "ltr", ("off", "on"), cpu)
    finally:
        fmsr.close()

    try:
        fmsr = PCStateConfigCtl.PCStateConfigCtl(pman=pman, cpuinfo=cpuinfo, msr=msr)
    except ErrorNotSupported as err:
        _LOG.notice(err)
    else:
        _find_feature_ioscope(cpuinfo, fmsr, "c1_demotion", ("off", "on"), cpu)
        _find_feature_ioscope(cpuinfo, fmsr, "c1_undemotion", ("off", "on"), cpu)

        if fmsr.read_cpu_feature("pkg_cstate_limit_lock", cpu) == "off":
            limits = fmsr.read_cpu_feature("pkg_cstate_limit", cpu)["pkg_cstate_limits"]
            _find_feature_ioscope(cpuinfo, fmsr, "pkg_cstate_limit", limits[:2], cpu)
    finally:
        fmsr.close()

    try:
        fmsr = PackagePowerLimit.PackagePowerLimit(pman=pman, cpuinfo=cpuinfo, msr=msr)
    except ErrorNotSupported as err:
        _LOG.notice(err)
    else:
        with contextlib.suppress(Error):
            limit = fmsr.read_cpu_feature("limit1", cpu)
            _find_feature_ioscope(cpuinfo, fmsr, "limit1", (limit - 1, limit + 1), cpu)
            _find_feature_ioscope(cpuinfo, fmsr, "limit1_enable", ("off", "on"), cpu)
            _find_feature_ioscope(cpuinfo, fmsr, "limit1_clamp", ("off", "on"), cpu)

        with contextlib.suppress(Error):
            limit = fmsr.read_cpu_feature("limit2", cpu)
            _find_feature_ioscope(cpuinfo, fmsr, "limit2", (limit - 1, limit + 1), cpu)
            _find_feature_ioscope(cpuinfo, fmsr, "limit2_enable", ("off", "on"), cpu)
            _find_feature_ioscope(cpuinfo, fmsr, "limit2_clamp", ("off", "on"), cpu)
    finally:
        fmsr.close()

    try:
        fmsr = MiscFeatureControl.MiscFeatureControl(pman=pman, cpuinfo=cpuinfo, msr=msr)
    except ErrorNotSupported as err:
        _LOG.notice(err)
    else:
        _find_feature_ioscope(cpuinfo, fmsr, "l2_hw_prefetcher", ("off", "on"), cpu)
        _find_feature_ioscope(cpuinfo, fmsr, "l2_adj_prefetcher", ("off", "on"), cpu)
        _find_feature_ioscope(cpuinfo, fmsr, "dcu_hw_prefetcher", ("off", "on"), cpu)
        _find_feature_ioscope(cpuinfo, fmsr, "dcu_ip_prefetcher", ("off", "on"), cpu)
    finally:
        fmsr.close()

    try:
        fmsr = PMEnable.PMEnable(pman=pman, cpuinfo=cpuinfo, msr=msr)
    except ErrorNotSupported as err:
        _LOG.notice(err)
    else:
        with contextlib.suppress(Error):
            _find_feature_ioscope(cpuinfo, fmsr, "hwp", ("off", "on"), cpu)
    finally:
        fmsr.close()

    try:
        fmsr = UncoreRatioLimit.UncoreRatioLimit(pman=pman, cpuinfo=cpuinfo, msr=msr)
    except ErrorNotSupported as err:
        _LOG.notice(err)
    else:
        with contextlib.suppress(ErrorNotSupported):
            ratio = fmsr.read_cpu_feature("min_ratio", cpu)
            ratio1, ratio2 = ratio, ratio + 1
            _find_feature_ioscope(cpuinfo, fmsr, "min_ratio", (ratio1, ratio2), cpu)

        with contextlib.suppress(ErrorNotSupported):
            ratio = fmsr.read_cpu_feature("max_ratio", cpu)
            ratio1, ratio2 = ratio, ratio + 1
            _find_feature_ioscope(cpuinfo, fmsr, "max_ratio", (ratio1, ratio2), cpu)
    finally:
        fmsr.close()

    try:
        fmsr = SwLTROvrd.SwLTROvrd(pman=pman, cpuinfo=cpuinfo, msr=msr)
    except ErrorNotSupported as err:
        _LOG.notice(err)
    else:
        _find_feature_ioscope(cpuinfo, fmsr, "sxl", (0, 1), cpu)
        _find_feature_ioscope(cpuinfo, fmsr, "sxlm", (0, 1), cpu)
        _find_feature_ioscope(cpuinfo, fmsr, "force_sxl", ("on", "off"), cpu)
        _find_feature_ioscope(cpuinfo, fmsr, "sxl_v", ("on", "off"), cpu)
    finally:
        fmsr.close()

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

    text = f"""{TOOLNAME} - find out the I/O scope of writable MSRs. I/O scope is defined by the
               observability MSR value changes, not by the functional impact of the MSR."""
    parser = ArgParse.SSHOptsAwareArgsParser(description=text, prog=TOOLNAME, ver=VERSION)
    ArgParse.add_ssh_options(parser)

    text = """MSR address or "all" to run tests for all writable MSRs supported by the 'pepc'
              project."""
    parser.add_argument("regaddr", help=text)

    text = """MSR bits to use for finding out the I/O scope, by default \"0:0\"."""
    parser.add_argument("--bits", help=text, default="0:0")
    text = """Two unique values to set the MSR bits specified with '--bits'."""
    parser.add_argument("--values", nargs=2, help=text, default=[0, 1])
    text = """CPU number to write to the MSR bits from, by default CPU 0."""
    parser.add_argument("--cpu", type=int, help=text, default=0)

    if argcomplete:
        argcomplete.autocomplete(parser)

    return parser

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

    parser = _build_arguments_parser()
    args = parser.parse_args()

    return args

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, \
             CPUInfo.CPUInfo(pman=pman) as cpuinfo, \
             MSR.MSR(cpuinfo, pman=pman, enable_cache=False) as msr:

            if cpuinfo.get_offline_cpus_count():
                raise Error("please online all CPUs")

            if args.regaddr == "all":
                _test_all_msrs(pman, cpuinfo, msr, args)
            else:
                _test_specific_msr(pman, cpuinfo, msr, args)

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

    return 0

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