#!python

import socket
import simplejson as json
import curses
import time
import atexit
import sys
import traceback

need_reset = True

def game_over():
    global need_reset
    if need_reset:
        curses.endwin()

def exc_hook(type, value, tb):
    global need_reset
    need_reset = False
    curses.endwin()
    traceback.print_exception(type, value, tb)

sys.excepthook = exc_hook

argc = len(sys.argv)

if argc < 2:
    raise Exception("You have to specify the uWSGI stats socket")

addr = sys.argv[1]
sfamily = socket.AF_UNIX
addr_tuple = addr
if ':' in addr:
    sfamily = socket.AF_INET
    addr_parts = addr.split(':')
    addr_tuple = (addr_parts[0], int(addr_parts[1]))

freq = 3
try:
    freg = int(sys.argv[2])
except:
    pass

screen = curses.initscr()
screen.timeout(freq*1000)
atexit.register(game_over)

curses.curs_set(0)
screen.clear()

def reqcount(a, b):
    if a['requests'] > b['requests']:
        return -1
    if a['requests'] < b['requests']:
        return 1
    return 0

def calc_percent(tot, req):
    if tot == 0:
        return 0.0
    return (100 *float(req))/float(tot)

while True:

    js = ''

    try:
        s = socket.socket(sfamily, socket.SOCK_STREAM)
        s.connect( addr_tuple )

        while True:
            data = s.recv(4096)
            if len(data) < 1:
                break
            js += data
    except:
        raise Exception("unable to get uWSGI statistics")

    dd = json.loads(js)

    tot = sum( [worker['requests'] for worker in dd['workers']] )
    tx = sum( [worker['tx'] for worker in dd['workers']] )

    
    uversion = ''
    if 'version' in dd:
        uversion = '-' + dd['version']

    if not 'listen_queue' in dd:
        dd['listen_queue'] = 0 
    screen.addstr(0, 0, "uwsgi%s@%s - %s - req: %d - lq: %d - tx: %d" % (uversion, socket.gethostname(), time.ctime(), tot, dd['listen_queue'], tx))
    screen.addstr(1, 0, "\tWID\t%\tPID\tREQ\tEXC\tSTATUS\tAVG\tRSS\tVSZ\tTX\tRT\t", curses.A_REVERSE)
    pos = 2

    dd['workers'].sort(reqcount)
    for worker in dd['workers']:
        screen.addstr(pos, 0, "\t%d\t%.2f\t%d\t%d\t%d\t%s\t%dms\t%d\t%d\t%d\t%d" % (
            worker['id'], calc_percent(tot, worker['requests']), worker['pid'], worker['requests'], worker['exceptions'], worker['status'], worker['avg_rt']/1000,
            worker['rss']/1024/1024, worker['vsz']/1024/1024,
            worker['tx'], worker['running_time']/1000
        ))
        pos += 1
    screen.refresh()

    s.close()
    if screen.getch() == ord('q'):
        game_over()
        break


