'''

Copyright (C) 2019 Vanessa Sochat.

This Source Code Form is subject to the terms of the
Mozilla Public License, v. 2.0. If a copy of the MPL was not distributed
with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

'''

from subprocess import (
    check_output, 
    CalledProcessError
)
from watchme.utils import get_watchme_env
from watchme.logger import bot
import psutil

def monitor_pid_task(**kwargs):
    '''monitor a specific process. This function can be used as a task, or
       is (most likely used) for the psutils.decorators. A pid parameter
       is required.

       Parameters
       ==========
       skip: an optional list of (comma separated) fields to skip. Can be in
             net_io_counters,net_connections,net_if_address,net_if_stats
       pid: the process id or name (required)
    '''
    pid = kwargs.get('pid', None)

    bot.debug(kwargs)

    # Helper function to get list from one,two,three
    def get_list(name):
        csv_list = kwargs.get(name, '')
        return [x for x in csv_list.split(',') if x]

    # A comma separated list of parameters to skip
    skip = get_list('skip')
    include = get_list('include')
    only = get_list('only')

    # Only continue given that argument is provided
    if pid is None:   
        bot.warning("A 'pid' parameter is required to use the monitor_pid_task function.")
        return pid

    # The user is allowed to provide a process name, or a number
    try:
        pid = int(pid)
    except ValueError:
        pid = _get_pid(pid)

    # But if it's stil no good (None) we exit.
    if pid is None:
        bot.warning("'pid' must be a running process or process name.")
        return pid

    ps = psutil.Process(pid)
    bot.debug(ps)

    results = {}

    for key, val in ps.as_dict().items():
 
        # If val is None, don't include
        if val is None:
            bot.debug('skipping %s, None' % key)
            continue

        if key == "connections":
            connections = []
            for net in val:
                entry = {'fd': net.fd,
                         'family': str(net.family),
                         'type': str(net.type),
                         'laddr_ip': net.laddr.ip,
                         'laddr_port': net.laddr.port,
                         'raddr': net.raddr,
                         'status': net.status}
                connections.append(entry)
            val = connections       

        # First priority goes to a custom set
        if len(only) > 0:
            if key in only:
                results[key] = val
            else:
                bot.debug('skipping %s' % key)
            continue 

        # The user requested to skip
        elif key in skip or key in ['threads']:
            continue

        # Keep count of openfiles
        elif key in ["open_files"]:
            results[key] = len(val)

        # Don't risk exposing sensitive information
        elif key == "environ":
            if key in include:
                results[key] = val

        # Skip over ones that are too detailed
        elif key in ['memory_maps']:
            continue

        # I assume this included all that are in pmem too
        elif isinstance(val, psutil._pslinux.pfullmem):
            results[key] = {"rss": val.rss,
                            "vms": val.vms,
                            "shared": val.shared,
                            "text": val.text,
                            "lib": val.lib,
                            "data": val.data,
                            "dirty": val.dirty,
                            "uss": val.uss,
                            "pss": val.pss,
                            "swap": val.swap}

        elif isinstance(val, psutil._common.pgids) or isinstance(val, psutil._common.puids):
            results[key] = {"real": val.real,
                            "effetive":val.effective,
                            "saved": val.saved}

        elif isinstance(val, psutil._common.pcputimes):
            results[key] = {"user": val.user,
                            "system": val.system,
                            "children_user": val.children_user,
                            "children_system": val.children_system}

        elif isinstance(val, psutil._common.pctxsw):
            results[key] = {"voluntary": val.voluntary,
                            "involuntary":val.involuntary}

        elif isinstance(val, psutil._common.pionice):
            results[key] = {"value": val.value}

            # Older Python version (2) doesn't have attribute
            if hasattr(val.ioclass, 'name'):
                results[key]["ioclass"] = val.ioclass.name

        # pfullmem (first above) should cover this
        elif isinstance(val, psutil._pslinux.pmem):
            continue

        elif isinstance(val, psutil._pslinux.pio):
            results[key] = {"read_count": val.read_count,
                            "write_count": val.write_count,
                            "read_bytes": val.read_bytes,
                            "write_bytes": val.write_bytes,
                            "read_chars": val.read_chars,
                            "write_chars": val.write_chars}
        else:
            results[key] = val

    # Add any environment variables prefixed wit WATCHMEENV_
    environ = get_watchme_env()
    results.update(environ)

    return results

def cpu_task(**kwargs):
    '''Get variables about the cpu of the host. No parameters are required.

       Parameters
       ==========
       skip: an optional list of (comma separated) fields to skip. Can be in
             net_io_counters,net_connections,net_if_address,net_if_stats
    '''
    result = {}

    # A comma separated list of parameters to not include
    skip = kwargs.get('skip', '')
    skip = skip.split(',')

    result['cpu_freq'] = dict(psutil.cpu_freq()._asdict())
    result['cpu_percent'] = psutil.cpu_percent()
    result['cpu_count'] = psutil.cpu_count()
    result['cpu_stats'] = dict(psutil.cpu_stats()._asdict())
    result['cpu_times'] = dict(psutil.cpu_times()._asdict())
    result['cpu_times_percent'] = dict(psutil.cpu_times_percent()._asdict())
    return _filter_result(result, skip)


def python_task(**kwargs):
    '''Get modules, version, etc. about currently running python

       Parameters
       ==========
       skip: an optional list of (comma separated) fields to skip. Can be in
             net_io_counters,net_connections,net_if_address,net_if_stats
    '''
    # A comma separated list of parameters to not include
    skip = kwargs.get('skip', '')
    skip = skip.split(',')

    result = {'modules': list(psutil.os.sys.modules),
              'version': psutil.os.sys.version,
              'path': psutil.os.sys.path,
              'prefix': psutil.os.sys.prefix,
              'executable': psutil.os.sys.executable,
              'copyright': psutil.os.sys.copyright }

    return _filter_result(result, skip)


def system_task(**kwargs):
    '''Get basic system info, unlikely to change.

       Parameters
       ==========
       skip: an optional list of (comma separated) fields to skip. Can be in
             net_io_counters,net_connections,net_if_address,net_if_stats
    '''
    # A comma separated list of parameters to not include
    skip = kwargs.get('skip', '')
    skip = skip.split(',')

    result = {}
    result['platform'] = psutil.os.sys.platform
    result['api_version'] = psutil.os.sys.api_version
    result['maxsize'] = psutil.os.sys.maxsize

    # Python 2 doesn't have these fields
    if hasattr(psutil.os.sys, 'int_info'):
        result['bits_per_digit'] = psutil.os.sys.int_info.bits_per_digit
        result['sizeof_digit'] = psutil.os.sys.int_info.sizeof_digit
 
    vinfo = psutil.os.sys.version_info
    result['version_info'] = {'major': vinfo.major,
                              'minor': vinfo.minor,
                              'micro': vinfo.micro,
                              'releaselevel': vinfo.releaselevel,
                              'serial': vinfo.serial}

    result['sep'] = psutil.os.sep
    uname = psutil.os.uname()
    
    # Python 2 returns a tuple 
    if not hasattr(uname, 'sysname'):
        result['uname'] = uname
    else:
        result['uname'] = {'sysname': uname.sysname,
                           'nodename': uname.nodename,
                           'version': uname.version,
                           'release': uname.release,
                           'machine': uname.machine}

    return _filter_result(result, skip)


def memory_task(**kwargs):
    '''Get values for memory of the host. No parameters are required.
    '''

    result = {}

    # gives an object with many fields
    result['virtual_memory'] = dict(psutil.virtual_memory()._asdict())

    # Add any environment variables prefixed wit WATCHMEENV_
    environ = get_watchme_env()
    result.update(environ)

    return result


def users_task(**kwargs):
    '''Get values for current users
    '''

    result = {'users': []}
    for user in psutil.users():
        result['users'].append(dict(user._asdict()))

    # Add any environment variables prefixed wit WATCHMEENV_
    environ = get_watchme_env()
    result.update(environ)

    return result


def sensors_task(**kwargs):
    '''Get values from system sensors like fans, temperature, battery
    '''

    # A comma separated list of parameters to not include
    skip = kwargs.get('skip', '')
    skip = skip.split(',')

    result = {}
    result['sensors_temperatures'] = {}
    result['sensors_fans'] = {}

    # battery
    battery = psutil.sensors_battery()
    if battery is not None:
        result['sensors_battery'] = {'percent': battery.percent,
                                     'secsleft': str(battery.secsleft),
                                     'power_plugged': battery.power_plugged}

    # fans
    for name, entry in psutil.sensors_fans().items():
        entries = []
        for e in entry:
            item = {'label': e.label,
                    'current': e.current}
            entries.append(item)
        result['sensors_fans'][name] = entries

    # temperature
    for name, entry in psutil.sensors_temperatures().items():
        entries = []
        for e in entry:
            temp = {'label': e.label,
                    'current': e.current,
                    'high': e.high,
                    'critical': e.critical}
            entries.append(temp)
        result['sensors_temperatures'][name] = entries

    return _filter_result(result, skip)



def net_task(**kwargs):
    '''Get values for networking on the host

       Parameters
       ==========
       skip: an optional list of (comma separated) fields to skip. Can be in
             net_io_counters,net_connections,net_if_address,net_if_stats
    '''
    result = {}

    # A comma separated list of parameters to not include
    skip = kwargs.get('skip', '')
    skip = skip.split(',')

    result['net_connections'] = []
    result['net_if_address'] = {}
    result['net_if_stats'] = {}

    for net in psutil.net_connections():
        net_result = {'fd': net.fd,
                      'family': str(net.family),
                      'type': str(net.type),
                      'laddr_ip': net.laddr.ip,
                      'laddr_port': net.laddr.port,
                      'raddr': net.raddr,
                      'status': net.status,
                      'pid': net.pid}
        result['net_connections'].append(net_result)

    for name, entry in psutil.net_if_addrs().items():
        entries = []
        for e in entry:
            item = {'family': str(e.family),
                    'address': e.address,
                    'netmask': e.netmask,
                    'broadcast': e.broadcast,
                    'ptp': e.ptp}
            entries.append(item)
        result['net_if_address'][name] = entries

    for name, entry in psutil.net_if_stats().items():
        item = {'isup': entry.isup,
                'duplex': str(entry.duplex),
                'speed': entry.speed,
                'mtu': entry.mtu}
        result['net_if_stats'][name] = item

    netio = psutil.net_io_counters()
    result['net_io_counters'] = {'bytes_sent': netio.bytes_sent,
                                 'bytes_recv': netio.bytes_recv,
                                 'packets_sent': netio.packets_sent,
                                 'packets_recv': netio.packets_recv,
                                 'errin': netio.errin,
                                 'errout': netio.errout,
                                 'dropin': netio.dropin,
                                 'dropout': netio.dropout}

    # Remove those not wanted
    return _filter_result(result, skip)


def disk_task(**kwargs):
    '''Get values for disks on the host

       Parameters
       ==========
       disk_usage: an optional path to get disk usage for.
       skip: an optional list of (comma separated) fields to skip. Can be in
             net_io_counters,net_connections,net_if_address,net_if_stats

    '''
    # A comma separated list of parameters to not include
    skip = kwargs.get('skip', '')
    skip = skip.split(',')
    path = kwargs.get('disk_usage', '/')

    result = {}
    result['disk_to_counters'] = dict(psutil.disk_io_counters()._asdict())
    result['disk_partitions'] = []
    for disk in psutil.disk_partitions():
        disk_result = {'device': disk.device,
                       'mountpoint': disk.mountpoint,
                       'fstype': disk.fstype,
                       'opts': disk.opts}
        result['disk_partitions'].append(disk_result)
    result['disk_usage'] = dict(psutil.disk_usage(path)._asdict())

    # Remove those not wanted
    return _filter_result(result, skip)


def _filter_result(result, skip):
    '''a helper function to filter a dictionary based on a list of keys to 
       skip. We also add variables from the environment.
    
       Parameters
       ==========
       result: a dictionary of results
       skip: a list of keys to remove/filter from the result.
    '''

    # Add any environment variables prefixed wit WATCHMEENV_
    environ = get_watchme_env()
    result.update(environ)

    for key in skip:
        if key in result:
            del result[key]

    return result


def _get_pid(name):
    '''used to get the pid of a process, by name
       
       Parameters
       ==========
       name: the name of the process to get.
    '''
    try:
        pid = check_output(["pidof", name])
        pid = pid.decode('utf-8').strip('\n').split(' ') 
        if len(pid) > 1:
            bot.warning("More than one pid found for %s, using first." % name)
        pid = int(pid[0])

    except CalledProcessError:
        bot.error("%s does not exist." % name)
        pid = None
    return pid
