#
# Copyright (C) 2017-2018 Marko Myllynen <myllynen@redhat.com>
# Copyright (C) 2018 Andreas Gerstmayr <andreas@gerstmayr.me>
# 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 biotop module """

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

import ctypes as ct
from threading import Lock, Thread
from time import sleep

from bcc import BPF

from pcp.pmapi import pmUnits
from cpmapi import PM_TYPE_U32, PM_TYPE_U64, PM_TYPE_STRING
from cpmapi import PM_SEM_INSTANT, PM_SPACE_BYTE, PM_TIME_USEC
from cpmda import PMDA_FETCH_NOVALUES

from modules.pcpbcc import PCPBCCBase

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

#
# PCP BCC PMDA constants
#
MODULE = 'biotop'
METRIC = 'proc.io.perdev.'
units_bytes = pmUnits(1, 0, 0, PM_SPACE_BYTE, 0, 0)
units_usecs = pmUnits(0, 1, 0, 0, PM_TIME_USEC, 0)
units_none = pmUnits(0, 0, 0, 0, 0, 0)

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

        self.sort_fns = {
            'bytes': lambda counts: counts[1].bytes,
            'io': lambda counts: counts[1].io,
            'duration': lambda counts: float(counts[1].us) / counts[1].io,
            'rw': lambda counts: counts[1].rwflag,
        }

        self.interval = 1
        self.include_residual = False
        self.process_count = 20
        sort_key = '-bytes'

        for opt in self.config.options(MODULE):
            if opt == 'interval':
                self.interval = int(self.config.get(MODULE, opt))
            if opt == 'include_residual':
                self.include_residual = self.config.getboolean(MODULE, opt)
            if opt == 'process_count':
                self.process_count = int(self.config.get(MODULE, opt))
            if opt == 'sort':
                sort_key = self.config.get(MODULE, opt)

        self.sort_reversed = False
        if sort_key[0] == "-":
            sort_key = sort_key[1:]
            self.sort_reversed = True
        if sort_key not in self.sort_fns.keys():
            raise RuntimeError("Sort key must be one of %s." % ", ".join(self.sort_fns.keys()))
        self.sort_fn = self.sort_fns[sort_key]

        self.cache = {}
        self.insts = {str(i) : ct.c_int(1) for i in range(self.process_count)}

        self.disklookup = None
        self.update_disk_info()

        self.lock = Lock()
        if self.interval > 0:
            self.thread = Thread(name="biotop_refresh", target=self.refresh_stats_loop)
            self.thread.daemon = True

        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 = (
            # Name - reserved - type - semantics - units - help
            (name + 'pid', None, PM_TYPE_U32, PM_SEM_INSTANT, units_none, 'PID'),
            (name + 'comm', None, PM_TYPE_STRING, PM_SEM_INSTANT, units_none, 'command'),
            (name + 'direction', None, PM_TYPE_STRING, PM_SEM_INSTANT, units_none, 'IO direction'),
            (name + 'major', None, PM_TYPE_U32, PM_SEM_INSTANT, units_none, 'disk major'),
            (name + 'minor', None, PM_TYPE_U32, PM_SEM_INSTANT, units_none, 'disk minor'),
            (name + 'disk', None, PM_TYPE_STRING, PM_SEM_INSTANT, units_none, 'disk name'),
            (name + 'io', None, PM_TYPE_U32, PM_SEM_INSTANT, units_none, 'IO count'),
            (name + 'bytes', None, PM_TYPE_U64, PM_SEM_INSTANT, units_bytes, 'number of bytes'),
            (name + 'duration', None, PM_TYPE_U32, PM_SEM_INSTANT, units_usecs, 'avg duration'),
        )
        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")
            if self.interval > 0:
                self.thread.start()
            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_stats(self):
        """ Refresh statistics from BPF table """
        self.lock.acquire()
        self.cache = {}
        self.lock.release()

        counts = self.bpf["counts"]
        sorted_counts = sorted(counts.items(), key=self.sort_fn, reverse=self.sort_reversed)

        idx = 0
        for k, v in sorted_counts:
            if not self.include_residual and k.pid == 0:
                continue

            disk = str(k.major) + "," + str(k.minor)

            if disk in self.disklookup:
                diskname = self.disklookup[disk]
            else:
                # Check for hot swapped devices
                self.update_disk_info()
                diskname = self.disklookup.get(disk, "?")

            self.lock.acquire()
            self.cache[idx] = [
                k.pid,
                k.name.decode(),
                "W" if k.rwflag else "R",
                k.major,
                k.minor,
                diskname,
                v.io,
                v.bytes,
                int(float(v.us) / v.io)
            ]
            self.lock.release()

            idx += 1
            if idx >= self.process_count:
                break

        counts.clear()
        return idx

    def refresh_stats_loop(self):
        """ Refresh statistics thread loop """
        while True:
            cnt = self.refresh_stats()
            if self.debug:
                self.log("Fetched data from BPF table, "
                         "current number of processes performing I/O: %d." % cnt)
            sleep(self.interval)

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

        if self.interval == 0:
            self.refresh_stats()

        return self.insts

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