#
# Copyright (C) 2018 Andreas Gerstmayr <andreas@gerstmayr.me>
# Based on the zfsdist BCC tool by Brendan Gregg:
# https://github.com/iovisor/bcc/blob/master/tools/zfsdist.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 zfsdist module """

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

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
#
bpf_src = "modules/fs/zfsdist.bpf"

#
# PCP BCC PMDA constants
#
MODULE = 'zfsdist'
BASENS = 'fs.zfs.latency.'
units_usecs = pmUnits(0, 1, 0, 0, PM_TIME_USEC, 0)

#
# PCP BCC Module
#
class PCPBCCModule(PCPBCCBase):
    """ PCP BCC zfsdist 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.item_identifiers = ['open', 'read', 'write', 'fsync']

        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("Initialized.")

    # pylint: disable=bad-continuation
    def metrics(self):
        """ Get metric definitions """
        for item in self.item_identifiers:
            self.items.append(
                # Name - reserved - type - semantics - units - help
                (BASENS + item, None, PM_TYPE_U64, PM_SEM_COUNTER, units_usecs, 'zfs'
                    '%s latency distribution' % item),
            )
        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()

            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, False)

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

            self.reset_cache()
            self.bpf = BPF(text=bpf_text)
            if self.get_kprobe_functions(b"zpl_iter"):
                self.bpf.attach_kprobe(event="zpl_iter_read", fn_name="trace_entry")
                self.bpf.attach_kprobe(event="zpl_iter_write", fn_name="trace_entry")
            elif self.get_kprobe_functions(b"zpl_aio"):
                self.bpf.attach_kprobe(event="zpl_aio_read", fn_name="trace_entry")
                self.bpf.attach_kprobe(event="zpl_aio_write", fn_name="trace_entry")
            else:
                self.bpf.attach_kprobe(event="zpl_read", fn_name="trace_entry")
                self.bpf.attach_kprobe(event="zpl_write", fn_name="trace_entry")
            self.bpf.attach_kprobe(event="zpl_open", fn_name="trace_entry")
            self.bpf.attach_kprobe(event="zpl_fsync", fn_name="trace_entry")
            if self.get_kprobe_functions(b"zpl_iter"):
                self.bpf.attach_kretprobe(event="zpl_iter_read", fn_name="trace_read_return")
                self.bpf.attach_kretprobe(event="zpl_iter_write", fn_name="trace_write_return")
            elif self.get_kprobe_functions(b"zpl_aio"):
                self.bpf.attach_kretprobe(event="zpl_aio_read", fn_name="trace_read_return")
                self.bpf.attach_kretprobe(event="zpl_aio_write", fn_name="trace_write_return")
            else:
                self.bpf.attach_kretprobe(event="zpl_read", fn_name="trace_read_return")
                self.bpf.attach_kretprobe(event="zpl_write", fn_name="trace_write_return")
            self.bpf.attach_kretprobe(event="zpl_open", fn_name="trace_open_return")
            self.bpf.attach_kretprobe(event="zpl_fsync", fn_name="trace_fsync_return")
            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_histograms(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)
            item_identifier = self.item_identifiers[item]
            if item_identifier in self.cache and key in self.cache[item_identifier]:
                return [self.cache[item_identifier][key], 1]
            elif self.insts:
                # BCC module is ready (we have histogram buckets),
                # but doesn't contain the specified histogram bucket
                return [0, 1]
            else:
                # BCC module isn't ready yet
                return [PMDA_FETCH_NOVALUES, 0]
        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):
        """ Instances labels """
        key = self.pmdaIndom.inst_name_lookup(inst)
        bounds = key.split('-')
        return '{"lower_bound":%s,"upper_bound":%s}' % (bounds[0], bounds[1])