#
# Copyright (C) 2017-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 module base class """

import re
import platform
import ctypes as ct
from os import kill, listdir, path
from collections import OrderedDict
from bcc.table import log2_index_max
from bcc import __version__, BPF

from cpmapi import PM_ERR_BADSTORE

# pylint: disable=too-many-instance-attributes, too-many-public-methods
class PCPBCCBase(object):
    """ PCP BCC Base Class """
    def __init__(self, module, config, log, err):
        """ Constructor """
        self._who = module
        self._log = log
        self._err = err

        self.bpf = None
        self.bpf_text = None
        self.insts = {}
        self.items = []

        self.pmdaIndom = None # pylint: disable=invalid-name
        self.config = config
        self.debug = False

        for opt in self.config.options(self._who):
            if opt == 'debug':
                self.debug = self.config.getboolean(self._who, opt)

        if self.debug:
            self.log("Debug logging enabled.")

    def log(self, msg):
        """ Log a message """
        self._log(self._who + ": " + msg)

    def err(self, msg):
        """ Log an error """
        self._err(self._who + ": " + msg)

    def metrics(self):
        """ Get metric definitions """
        raise NotImplementedError

    def helpers(self, pmdaIndom): # pylint: disable=invalid-name
        """ Register helper function references """
        self.pmdaIndom = pmdaIndom

    def compile(self):
        """ Compile BPF """
        raise NotImplementedError

    def refresh(self):
        """ Refresh BPF data """
        raise NotImplementedError

    def bpfdata(self, item, inst):
        """ Return BPF data as PCP metric value """
        raise NotImplementedError

    def label_cluster(self): # pylint: disable=no-self-use
        """ Cluster labels """
        return '{"bcc_module":"' + self._who + '"}'

    def label_indom(self): # pylint: disable=no-self-use
        """ Instance domain labels """
        return '{}'

    def label_instance(self, inst): # pylint: disable=no-self-use,unused-argument
        """ Instance labels """
        return '{}'

    def store(self, item, inst, val): # pylint: disable=no-self-use,unused-argument
        """ Store callback """
        return PM_ERR_BADSTORE

    def free_memory(self):
        """ Free module compilation memory """
        if self.bpf is not None:
            self.bpf.free_bcc_memory()
            self.log("Compilation memory freed.")

    def cleanup(self):
        """ Clean up at exit """
        if self.bpf is not None:
            self.bpf.cleanup()
            self.log("BPF detached.")
            self.undef_cache()
            self.bpf = None

    def reset_cache(self): # pylint: disable=no-self-use
        """ Reset internal cache """
        return

    def undef_cache(self): # pylint: disable=no-self-use
        """ Undefine internal cache """
        return

    #
    # Helpers for modules
    #

    @classmethod
    def read_log2_histogram_section(cls, hist_data, cache):
        """
        Read log2 histogram section

        Adapted from https://github.com/iovisor/bcc/blob/master/src/python/bcc/table.py

        Cache is filled until the top end (zero value buckets).
        Always populate the cache with the first element, even if the value is
        zero (to have at least a single instance instead of no values at all).
        """
        idx_max = -1
        val_max = 0

        for i, v in enumerate(hist_data): # pylint: disable=invalid-name
            if v > 0:
                idx_max = i
            if v > val_max:
                val_max = v

        for i in range(1, max(idx_max, 1) + 1):
            low = (1 << i) >> 1
            high = (1 << i) - 1
            if low == high:
                low -= 1

            key = str(low) + "-" + str(high)
            if key not in cache:
                cache[key] = 0
            cache[key] += hist_data[i]

    # pylint: disable=protected-access
    @classmethod
    def read_log2_histograms(cls, table, cache):
        """
        Read multiple log2 histograms

        Adapted from https://github.com/iovisor/bcc/blob/master/src/python/bcc/table.py
        """

        if not isinstance(table.Key(), ct.Structure):
            raise Exception("Histogram does not contain multiple sections.")

        hist_sections = {}
        hist_key_field1 = table.Key._fields_[0][0]
        hist_key_field2 = table.Key._fields_[1][0]

        # Note: table contains sections with changed values only
        for k, v in table.items(): # pylint: disable=invalid-name
            section = getattr(k, hist_key_field1).decode("UTF-8")
            if section not in hist_sections:
                hist_sections[section] = [0] * log2_index_max
            slot = getattr(k, hist_key_field2)
            hist_sections[section][slot] = v.value

        for section, hist_data in hist_sections.items():
            if section not in cache:
                cache[section] = {}
            cls.read_log2_histogram_section(hist_data, cache[section])

        all_cache_keys = set()
        for section_cache in cache.values():
            all_cache_keys.update(section_cache.keys())
        sorted_cache_keys = sorted(all_cache_keys, key=lambda k: int(k.split('-')[0]))
        return OrderedDict([(key, ct.c_int(1)) for key in sorted_cache_keys])

    @classmethod
    def read_log2_histogram(cls, table, cache):
        """ Read single log2 histogram """
        if isinstance(table.Key(), ct.Structure):
            raise Exception("Histogram contains multiple sections.")

        hist_data = [0] * log2_index_max
        for k, v in table.items(): # pylint: disable=invalid-name
            hist_data[k.value] = v.value
        cls.read_log2_histogram_section(hist_data, cache)
        sorted_cache_keys = sorted(cache.keys(), key=lambda k: int(k.split('-')[0]))
        return OrderedDict([(key, ct.c_int(1)) for key in sorted_cache_keys])

    @staticmethod
    def read_probe_conf(conf):
        """ Read tracepoint/uprobe/USDT probes from a file """
        if ":" in conf:
            return conf.split(",")
        if not conf.startswith("/"):
            conf = path.dirname(__file__) + "/../" + conf
        return [l.strip() for l in open(conf).readlines() if l.strip() != "" and '#' not in l]

    @staticmethod
    def get_proc_info(proc_filter):
        """ Get process info for given process filter """
        if not proc_filter:
            return None

        filters = []
        for filt in proc_filter.split(","):
            if filt.isalnum():
                filters.append(filt)
            else:
                filters.append(re.compile(r'\A' + filt + r'\Z'))

        procinfo = []
        retype = type(re.compile('test'))
        for dirname in listdir("/proc"):
            if not dirname.isdigit():
                continue
            try:
                pid = int(dirname)
                with open("/proc/%s/cmdline" % dirname, "rb") as fcmd:
                    cmdline = fcmd.read().decode().split("\x00")
                with open("/proc/%s/comm" % dirname) as fcomm:
                    comm = fcomm.read().strip()
                cmdline = cmdline if cmdline[0] else ["(" + comm + ")", ""]
                cmdline_full = " ".join(cmdline[:-1])
                cmdline_args = " ".join(cmdline[1:-1])
                for filt in filters:
                    if filt == dirname or filt == comm or \
                       (isinstance(filt, retype) and re.match(filt, cmdline_full)):
                        procinfo.append([pid, cmdline[0], cmdline_args])
            except Exception: # pylint: disable=broad-except
                continue

        return procinfo

    # pylint: disable=access-member-before-definition, attribute-defined-outside-init, no-member
    def update_pids(self, procs):
        """ Update PIDs to trace """
        prev = self.pid if hasattr(self, 'pid') else self.pids
        if not self.proc_filter:
            return 0
        if not prev and not procs:
            return 0
        if not procs and self.proc_filter:
            if hasattr(self, 'pid'):
                self.pid = None
            else:
                self.pids = []
            return -1
        if hasattr(self, 'pid'):
            if self.pid not in [p[0] for p in procs]:
                if len(procs) > 1:
                    self.log("Several PIDs found, tracing the first one.")
                info = procs[0][1] + " " + procs[0][2] if procs[0][2] else procs[0][1]
                self.log("Tracing PID %s: %s." % (str(procs[0][0]), info))
                self.pid = procs[0][0]
                return 1
        else:
            if [p[0] for p in procs if p[0] not in self.pids]:
                self.pids = []
                for proc in procs:
                    info = proc[1] + " " + proc[2] if proc[2] else proc[1]
                    self.log("Tracing PID %s: %s." % (str(proc[0]), info))
                    self.pids.append(proc[0])
                return 1
        return 0

    @staticmethod
    def apply_pid_filter(bpf_text, pids, need_pid_def=True, need_shift=False):
        """ Apply PID filtering for BPF text """
        if pids and need_pid_def:
            pid_filter = "u32 pid = bpf_get_current_pid_tgid() >> 32; if (PID_CHECK) return 0;"
        else:
            pid_filter = "if (PID_CHECK) return 0;"
        # Try to cover all known variants
        if "if (FILTER_PID)" in bpf_text:
            bpf_text = bpf_text.replace("FILTER_PID", "PID_CHECK")
        bpf_text = bpf_text.replace("//FILTER_PID", pid_filter)
        bpf_text = bpf_text.replace("FILTER_PID", pid_filter)
        bpf_text = bpf_text.replace("PID_FILTER", pid_filter)
        if not pids:
            return bpf_text.replace("PID_CHECK", "0")
        pid_check = "pid >> 32 != %d" if need_shift else "pid != %d"
        check_str = " && ".join([pid_check % pid for pid in pids])
        return bpf_text.replace("PID_CHECK", check_str)

    @staticmethod
    def pid_alive(pid):
        """ Test liveliness of PID """
        try:
            kill(int(pid), 0)
            return True
        except Exception: # pylint: disable=broad-except
            return False

    @staticmethod
    def parse_inst_name(string):
        """ Parse string into valid instance name """
        name = ""
        started = 0
        ok_chars = ['.', '/', '<', '>', '[', ']', '_']
        for char in string:
            if not started and not char.isalnum():
                continue
            if not char.isalnum() and char not in ok_chars:
                break
            name += char
            started = 1
        return name.replace("/", ".").replace("[", "{").replace("]", "}")

    @staticmethod
    def get_instance_name_for_pid(pid):
        """
        Returns an instance name for a given process
        in the same format as the proc PMDA, for example:

            001379 /usr/sbin/NetworkManager
        """
        process_name = ""
        try:
            cmdline = open("/proc/{}/cmdline".format(pid), "r").read()
            process_name = " " + re.split("[\x00 ]", cmdline)[0]
        except IOError:
            pass
        return "{:06d}{}".format(pid, process_name)

    @staticmethod
    def bcc_version():
        """ Check BCC version """
        try:
            return __version__.split("-", maxsplit=1)[0]
        except Exception: # pylint: disable=broad-except
            if 'perf_buffer_poll' in dir(BPF):
                return "0.6.0"
            else:
                return "0.5.0"

    @staticmethod
    def bcc_version_tuple():
        """ Returns BCC version as an int tuple (for comparisons) """
        return tuple(map(int, PCPBCCBase.bcc_version().split('.')))

    @staticmethod
    def kernel_version():
        """Returns the kernel version"""
        version_str = platform.release()
        m = re.match(r'^(\d+)\.(\d+)\.(\d+)', version_str)
        if m:
            return tuple(map(int, m.groups()))
        else:
            return (0, 0, 0)

    def perf_buffer_poller(self):
        """ BPF poller """
        try:
            # pylint: disable=no-member
            if self.bcc_version() == "0.5.0":
               # Compat: bcc < 0.6.0
                poll = self.bpf.kprobe_poll
            else:
                poll = self.bpf.perf_buffer_poll
            bpf = self.bpf
            while bpf == self.bpf:
                poll()
        except Exception as error: # pylint: disable=broad-except
            self.err(str(error))
            self.err("BPF perf buffer poll failed!")
        self.log("Poller thread exiting.")

    def get_syscall_prefix(self):
        """
        Get syscall prefix
        Compat: bcc < 0.6.0
        source: https://github.com/iovisor/bcc/blame/master/src/python/bcc/__init__.py
        """
        # Test BPF syscall kernel function name
        if self.bpf.ksymname("sys_bpf") != -1:
            return "sys_"
        if self.bpf.ksymname("__x64_sys_bpf") != -1:
            return "__x64_sys_"
        if self.bpf.ksymname("__x32_compat_sys_bpf") != -1:
            return "__x32_compat_sys_"
        if self.bpf.ksymname("__ia32_compat_sys_bpf") != -1:
            return "__ia32_compat_sys_"
        # None of them, just return "sys_",
        # later API calls will return error
        return "sys_"

    def get_syscall_fnname(self, name):
        """
        Get syscall function name
        Compat: bcc < 0.6.0
        source: https://github.com/iovisor/bcc/blame/master/src/python/bcc/__init__.py
        """
        if hasattr(self.bpf, 'get_syscall_fnname'):
            return self.bpf.get_syscall_fnname(name)
        else:
            return self.get_syscall_prefix() + name

    def get_kprobe_functions(self, event_re):
        """
        Get list of kprobe functions matching a regex
        event_re: regex pattern (bytes object)
        Compat: bcc 0.5.0 expects a string pattern, bcc 0.6.0+ a bytes pattern
        """
        try:
            return self.bpf.get_kprobe_functions(event_re)
        except TypeError:
            return self.bpf.get_kprobe_functions(event_re.decode())