#!/usr/bin/python3 -tt
# -*- coding: utf-8 -*-
# SPDX-License-Identifier: GPL-2.0-or-later
#
#   rteval - script for evaluating platform suitability for RT Linux
#
#           This program is used to determine the suitability of
#           a system for use in a Real Time Linux environment.
#           It starts up various system loads and measures event
#           latency while the loads are running. A report is generated
#           to show the latencies encountered during the run.
#
#   Copyright 2009 - 2013   Clark Williams <williams@redhat.com>
#   Copyright 2009 - 2013   David Sommerseth <davids@redhat.com>
#   Copyright 2012 - 2013   Raphaël Beamonte <raphael.beamonte@gmail.com>
#
""" Main module of the rteval program """

import sys
import os
import time
import re
import shutil
import argparse
import tempfile
import requests
import lxml.etree
from rteval.Log import Log
from rteval import RtEval, rtevalConfig
from rteval.modules.loads import LoadModules
from rteval.modules.measurement import MeasurementModules
from rteval import cpupower
from rteval.version import RTEVAL_VERSION
from rteval.systopology import SysTopology, parse_cpulist_from_config
from rteval.modules.loads.kcompile import ModuleParameters
import rteval.cpulist_utils as cpulist_utils

compress_cpulist = cpulist_utils.compress_cpulist
expand_cpulist = cpulist_utils.expand_cpulist
collapse_cpulist = cpulist_utils.collapse_cpulist

def summarize(repfile, xslt):
    """ Summarize an already existing XML report """
    isarchive = False
    summaryfile = repfile
    if repfile.endswith(".tar.bz2"):
        import tarfile
        try:
            t = tarfile.open(repfile)
        except:
            print(f"Don't know how to summarize {repfile} (tarfile open failed)")
            return
        element = None
        for f in t.getnames():
            if f.find('summary.xml') != -1:
                element = f
                break
        if element is None:
            print(f"No summary.xml found in tar archive {repfile}")
            return
        tmp = tempfile.gettempdir()
        t.extract(element, path=tmp)
        summaryfile = os.path.join(tmp, element)
        isarchive = True

    # Load the XSLT template
    with open(xslt, "r") as xsltfp:
        xsltdoc = lxml.etree.parse(xsltfp)
        xsltprs = lxml.etree.XSLT(xsltdoc)

    # Load the summary.xml report - with some simple sanity checks
    with open(summaryfile, "r") as xmlfp:
        xmldoc = lxml.etree.parse(xmlfp)

    if xmldoc.docinfo.root_name != 'rteval':
        raise RuntimeError("The report doesn't seem like a rteval summary report")

    # Parse and print the report through the XSLT template - preserve proper encoding
    resdoc = xsltprs(xmldoc)
    print(str(resdoc))

    # Clean up
    del resdoc
    del xmldoc
    del xsltprs
    del xsltdoc

    if isarchive:
        os.unlink(summaryfile)



def parse_options(cfg, parser, cmdargs):
    '''parse the command line arguments'''

    rtevcfg = cfg.GetSection('rteval')
    #
    # All the destination variables here should go into the 'rteval' section,
    # thus they are prefixed with 'rteval___'.
    # See rteval/rtevalConfig::UpdateFromOptionParser() method for more info
    #
    parser.add_argument("-d", "--duration", dest="rteval___duration",
                      type=str, default=rtevcfg.duration, metavar="DURATION",
                      help=f"specify length of test run (default: {rtevcfg.duration})")
    parser.add_argument("-v", "--verbose", dest="rteval___verbose",
                      action="store_true", default=rtevcfg.verbose,
                      help=f"turn on verbose prints (default: {rtevcfg.verbose})")
    parser.add_argument("-q", "--quiet", dest="rteval___quiet",
                      action="store_true", default=rtevcfg.quiet,
                      help=f"turn on quiet mode (default: {rtevcfg.quiet})")
    parser.add_argument("-w", "--workdir", dest="rteval___workdir",
                      type=str, default=rtevcfg.workdir, metavar="DIRECTORY",
                      help=f"top directory for rteval data (default: {rtevcfg.workdir})")
    parser.add_argument("-l", "--loaddir", dest="rteval___srcdir",
                      type=str, default=rtevcfg.srcdir, metavar="DIRECTORY",
                      help=f"directory for load source tarballs (default: {rtevcfg.srcdir})")
    parser.add_argument("-i", "--installdir", dest="rteval___installdir",
                      type=str, default=rtevcfg.installdir, metavar="DIRECTORY",
                      help=f"place to locate installed templates (default: {rtevcfg.installdir})")
    parser.add_argument("-s", "--sysreport", dest="rteval___sysreport",
                      action="store_true", default=rtevcfg.sysreport,
                      help=f'run sysreport to collect system data (default: {rtevcfg.sysreport})')
    parser.add_argument("-D", '--debug', dest='rteval___debugging',
                      action='store_true', default=rtevcfg.debugging,
                      help=f'turn on debug prints (default: {rtevcfg.debugging})')
    parser.add_argument("-Z", '--summarize', dest='rteval___summarize',
                      nargs='+', default=None, metavar='XMLFILE',
                      help='summarize one or more already existing XML reports')
    parser.add_argument("-H", '--raw-histogram', dest='rteval___rawhistogram',
                      nargs='+', default=None, metavar='XMLFILE',
                      help='Generate raw histogram data for one or more already existing XML reports')
    parser.add_argument("-f", "--inifile", dest="rteval___inifile",
                      type=str, default=None, metavar="FILE",
                      help="initialization file for configuring loads and behavior")
    parser.add_argument("-a", "--annotate", dest="rteval___annotate",
                      type=str, default=None, metavar="STRING",
                      help="Add a little annotation which is stored in the report")
    parser.add_argument("-L", "--logging", dest="rteval___logging",
                      action='store_true', default=False,
                      help='log the output of the loads in the report directory')
    parser.add_argument("-O", "--onlyload", dest="rteval___onlyload",
                      action='store_true', default=False,
                      help="only run the loads (don't run measurement threads)")
    parser.add_argument("-V", "--version", dest="rteval___version",
                      action='store_true', default=False,
                      help='print rteval version and exit')
    parser.add_argument("-S", "--source-download", nargs="?", dest="rteval___srcdownload",
                        type=str, default=None, const='', metavar="KERNEL_VERSION",
                        help='download a source kernel from kernel.org and exit')
    parser.add_argument("--noload", dest="rteval___noload",
                        action="store_true", default=False,
                        help="only run the measurements (don't run loads)")

    cmd_opts = parser.parse_args(args=cmdargs)

    # if no kernel version was provided for --source-download, set version to default
    if cmd_opts.rteval___srcdownload is not None:
        if cmd_opts.rteval___srcdownload == '':
            cmd_opts.rteval___srcdownload = ModuleParameters()["source"]["default"].replace(".tar.xz", "")

    if cmd_opts.rteval___version:
        print(f"rteval version {RTEVAL_VERSION}")
        sys.exit(0)

    if cmd_opts.rteval___duration:
        mult = 1.0
        v = cmd_opts.rteval___duration.lower()
        if v.endswith('s'):
            v = v[:-1]
        elif v.endswith('m'):
            v = v[:-1]
            mult = 60.0
        elif v.endswith('h'):
            v = v[:-1]
            mult = 3600.0
        elif v.endswith('d'):
            v = v[:-1]
            mult = 3600.0 * 24.0
        cmd_opts.rteval___duration = float(v) * mult

    # Update the config object with the parsed arguments
    cfg.UpdateFromOptionParser(cmd_opts)

    return cmd_opts

if __name__ == '__main__':
    from rteval.sysinfo import dmi

    # set LD_BIND_NOW to resolve shared library symbols
    # note: any string will do, nothing significant about 'rteval'

    os.environ['LD_BIND_NOW'] = 'rteval'

    try:
        # Prepare logging
        logger = Log()
        logger.SetLogVerbosity(Log.NONE)

        # setup initial configuration
        config = rtevalConfig.rtevalConfig(logger=logger)

        # Before really parsing options, see if we have been given a config file in the args
        # and load it - just so that default values are according to the config file
        # Use a minimal parser to extract just the -f option
        config_parser = argparse.ArgumentParser(add_help=False)
        config_parser.add_argument('-f', '--inifile', dest='inifile')
        early_args, _ = config_parser.parse_known_args()

        if early_args.inifile:
            config.Load(early_args.inifile)
        else:
            # No configuration file given, load defaults
            config.Load()

        if not config.HasSection('loads'):
            config.AppendConfig('loads', {
                'kcompile'   : 'module',
                'hackbench'  : 'module',
                'stressng'   : 'module'})

        if not config.HasSection('measurement'):
            config.AppendConfig('measurement', {
                'timerlat' : 'module',
                'sysstat' : 'module'})

        # Check for --measurement-module argument early to override config file
        measurement_parser = argparse.ArgumentParser(add_help=False)
        measurement_parser.add_argument('--measurement-module', dest='measurement_module',
                                        type=str, choices=['cyclictest', 'timerlat'])
        early_args, _ = measurement_parser.parse_known_args()

        # For help generation, we want to load both cyclictest and timerlat modules
        # so users can see all available options regardless of which one is selected
        msrcfg = config.GetSection('measurement')

        # Determine which module should actually run before we load both for help
        selected_measurement_module = None
        if early_args.measurement_module:
            # Command-line argument takes precedence
            selected_measurement_module = early_args.measurement_module
            logger.log(Log.INFO, f"Using measurement module: {selected_measurement_module} (from command line)")
        else:
            # Check config file to see what's configured
            cyclictest_in_config = hasattr(msrcfg, 'cyclictest') and msrcfg.cyclictest == 'module'
            timerlat_in_config = hasattr(msrcfg, 'timerlat') and msrcfg.timerlat == 'module'

            if cyclictest_in_config and not timerlat_in_config:
                selected_measurement_module = 'cyclictest'
            elif timerlat_in_config and not cyclictest_in_config:
                selected_measurement_module = 'timerlat'
            elif cyclictest_in_config and timerlat_in_config:
                # Both in config - prefer timerlat as default
                selected_measurement_module = 'timerlat'
            else:
                # Neither in config - default to timerlat
                selected_measurement_module = 'timerlat'

        # Now ensure both modules are loaded for option display
        msrcfg.cyclictest = 'module'
        msrcfg.timerlat = 'module'

        # Prepare log levels before loading modules, not to have unwanted log messages
        # Use a minimal parser to extract logging-related flags early
        rtevcfg = config.GetSection('rteval')
        log_parser = argparse.ArgumentParser(add_help=False)
        log_parser.add_argument('-v', '--verbose', action='store_true')
        log_parser.add_argument('-D', '--debug', action='store_true')
        log_parser.add_argument('-q', '--quiet', action='store_true')
        log_parser.add_argument('--idle-set', action='store_true')
        early_args, _ = log_parser.parse_known_args()

        if early_args.verbose:
            rtevcfg.verbose = True
        if early_args.debug:
            rtevcfg.debugging = True
        if early_args.quiet:
            rtevcfg.quiet = True

        loglev = (not rtevcfg.quiet and (Log.ERR | Log.WARN)) \
                | (rtevcfg.verbose and Log.INFO) \
                | (rtevcfg.debugging and Log.DEBUG)
        logger.SetLogVerbosity(loglev)

        # check if cpupower is being used
        if early_args.idle_set:
            rtevcfg.update({'usingCpupower': True})

        # Load modules (both cyclictest and timerlat are loaded for help display)
        loadmods = LoadModules(config, logger=logger)
        measuremods = MeasurementModules(config, logger=logger)

        # parse command line options
        parser = argparse.ArgumentParser()
        loadmods.SetupModuleOptions(parser)
        measuremods.SetupModuleOptions(parser)
        cmd_opts = parse_options(config, parser, sys.argv[1:])

        # After parsing options, disable the measurement module that won't be used
        # This must be done after option parsing so both modules' options appear in help
        msrcfg = config.GetSection('measurement')
        if hasattr(cmd_opts, 'measurement___measurement_module') and cmd_opts.measurement___measurement_module:
            selected_measurement_module = cmd_opts.measurement___measurement_module

        # Disable the non-selected measurement module
        if selected_measurement_module == 'timerlat':
            if hasattr(msrcfg, 'cyclictest'):
                msrcfg.cyclictest = None
        elif selected_measurement_module == 'cyclictest':
            if hasattr(msrcfg, 'timerlat'):
                msrcfg.timerlat = None

        if rtevcfg.noload:
            if rtevcfg.onlyload:
                # Make up your mind!
                raise RuntimeError('The --noload and --onlyload options are incompatible.')
            loadmods = None

        # download kernel tarball
        if rtevcfg.srcdownload:
            logger.log(Log.DEBUG, f"Kernel Version to download = {rtevcfg.srcdownload}")

            # handle a kernel version like linux-5.19-rc5
            if 'rc' in rtevcfg.srcdownload:
                kernel_prefix = re.search(r"\d{1,2}\.\d{1,3}\-[a-z]*\d{1,2}", rtevcfg.srcdownload).group(0)
                url = "https://git.kernel.org/torvalds/t/"
            else:
                kernel_prefix = re.search(r"(\d{1,2}\.\d{1,3}\.\d{1,3})|(\d{1,2}\.\d{1,3})", rtevcfg.srcdownload).group(0)
                major_version = re.search(r"\d{1,2}", kernel_prefix).group(0)
                url = "https://kernel.org/pub/linux/kernel/v" + major_version + ".x/"

            file_ext = rtevcfg.srcdownload.split(kernel_prefix)[-1]

            if file_ext and file_ext not in ('.tar.xz', '.tar.gz'):
                sys.exit("Invalid file extension for the kernel source. Exiting")

            if rtevcfg.srcdownload.endswith(".gz") or 'rc' in rtevcfg.srcdownload:
                rtevcfg.srcdownload = "linux-" + kernel_prefix + ".tar.gz"
            else:
                rtevcfg.srcdownload = "linux-" + kernel_prefix + ".tar.xz"
            tarfl = os.path.join(rtevcfg.srcdir, rtevcfg.srcdownload)

            # if default kernel packages with rteval-loads exists, do not download/overwrite
            default_kernel_file = ModuleParameters().get('source').get('default')
            if os.path.exists(tarfl):
                if rtevcfg.srcdownload == default_kernel_file:
                    sys.exit(f"Default kernel {default_kernel_file} already exists, will not download")
                prompt = input("Kernel already exists, download and overwrite anyway? (y/n)  ")
                prompt = prompt.lower()
                if prompt in ('no', 'n'):
                    sys.exit("Exiting")
                elif prompt in ('yes','y'):
                    # backup the existing kernel in case it needs to be restored later
                    shutil.move(tarfl, tarfl + ".bkup")
                else:
                    sys.exit("Invalid option. Exiting")

            url = url + rtevcfg.srcdownload
            print(f"Downloading kernel {url}")
            downloaded_file = requests.get(url)
            if downloaded_file.status_code != 200:
                # restore the kernel file if it exists
                if os.path.exists(tarfl + ".bkup"):
                    shutil.move(tarfl + ".bkup", tarfl)
                sys.exit(f"Could not download tar file {rtevcfg.srcdownload}, status code {downloaded_file.status_code}")
            with open(tarfl, 'wb') as fd:
                fd.write(downloaded_file.content)
            logger.log(Log.DEBUG, f"Kernel source {rtevcfg.srcdownload} downloaded successfully")
            logger.log(Log.DEBUG, f"Downloaded to directory location: {rtevcfg.srcdir}")
            # download was successful, delete the backup file if it exists
            if os.path.exists(tarfl + ".bkup"):
                os.remove(tarfl + ".bkup")
            sys.exit(0)


        ldcfg = config.GetSection('loads')
        msrcfg = config.GetSection('measurement')
        # Remember if cpulists were explicitly set by the user before running
        # parse_cpulist_from_config, which generates default value for them
        msrcfg_cpulist_present = msrcfg.cpulist != ""
        ldcfg_cpulist_present = ldcfg.cpulist != ""
        # Parse cpulists using parse_cpulist_from_config to account for
        # run-on-isolcpus and relative cpusets
        cpulist = parse_cpulist_from_config(msrcfg.cpulist, msrcfg.run_on_isolcpus)
        if msrcfg_cpulist_present and not cpulist_utils.is_relative(msrcfg.cpulist) and msrcfg.run_on_isolcpus:
            logger.log(Log.WARN, "ignoring --measurement-run-on-isolcpus, since cpulist is specified")
        msrcfg.cpulist = collapse_cpulist(cpulist)
        cpulist = parse_cpulist_from_config(ldcfg.cpulist)
        ldcfg.cpulist = collapse_cpulist(cpulist)
        # if we only specified one set of cpus (loads or measurement)
        # default the other to the inverse of the specified list
        if not ldcfg_cpulist_present and msrcfg_cpulist_present:
            tmplist = expand_cpulist(msrcfg.cpulist)
            tmplist = SysTopology().invert_cpulist(tmplist)
            tmplist = cpulist_utils.online_cpulist(tmplist)
            ldcfg.cpulist = collapse_cpulist(tmplist)
        if not msrcfg_cpulist_present and ldcfg_cpulist_present:
            tmplist = expand_cpulist(ldcfg.cpulist)
            tmplist = SysTopology().invert_cpulist(tmplist)
            tmplist = cpulist_utils.online_cpulist(tmplist)
            msrcfg.cpulist = collapse_cpulist(tmplist)

        if ldcfg_cpulist_present:
            logger.log(Log.DEBUG, f"loads cpulist: {ldcfg.cpulist}")
        # if --onlyload is specified msrcfg.cpulist is unused
        if msrcfg_cpulist_present and not rtevcfg.onlyload:
            logger.log(Log.DEBUG, f"measurement cpulist: {msrcfg.cpulist}")
        logger.log(Log.DEBUG, f"workdir: {rtevcfg.workdir}")

        # if --summarize was specified then just parse the XML, print it and exit
        if cmd_opts.rteval___summarize:
            for xmlfile in cmd_opts.rteval___summarize:
                summarize(xmlfile, rtevcfg.xslt_report)
            sys.exit(0)

        # if --raw-histogram was specified then just parse the XML, print it and exit
        if cmd_opts.rteval___rawhistogram:
            for xmlfile in cmd_opts.rteval___rawhistogram:
                summarize(xmlfile, rtevcfg.xslt_histogram)
            sys.exit(0)

        if os.getuid() != 0:
            print("Must be root to run rteval!")
            sys.exit(-1)

        logger.log(Log.DEBUG, f'''rteval options:
     workdir: {rtevcfg.workdir}
     loaddir: {rtevcfg.srcdir}
     reportdir: {rtevcfg.reportdir}
     verbose: {rtevcfg.verbose}
     debugging: {rtevcfg.debugging}
     logging:  {rtevcfg.logging}
     duration: {rtevcfg.duration}
     sysreport: {rtevcfg.sysreport}''')

        if not os.path.isdir(rtevcfg.workdir):
            raise RuntimeError(f"work directory {rtevcfg.workdir} does not exist")

        # if idle-set has been specified, enable the idle state via cpupower
        if msrcfg.idlestate:
            cpupower_controller = cpupower.Cpupower(msrcfg.cpulist, msrcfg.idlestate, logger=logger)
            cpupower_controller.enable_idle_state()

        rteval = RtEval(config, loadmods, measuremods, logger)
        rteval.Prepare(rtevcfg.onlyload)

        if rtevcfg.onlyload:
            # If --onlyload were given, just kick off the loads and nothing more
            # No reports will be created.
            loadmods.Start()
            nthreads = loadmods.Unleash()
            logger.log(Log.INFO, f"Started {nthreads} load threads - will run for {rtevcfg.duration} seconds")
            logger.log(Log.INFO, "No measurements will be performed, due to the --onlyload option")
            time.sleep(rtevcfg.duration)
            loadmods.Stop()
            ec = 0
        else:
            # ... otherwise, run the full measurement suite with loads
            ec = rteval.Measure()
            logger.log(Log.DEBUG, f"exiting with exit code: {ec}")

        # restore previous idle state settings
        if msrcfg.idlestate:
            cpupower_controller.restore_idle_states()

        sys.exit(ec)
    except KeyboardInterrupt:
        sys.exit(0)
