#
# 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 tracepoint 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/tracepoint_hits.bpf"
# Individual tracepoint probe
tp_probe = """
static char *HASH_KEY = "TRACEPOINT_NAME";
TRACEPOINT_PROBE(TRACEPOINT_CATEGORY, TRACEPOINT_EVENT) {
    FILTER_PID
    struct tp_t key = {};
    __builtin_memcpy(&key.tp, HASH_KEY, sizeof(key.tp));
    u64 zero = 0, *val;
    val = stats.lookup_or_init(&key, &zero);
    if (val) {
      (*val)++;
    }
    return 0;
}
"""

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

#
# PCP BCC Module
#
class PCPBCCModule(PCPBCCBase):
    """ PCP BCC tracepoint 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.compile_test = False
        self.tracepoints = []

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

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

        found = []
        for item in self.tracepoints:
            # commit https://github.com/iovisor/bcc/commit/d2a4626 changed
            # signature of BPF.get_tracepoints to require a bytes object
            # instead of a string
            if self.bcc_version_tuple() >= (0, 10, 0):
                tps = BPF.get_tracepoints(item.encode('utf-8') + b"$")
            else:
                tps = BPF.get_tracepoints(item + "$")

            if not tps:
                self.log("Warning: ignoring unknown tracepoint '%s'." % item)
            for tp in tps:
                if tp not in found:
                    found.append(tp)
                    self.insts[tp] = c_int(1)
        self.tracepoints = found

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

        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, 'tracepoint 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 trace found.")

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

                tps = []
                tested_probes = ""
                for tp in self.tracepoints:
                    hash_key = tp.replace(":", "__")
                    tp_category, tp_event = tp.split(":")
                    probe = tp_probe.replace("HASH_KEY", hash_key)
                    probe = probe.replace("TRACEPOINT_NAME", tp)
                    probe = probe.replace("TRACEPOINT_CATEGORY", tp_category)
                    probe = probe.replace("TRACEPOINT_EVENT", tp_event)
                    if self.compile_test:
                        self.log("Compilation test for tracepoint probe: '%s'." % tp)
                        test_probe = self.bpf_text + probe.replace("FILTER_PID", "")
                        test_probe = test_probe.replace("TRACEPOINT_COUNT", "1")
                        try:
                            test_bpf = BPF(text=test_probe)
                            test_bpf.cleanup()
                            tps.append(tp)
                            tested_probes += probe
                        except Exception: # pylint: disable=broad-except
                            self.err("Failed to compile tracepoint probe '%s', disabling." % tp)
                    else:
                        tps.append(tp)
                        self.bpf_text += probe

                    self.bpf_text += tested_probes
                    self.bpf_text = self.bpf_text.replace("TRACEPOINT_COUNT", str(len(tps)))
                    self.tracepoints = tps

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

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

            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)
            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.tp.decode("ASCII")] = 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]