#
# Copyright (C) 2018 Marko Myllynen <myllynen@redhat.com>
#
# 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 uprobe hits module """

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

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_COUNT_ONE
from cpmda import PMDA_FETCH_NOVALUES

from modules.pcpbcc import PCPBCCBase

#
# BPF program
#
bpf_src = "modules/uprobe_hits.bpf"
# Individual uprobe
uprobe_txt = """
static char *HASH_KEY = "UPROBE_NAME";
int trace_HASH_KEY(void *ctx) {
    FILTER_PID
    struct uprobe_t key = {};
    __builtin_memcpy(&key.uprobe, HASH_KEY, sizeof(key.uprobe));
    u64 zero = 0, *val;
    val = stats.lookup_or_init(&key, &zero);
    if (val) {
      (*val)++;
    }
    return 0;
}
"""

#
# PCP BCC PMDA constants
#
MODULE = 'uprobe_hits'
METRIC = 'uprobe.hits'
units_count = pmUnits(0, 0, 1, 0, 0, PM_COUNT_ONE)

#
# PCP BCC Module
#
class PCPBCCModule(PCPBCCBase):
    """ PCP BCC uprobe hits 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
        self.uprobes = []

        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 == 'uprobes':
                self.uprobes = self.read_probe_conf(self.config.get(MODULE, opt))

        self.log("Configured uprobes: " + str(self.uprobes))

        found = []
        for uprobe in self.uprobes:
            if uprobe not in found:
                found.append(uprobe)
                self.insts[uprobe.replace(":", "::")] = c_int(1)
        self.uprobes = found

        if not self.uprobes:
            raise RuntimeError("No matching uprobes found.")
        self.log("Found %s uprobes: %s." % (str(len(self.uprobes)), str(self.uprobes)))

        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_count, 'uprobe hit count'),
        )
        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()

                for uprobe in self.uprobes:
                    hash_key = uprobe.replace(":", "__")
                    probe = uprobe_txt.replace("HASH_KEY", hash_key)
                    probe = probe.replace("UPROBE_NAME", uprobe)
                    if self.debug:
                        self.log("Generated function:\n%s" % probe)
                    self.bpf_text += probe
                self.bpf_text = self.bpf_text.replace("UPROBE_COUNT", str(len(self.uprobes)))

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

            self.log("Compiling %s uprobes: %s" % (str(len(self.uprobes)), str((self.uprobes))))

            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)
            for uprobe in self.uprobes:
                lib, sym = uprobe.split(":")
                hash_key = uprobe.replace(":", "__")
                func = "trace_HASH_KEY".replace("HASH_KEY", hash_key)
                if self.debug:
                    self.log("Attaching uprobe (%s:%s): %s" % (lib, sym, func))
                self.bpf.attach_uprobe(lib, sym, fn_name=func)
            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

        for k, v in self.bpf["stats"].items():
            self.cache[k.uprobe.decode("ASCII").replace(":", "::")] = v.value

        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]