#!/usr/bin/env python

import re
import time
import argparse
import logging

import psutil
import statsd


def find_process(pid, regex):
    if pid:
        try:
            return psutil.Process(pid)
        except psutil.NoSuchProcess:
            pass
    elif regex:
        current_pid = psutil.Process().pid
        for p in psutil.process_iter():
            if regex.search(' '.join(p.cmdline())) and p.pid != current_pid:
                return p
    else:
        raise ValueError('No pid or regex')


def main(host, prefix, interval, pid, regex):
    client = statsd.StatsClient(host=host, prefix='.'.join(filter(bool, ['procmon', prefix])))
    while True:
        try:
            proc = find_process(pid, regex)
            if not proc:
                logging.warning('No target process found')
                time.sleep(interval)
                continue
            else:
                logging.debug('Found process %s, collecting metrics', proc)

            try:
                pct = proc.cpu_percent(interval) # the call is blocking
                mem = proc.memory_info().rss
            except psutil.NoSuchProcess:
                logging.warning('Target process was terminated during collection')
            else:
                client.gauge('cpu', pct)
                client.gauge('rss', mem)
                logging.debug('Sent CPU:%.0f%% RSS:%d', pct, mem)
        except KeyboardInterrupt:
            break


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Process monitor')
    parser.add_argument('-p', '--pid', help='Target process pid', default=0)
    parser.add_argument('-r', '--regex', help='Target process cmdline regex', default='')
    parser.add_argument('-f', '--prefix', help='Statsd prefix', default='')
    parser.add_argument('-s', '--host', help='Statsd host', default='localhost')
    parser.add_argument('-i', '--interval', help='Collection interval in seconds', default=1)
    parser.add_argument('-v', '--verbose', help='Be verbose', action='store_true', default=False)

    args = parser.parse_args()
    if not any([args.pid, args.regex]):
        parser.error('Either PID or cmdline regex must be provided')

    level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(level=level, format='%(asctime)s %(levelname)s %(name)s %(message)s')

    main(args.host, args.prefix, float(args.interval), int(args.pid), re.compile(args.regex))