#
# Copyright (C) 2019 Marko Myllynen <myllynen@redhat.com>
# Based on the klockstat BCC tool by David Valin and Jiri Olsa:
# https://github.com/iovisor/bcc/blob/master/tools/klockstat.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 klockstat module """

# pylint: disable=invalid-name, too-many-instance-attributes, too-many-return-statements

from ctypes import c_int
from os import path

from bcc import BPF

from pcp.pmapi import pmUnits
from cpmapi import PM_TYPE_U64, PM_SEM_COUNTER, PM_SEM_INSTANT, PM_COUNT_ONE, PM_TIME_NSEC
from cpmapi import PM_ERR_PMID
from cpmda import PMDA_FETCH_NOVALUES

from modules.pcpbcc import PCPBCCBase

#
# BPF program
#
bpf_src = "modules/klockstat.bpf"

#
# PCP BCC PMDA constants
#
MODULE = 'klockstat'
BASENS = 'kernel.lock.mutex.'
units_count = pmUnits(0, 0, 1, 0, 0, PM_COUNT_ONE)
units_nsecs = pmUnits(0, 1, 0, 0, PM_TIME_NSEC, 0)

#
# PCP BCC Module
#
class PCPBCCModule(PCPBCCBase):
    """ PCP BCC klockstat 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.caller_filter = None
        self.stack_depth = 1
        self.stack_storage_size = 16384
        self.use_caller_offset = False

        self.stack_traces = None
        self.aq_counts = None
        self.aq_maxs = None
        self.aq_totals = None
        self.hl_counts = None
        self.hl_maxs = None
        self.hl_totals = 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))
            if opt == 'caller_filter':
                self.caller_filter = self.config.get(MODULE, opt).encode()
            if opt == 'stack_depth':
                self.stack_depth = int(self.config.get(MODULE, opt))
            if opt == 'stack_storage_size':
                self.stack_storage_size = int(self.config.get(MODULE, opt))
            if opt == 'use_caller_offset':
                self.use_caller_offset = self.config.getboolean(MODULE, opt)

        if self.stack_depth < 1:
            raise RuntimeError("stack_depth must be greater than zero.")
        self.stack_depth += 1

        self.cache = None
        self.insts = None

        self.log("Initialized.")

    def metrics(self):
        """ Get metric definitions """
        name = BASENS
        self.items = (
            # Name - reserved - type - semantics - units - help
            (name + 'spin.avg', None, PM_TYPE_U64, PM_SEM_INSTANT, units_nsecs, 'spin avg'),
            (name + 'spin.count', None, PM_TYPE_U64, PM_SEM_COUNTER, units_count, 'spin count'),
            (name + 'spin.max', None, PM_TYPE_U64, PM_SEM_INSTANT, units_nsecs, 'spin max'),
            (name + 'spin.total', None, PM_TYPE_U64, PM_SEM_COUNTER, units_nsecs, 'spin total'),
            (name + 'hold.avg', None, PM_TYPE_U64, PM_SEM_INSTANT, units_nsecs, 'hold avg'),
            (name + 'hold.count', None, PM_TYPE_U64, PM_SEM_COUNTER, units_count, 'hold count'),
            (name + 'hold.max', None, PM_TYPE_U64, PM_SEM_INSTANT, units_nsecs, 'hold max'),
            (name + 'hold.total', None, PM_TYPE_U64, PM_SEM_COUNTER, units_nsecs, 'hold total'),
        )
        return True, self.items

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

    def undef_cache(self):
        """ Undefine internal cache """
        self.cache = None
        self.insts = 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()

                self.bpf_text = self.bpf_text.replace("STACK_STORAGE_SIZE",
                                                      str(self.stack_storage_size))

            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.bpf_text
            bpf_text = self.apply_pid_filter(bpf_text, self.pids, False, True)

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

            self.reset_cache()
            self.bpf = BPF(text=bpf_text)
            self.bpf.attach_kprobe(event="mutex_unlock", fn_name="mutex_unlock_enter")
            self.bpf.attach_kretprobe(event="mutex_lock", fn_name="mutex_lock_return")
            self.bpf.attach_kprobe(event="mutex_lock", fn_name="mutex_lock_enter")
            self.stack_traces = self.bpf["stack_traces"]
            self.aq_counts = self.bpf["aq_report_count"]
            self.aq_maxs = self.bpf["aq_report_max"]
            self.aq_totals = self.bpf["aq_report_total"]
            self.hl_counts = self.bpf["hl_report_count"]
            self.hl_maxs = self.bpf["hl_report_max"]
            self.hl_totals = self.bpf["hl_report_total"]
            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

        def get_fnname(name):
            """ Get function name as string """
            try:
                return name.decode()
            except AttributeError:
                return name

        for key in self.aq_counts.keys():
            stack = []
            caller = "<Missed Kernel Stack>"
            if key.value > 0:
                stack = list(self.stack_traces.walk(key.value))
                caller = self.bpf.ksym(stack[1], show_offset=self.use_caller_offset)
                if self.caller_filter and caller.find(self.caller_filter):
                    continue
                caller = get_fnname(caller)
                for addr in stack[2:self.stack_depth]:
                    caller += "+"
                    caller += get_fnname(self.bpf.ksym(addr, show_offset=self.use_caller_offset))
            self.cache[caller] = key
            self.insts[caller] = c_int(1)

        return self.insts

    def bpfdata(self, item, inst):
        """ Return BPF data as PCP metric value """
        try:
            key = self.cache[self.pmdaIndom.inst_name_lookup(inst)]
            if item == 0:
                return [int(self.aq_totals[key].value / self.aq_counts[key].value), 1]
            elif item == 1:
                return [self.aq_counts[key].value, 1]
            elif item == 2:
                return [self.aq_maxs[key].value, 1]
            elif item == 3:
                return [self.aq_totals[key].value, 1]
            elif item == 4:
                return [int(self.hl_totals[key].value / self.hl_counts[key].value), 1]
            elif item == 5:
                return [self.hl_counts[key].value, 1]
            elif item == 6:
                return [self.hl_maxs[key].value, 1]
            elif item == 7:
                return [self.hl_totals[key].value, 1]
            else:
                return [PM_ERR_PMID, 0]
        except Exception: # pylint: disable=broad-except
            return [PMDA_FETCH_NOVALUES, 0]