#
# Copyright (C) 2017-2018 Marko Myllynen <myllynen@redhat.com>
# Based on the biotop BCC tool by Brendan Gregg:
# https://github.com/iovisor/bcc/blob/master/tools/biotop.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 bioperpid module """

# pylint: disable=invalid-name, line-too-long

from ctypes import c_int

from bcc import BPF

from pcp.pmapi import pmUnits
from cpmapi import PM_TYPE_U64, PM_SEM_COUNTER, PM_SPACE_BYTE
from cpmda import PMDA_FETCH_NOVALUES

from modules.pcpbcc import PCPBCCBase

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

#
# PCP BCC PMDA constants
#
MODULE = 'bioperpid'
METRIC = 'proc.io.total'
units_bytes = pmUnits(1, 0, 0, PM_SPACE_BYTE, 0, 0)

#
# PCP BCC Module
#
class PCPBCCModule(PCPBCCBase):
    """ PCP BCC bioperpid module """
    def __init__(self, config, log, err, _):
        """ Constructor """
        PCPBCCBase.__init__(self, MODULE, config, log, err)

        self.cache = {}
        self.stale_pids = []

        self.disklookup = None
        self.update_disk_info()

        self.log("Initialized.")

    def update_disk_info(self):
        """ Update disk info cache """
        # Cache disk major,minor -> diskname
        self.disklookup = {}
        if self.debug:
            self.log("Updating disk cache...")
        with open("/proc/diskstats") as stats:
            for line in stats:
                a = line.split()
                self.disklookup[a[0] + "," + a[1]] = a[2]

    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_bytes, 'device io per pid'),
        )
        return True, self.items

    def compile(self):
        """ Compile BPF """
        try:
            self.bpf = BPF(src_file=bpf_src)
            self.bpf.attach_kprobe(event="blk_account_io_start", fn_name="trace_pid_start")
            # Compat: bcc < 0.6.0 (first check)
            if 'get_kprobe_functions' not in dir(self.bpf) or \
               self.get_kprobe_functions(b"blk_start_request"):
                self.bpf.attach_kprobe(event="blk_start_request", fn_name="trace_req_start")
            self.bpf.attach_kprobe(event="blk_mq_start_request", fn_name="trace_req_start")
            self.bpf.attach_kprobe(event="blk_account_io_done", fn_name="trace_req_done")
            self.log("Compiled.")
        except Exception as error: # pylint: disable=broad-except
            self.bpf = None
            self.err(str(error))
            self.err("Module NOT active!")
            raise

    def refresh(self):
        """ Refresh BPF data """
        if self.bpf is None:
            return None

        # Clean stale data
        for key in list(self.stale_pids):
            del self.cache[key]
            del self.insts[key]
            self.stale_pids.remove(key)
        for key in self.cache:
            if not self.pid_alive(key.split("::")[1]):
                self.stale_pids.append(key)

        counts = self.bpf["counts"]

        # Update current data
        for k, v in counts.items():
            disk = str(k.major) + "," + str(k.minor)

            # Unnamed devices (e.g. non-device mounts)
            if k.major == 0:
                continue
            elif disk not in self.disklookup:
                # Check for hot swapped devices
                self.update_disk_info()
                if disk not in self.disklookup:
                    self.log("Traced unknown device (major: %s minor: %s)" % (k.major, k.minor))
                    continue

            key = self.disklookup[disk] + "::" + str(k.pid).zfill(6)
            value = v.bytes if key not in self.cache else v.bytes + self.cache[key]
            self.cache[key] = value
            self.insts[key] = c_int(1)

        counts.clear()

        return self.insts

    def bpfdata(self, item, inst):
        """ Return BPF data as PCP metric value """
        try:
            key = self.pmdaIndom.inst_name_lookup(inst).zfill(6)
            return [self.cache[key], 1]
        except Exception: # pylint: disable=broad-except
            return [PMDA_FETCH_NOVALUES, 0]