#!/usr/bin/env python
from argparse import ArgumentParser
import os
import re
import sha
import subprocess
import sys


try:
    import json
except ImportError:
    import simplejson as json


try:
    parser = ArgumentParser()
    parser.add_argument(
        "-s", dest="services", required=True,
        help="Services to check, comma-separated",
    )
    parser.add_argument(
        "-u", "--require-up", dest="require_up", action="store_true",
        help="If set, require all services to be running",
    )
    parser.add_argument(
        "-p", "--path", dest="service_path", default="/etc/service/",
        help="Prefix path of where daemontools services live",
    )
    parser.add_argument(
        "-n", "--spins", dest="max_spins", type=int, default=3,
        help="Maximum tolerated number of spins",
    )
    args = parser.parse_args()

    ident = sha.new(args.services).digest().encode('hex')
    spinfile = '/tmp/spincheck_%s' % (ident,)

    try:
        with open(spinfile, 'r') as f:
            old = json.load(f)
    except (IOError, ValueError):
        old = {}

    check_command = 'svstat %s*' % (args.service_path,)

    p = subprocess.Popen(check_command, shell=True, stdout=subprocess.PIPE)
    path_services = p.communicate()[0].splitlines()
    if p.returncode != 0:
        sys.exit(2)

    new = {}
    check_services = set(args.services.split(','))
    running_services = set()

    for line in path_services:
        stat = line.split(':')[0]
        service_name = os.path.split(stat)[-1]
        if service_name not in check_services:
            continue

        count, oldpid = old.get(service_name, [0, 0])
        match = re.match('.*\(pid ([0-9]*)\)', line)

        if not match:
            continue

        running_services.add(service_name)

        newpid = int(match.group(1))

        if newpid == oldpid:
            new[service_name] = [0, newpid]
        else:
            new[service_name] = [count + 1, newpid]

    with open(spinfile, 'w') as f:
        json.dump(new, f)

    down_services = check_services - running_services
    spinning = []
    for k, v in new.iteritems():
        if v[0] > args.max_spins:
            spinning.append(k)

    down_summary = 'Down: %s' % ', '.join(down_services)
    running_summary = 'Running: %s' % ', '.join(running_services)

    if len(spinning) > 0:
        print 'Spinning: %s\n%s\n%s' % (', '.join(spinning), down_summary, running_summary)
        sys.exit(2)

    if args.require_up and down_services:
        print '%s\n%s' % (down_summary, running_summary)
        sys.exit(2)

    print "OK: no workers spinning.\n%s\n%s" % (running_summary, down_summary)

except Exception, e:
    print repr(e)
    sys.exit(3)
