__author__ = "Vanessa Sochat"
__copyright__ = "Copyright 2020-2021, Vanessa Sochat"
__license__ = "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
