#!/usr/bin/env python

import os, sys, signal, time, socket, errno, argparse

RECV_BUFSIZE = 2048

def parse_args():
    try: default_port = int(os.environ['RUN_PORT'])
    except (KeyError, ValueError): default_port = 11211

    parser = argparse.ArgumentParser(description='Gather stats from memcached-compatible server')
    parser.add_argument('--host', '-H', type=str, help='Host server is running on (default: localhost).', default='localhost')
    parser.add_argument('--port', '-p', type=int, help='Port server is running on (default: %d).' % default_port, default=default_port)
    # iostat and vmstat get their argument in seconds. Do we need to support a shorter interval than that?
    parser.add_argument('interval', nargs='?', type=int, help='Interval (in seconds) between updates. If not specified, will print stats once and then exit.')
    parser.add_argument('count', nargs='?', type=int, help='If interval is specified, number of updates. If not specified, will continue indefinitely.')
    return parser.parse_args()

def read_print_stats(sock):
    sock.send('rdb stats\r\n')
    stats = ''
    while stats[-7:] != '\r\nEND\r\n':
        buf = sock.recv(RECV_BUFSIZE)
        if buf == '': raise socket.error('Socket closed')
        else: stats += buf
    sys.stdout.write(stats)
    sys.stdout.flush()

def start(args):
    sock = socket.socket()
    try: sock.connect((args.host, args.port))
    except socket.error, e:
        if e.errno != errno.ECONNREFUSED: raise
        print >> sys.stderr, 'Unable to connect to %s:%d (connection refused).' % (args.host, args.port)
        sys.exit(1)
    return sock

def loop(args, sock):
    if args.interval == None:
        read_print_stats(sock)
    else:
        if args.count == None:
            while True:
                read_print_stats(sock)
                time.sleep(args.interval)
        else:
            for i in xrange(args.count - 1):
                read_print_stats(sock)
                time.sleep(args.interval)
            read_print_stats(sock)


def cleanup(sock):
    try:
        sock.shutdown(socket.SHUT_RDWR)
        sock.close()
    except socket.error:
        pass
    sys.exit(0)


args = parse_args()
sock = start(args)
try: loop(args, sock)
except KeyboardInterrupt: pass
except socket.error: pass
cleanup(sock)
