# coding=utf-8

"""
A Diamond collector that collects memory usage of each process defined in it's
config file by matching them with their executable filepath or the process name.
This collector can also be used to collect memory usage for the Diamond process.

Example config file ProcessResourcesCollector.conf

```
enabled=True
unit=B
cpu_interval=0.1
[process]
[[postgres]]
exe=^\/usr\/lib\/postgresql\/+d.+d\/bin\/postgres$
name=^postgres,^pg

[[diamond]]
selfmon=True
```

exe and name are both lists of comma-separated regexps.

count_workers defined under [process] will determine whether to count how many
workers are there of processes which match this [process],
for example: cgi workers.

"""

import os
import re
from collections import defaultdict

import diamond.collector
import diamond.convertor

try:
    import psutil
    psutil
except ImportError:
    psutil = None


def match_process(pid, name, cmdline, exe, cfg):
    """
    Decides whether a process matches with a given process descriptor

    :param pid: process pid
    :param exe: process executable
    :param name: process name
    :param cmdline: process cmdline
    :param cfg: the dictionary from processes that describes with the
        process group we're testing for
    :return: True if it matches
    :rtype: bool
    """
    if cfg['selfmon'] and pid == os.getpid():
        return True
    for exe_re in cfg['exe']:
        if exe_re.search(exe):
            return True
    for name_re in cfg['name']:
        if name_re.search(name):
            return True
    for cmdline_re in cfg['cmdline']:
        if cmdline_re.search(' '.join(cmdline)):
            return True
    return False

def process_info(process, info_keys):
    results = {}
    process_info = process.as_dict()
    metrics = ((key, process_info.get(key, None)) for key in info_keys)
    for key, value in metrics:
        if type(value) in [float, int]:
            results.update({key: value})
        elif hasattr(value, '_asdict'):
            for subkey, subvalue in value._asdict().iteritems():
                results.update({"%s.%s" % (key, subkey): subvalue})
    return results

class ProcessResourcesCollector(diamond.collector.Collector):
    def __init__(self, *args, **kwargs):
        """
        prepare self.processes, which is a descriptor dictionary in
        pg_name: {
            exe: [regex],
            name: [regex],
            cmdline: [regex],
            selfmon: [boolean],
            procs: [psutil.Process],
            count_workers: [boolean]
        }
        """
        super(ProcessResourcesCollector, self).__init__(*args, **kwargs)
        self.processes = {}
        self.processes_info = {}
        for pg_name, cfg in self.config['process'].items():
            pg_cfg = {}
            for key in ('exe', 'name', 'cmdline'):
                pg_cfg[key] = cfg.get(key, [])
                if not isinstance(pg_cfg[key], list):
                    pg_cfg[key] = [pg_cfg[key]]
                pg_cfg[key] = [re.compile(e) for e in pg_cfg[key]]
            pg_cfg['selfmon'] = cfg.get('selfmon', '').lower() == 'true'
            pg_cfg['count_workers'] = cfg.get(
                'count_workers', '').lower() == 'true'
            self.processes[pg_name] = pg_cfg
            self.processes_info[pg_name] = defaultdict(int)

    def get_default_config_help(self):
        config_help = super(ProcessResourcesCollector,
                            self).get_default_config_help()
        config_help.update({
            'unit': 'The unit in which memory data is collected.',
            'process': ("A subcategory of settings inside of which each "
                        "collected process has it's configuration"),
        })
        return config_help

    def get_default_config(self):
        """
        Default settings are:
            path: 'process'
            unit: 'B'
        """
        config = super(ProcessResourcesCollector, self).get_default_config()
        config.update({
            'path': 'process',
            'unit': 'B',
            'process': '',
        })
        return config

    default_info_keys = [
        'num_ctx_switches',
        'cpu_percent',
        'cpu_times',
        'io_counters',
        'num_threads',
        'memory_percent',
        'ext_memory_info',
    ]

    def save_process_info(self, pg_name, process_info):
        for key, value in process_info.iteritems():
            self.processes_info[pg_name][key] += value

    def collect_process_info(self, process):
        try:
            pid = process.pid
            name = process.name
            cmdline = process.cmdline
            try:
                exe = process.exe
            except psutil.AccessDenied:
                exe = ""
            for pg_name, cfg in self.processes.items():
                if match_process(pid, name, cmdline, exe, cfg):
                    self.save_process_info(pg_name, process_info(process, self.default_info_keys))
        except psutil.NoSuchProcess, e:
            self.log.warning("Process exited while trying to get info: %s", e)

    def collect(self):
        """
        Collects resources usage of each process defined under the
        `process` subsection of the config file
        """
        if not psutil:
            self.log.error('Unable to import psutil')
            self.log.error('No process resource metrics retrieved')
            return None


        
        for process in psutil.process_iter():
            self.collect_process_info(process)

        # publish results
        for pg_name, counters in self.processes_info.iteritems():
            metrics = (("%s.%s" % (pg_name, key), value) for key, value in counters.iteritems())
            [self.publish(*metric) for metric in metrics]
            self.processes_info[pg_name] = defaultdict(int)  # reinitialize process info

@@ -4896,39 +4896,12 @@
     
-self.save_process_info(pg_name,
+pi =
  pro
@@ -4942,16 +4942,172 @@
 fo_keys)
+%0A                    if cfg%5B'count_workers'%5D:%0A                        pi.update(%7B'workers_count': 1%7D)%0A                    self.save_process_info(pg_name, pi
 )%0A      

support workers counter