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

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

import errno
import subprocess
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_SEM_INSTANT, PM_COUNT_ONE, PM_TIME_NSEC
from cpmapi import PM_ERR_PMID
from cpmda import PMDA_FETCH_NOVALUES

from modules.pcpbcc import PCPBCCBase

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

#
# PCP BCC PMDA constants
#
MODULE = 'syscount'
BASENS = 'proc.syscall.'
units_count = pmUnits(0, 0, 1, 0, 0, PM_COUNT_ONE)
units_nsecs = pmUnits(0, 1, 0, 0, PM_TIME_NSEC, 0)

#
# PCP BCC Module
#
class PCPBCCModule(PCPBCCBase):
    """ PCP BCC syscount 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.errno = 0
        self.failed = False
        self.details = False
        self.latency = False

        self.max_pids = 128

        def parse_syscall(line):
            """ Helper to parse ausyscall(8) output """
            parts = line.split()
            return (int(parts[0]), parts[1].strip().decode("UTF-8"))

        try:
            out = subprocess.check_output('ausyscall --dump | tail -n +2', shell=True)
            self.syscalls = dict(map(parse_syscall, out.strip().split(b'\n')))
        except Exception: # pylint: disable=broad-except
            raise RuntimeError("ausyscall: command not found")

        def read_errno(errstr):
            """ Helper to read errno """
            try:
                return abs(int(errstr))
            except ValueError:
                return getattr(errno, errstr)

        for opt in self.config.options(MODULE):
            if opt == 'errno':
                self.errno = read_errno(self.config.get(MODULE, opt))
            if opt == 'failed':
                self.failed = self.config.getboolean(MODULE, opt)
            if opt == 'details':
                self.details = self.config.getboolean(MODULE, opt)
            if opt == 'latency':
                self.latency = 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))

        self.cnt_cache = None
        self.avg_cache = None
        self.cml_cache = None
        self.insts = None

        self.log("Initialized.")

    def metrics(self):
        """ Get metric definitions """
        name = BASENS
        self.items = (
            # Name - reserved - type - semantics - units - help
            (name + 'count', None, PM_TYPE_U64, PM_SEM_COUNTER, units_count, 'syscall count'),
            (name + 'latency.avg', None, PM_TYPE_U64, PM_SEM_INSTANT, units_nsecs, 'syscall avg'
                                                                                   'latency'),
            (name + 'latency.cml', None, PM_TYPE_U64, PM_SEM_COUNTER, units_nsecs, 'syscall cml'
                                                                                   'latency'),
        )
        return True, self.items

    def init_insts(self):
        """ Initialize insts """
        self.insts = {}
        for i in range(list(self.syscalls.keys())[-1]):
            if i in self.syscalls:
                self.insts[self.syscalls[i]] = c_int(1)
            else:
                self.insts["UNDEF" + str(i)] = c_int(1)

    def reset_cache(self):
        """ Reset internal cache """
        self.cnt_cache = {}
        self.avg_cache = {}
        self.cml_cache = {}
        self.init_insts()

    def undef_cache(self):
        """ Undefine internal cache """
        self.cnt_cache = None
        self.avg_cache = None
        self.cml_cache = None
        self.insts = None

    def update_pids(self, procs):
        """ Update PIDs to trace """
        prev = self.pids
        if not self.proc_filter:
            return 0
        if not prev and not procs:
            return 0
        if not procs and self.proc_filter:
            self.pids = []
            return -1
        trace_pids = []
        for pid in self.pids:
            if self.pid_alive(pid):
                trace_pids.append(pid)
        if len(trace_pids) == self.max_pids:
            return 0
        found_procs = [p for p in procs if p[0] not in trace_pids]
        if found_procs:
            for proc in found_procs:
                if len(trace_pids) < self.max_pids:
                    trace_pids.append(proc[0])
            self.pids = trace_pids
            for proc in procs:
                if proc[0] in self.pids:
                    info = proc[1] + " " + proc[2] if proc[2] else proc[1]
                    self.log("Tracing PID %s: %s." % (str(proc[0]), info))
            return 1
        return 0


    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 self.errno:
                    errno_filter = "if (args->ret != -%d) return 0;" % self.errno
                    self.bpf_text = self.bpf_text.replace("//FILTER_ERRNO", errno_filter)
                if self.failed:
                    failed_filter = "if (args->ret >= 0) return 0;"
                    self.bpf_text = self.bpf_text.replace("//FILTER_FAILED", failed_filter)
                self.bpf_text = self.bpf_text.replace("NR_SYSCALLS", str(len(self.syscalls)))
                self.bpf_text = self.bpf_text.replace("MAX_PIDS", str(self.max_pids))

            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.bpf_text
            details = "#define DETAILS" if self.details and self.pids else ""
            latency = "#define LATENCY" if self.latency else ""
            bpf_text = bpf_text.replace("DEFINE_DETAILS", details)
            bpf_text = bpf_text.replace("DEFINE_LATENCY", latency)

            bpf_text = self.apply_pid_filter(bpf_text, self.pids, False, True)

            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

        self.init_insts()

        stats = self.bpf["stats"]
        for i in stats:
            val = stats[i].value
            if val:
                key = self.syscalls[i.value]
                self.cnt_cache[key] = val

        if self.latency:
            latstats = self.bpf["latstats"]
            for i in latstats:
                val = latstats[i].value
                if val:
                    key = self.syscalls[i.value]
                    self.cml_cache[key] = val
                    val = val if key not in self.avg_cache else int(val / self.cnt_cache[key])
                    self.avg_cache[key] = val

        check_pids = set()
        stale_pids = set()

        def use_pid(pid):
            """ Helper to quickly check whether to use PID info """
            if pid not in check_pids:
                if pid in stale_pids or not self.pid_alive(pid):
                    stale_pids.add(pid)
                    return False
                else:
                    check_pids.add(pid)
            return True

        if self.details and self.pids:
            pidstats = self.bpf["pidstats"]
            for k, v in pidstats.items():
                if not use_pid(k.pid):
                    continue
                key = str(k.pid) + "::" + self.syscalls[k.id]
                self.cnt_cache[key] = v.value
                self.insts[key] = c_int(1)

        if self.details and self.pids and self.latency:
            pidlatstats = self.bpf["pidlatstats"]
            for k, v in pidlatstats.items():
                if not use_pid(k.pid):
                    continue
                key = str(k.pid) + "::" + self.syscalls[k.id]
                self.cml_cache[key] = v.value
                val = v.value if key not in self.avg_cache else int(v.value / self.cnt_cache[key])
                self.avg_cache[key] = val

        return self.insts

    def bpfdata(self, item, inst):
        """ Return BPF data as PCP metric value """
        try:
            key = self.pmdaIndom.inst_name_lookup(inst)
            if item == 0:
                return [self.cnt_cache[key], 1]
            elif item == 1:
                return [self.avg_cache[key], 1]
            elif item == 2:
                return [self.cml_cache[key], 1]
            else:
                return [PM_ERR_PMID, 0]
        except Exception: # pylint: disable=broad-except
            return [PMDA_FETCH_NOVALUES, 0]