#
# Copyright (C) 2018 Andreas Gerstmayr <andreas@gerstmayr.me>
# Based on the runqlat BCC tool by Brendan Gregg:
# https://github.com/iovisor/bcc/blob/master/tools/runqlat.py
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
""" PCP BCC PMDA runqlat module """

# pylint: disable=invalid-name

from os import path

from bcc import BPF

from pcp.pmapi import pmUnits
from cpmapi import PM_TYPE_U64, PM_SEM_COUNTER, PM_TIME_USEC
from cpmda import PMDA_FETCH_NOVALUES

from modules.pcpbcc import PCPBCCBase

#
# BPF program
#
is_support_raw_tp = BPF.support_raw_tracepoint()
if is_support_raw_tp:
    bpf_src = "modules/runqlat_tp.bpf"
else:
    bpf_src = "modules/runqlat_kp.bpf"

#
# PCP BCC PMDA constants
#
MODULE = 'runqlat'
METRIC = 'runq.latency'
units_usecs = pmUnits(0, 1, 0, 0, PM_TIME_USEC, 0)

#
# PCP BCC Module
#
class PCPBCCModule(PCPBCCBase):
    """ PCP BCC runqlat module """
    def __init__(self, config, log, err, proc_refresh):
        """ Constructor """
        PCPBCCBase.__init__(self, MODULE, config, log, err)

        self.pids = []
        self.proc_filter = None
        self.proc_refresh = proc_refresh

        self.cache = None

        for opt in self.config.options(MODULE):
            if opt == 'process':
                self.proc_filter = self.config.get(MODULE, opt)
                self.update_pids(self.get_proc_info(self.proc_filter))

        self.log("Using BPF source file %s." % bpf_src)
        self.log("Initialized.")

    def metrics(self):
        """ Get metric definitions """
        name = METRIC
        self.items.append(
            # Name - reserved - type - semantics - units - help
            (name, None, PM_TYPE_U64, PM_SEM_COUNTER, units_usecs, 'run queue (scheduler)'
                                                                   'latency distribution'),
        )
        return True, self.items

    def reset_cache(self):
        """ Reset internal cache """
        self.cache = {}

    def undef_cache(self):
        """ Undefine internal cache """
        self.cache = None

    def compile(self):
        """ Compile BPF """
        try:
            if not self.pids and self.proc_filter and not self.proc_refresh:
                raise RuntimeError("No process to attach found.")

            if not self.bpf_text:
                with open(path.dirname(__file__) + '/../' + bpf_src) as src:
                    self.bpf_text = src.read()

                # BPF.kernel_struct_has_field requires BCC v0.23.0
                # use kernel version check as alternative
                # pylint: disable=no-member
                if (
                    hasattr(BPF, "kernel_struct_has_field")
                    and BPF.kernel_struct_has_field(b"task_struct", b"__state") == 1
                ) or self.kernel_version() >= (4, 18, 0):
                    self.bpf_text = self.bpf_text.replace('STATE_FIELD', '__state')
                else:
                    self.bpf_text = self.bpf_text.replace('STATE_FIELD', 'state')

                self.bpf_text = self.bpf_text.replace("FILTER", "PID_CHECK")
                self.bpf_text = self.bpf_text.replace('FACTOR', 'delta /= 1000;')

                self.bpf_text = self.bpf_text.replace('STORAGE', 'BPF_HISTOGRAM(dist);')
                self.bpf_text = self.bpf_text.replace('STORE',
                    'dist.increment(bpf_log2l(delta));')

            if not self.pids and self.proc_filter and self.proc_refresh:
                self.log("No process to attach found, activation postponed.")
                return

            bpf_text = self.apply_pid_filter(self.bpf_text, self.pids)

            if self.debug:
                self.log("BPF to be compiled:\n" + bpf_text.strip())

            self.reset_cache()
            self.bpf = BPF(text=bpf_text)
            if not is_support_raw_tp:
                self.bpf.attach_kprobe(event="ttwu_do_wakeup", fn_name="trace_ttwu_do_wakeup")
                self.bpf.attach_kprobe(event="wake_up_new_task", fn_name="trace_wake_up_new_task")
                self.bpf.attach_kprobe(event_re=r"^finish_task_switch$|^finish_task_switch\.isra\.\d$",
                                       fn_name="trace_run")
            self.log("Compiled.")
        except Exception as error: # pylint: disable=broad-except
            self.bpf = None
            self.undef_cache()
            self.err(str(error))
            self.err("Module NOT active!")
            raise

    def refresh(self):
        """ Refresh BPF data """
        if self.bpf is None:
            return None

        dist = self.bpf["dist"]
        self.insts = self.read_log2_histogram(dist, self.cache)
        dist.clear()

        return self.insts

    def bpfdata(self, item, inst):
        """ Return BPF data as PCP metric value """
        try:
            key = self.pmdaIndom.inst_name_lookup(inst)
            return [self.cache[key], 1]
        except Exception: # pylint: disable=broad-except
            return [PMDA_FETCH_NOVALUES, 0]

    def label_indom(self):
        """ Instance domain labels """
        return '{"statistic":"histogram"}'

    def label_instance(self, inst):
        """ Instance labels """
        key = self.pmdaIndom.inst_name_lookup(inst)
        bounds = key.split('-')
        return '{"lower_bound":%s,"upper_bound":%s}' % (bounds[0], bounds[1])