#!/usr/bin/python

import socket

try:
    import simplejson as json
except ImportError:
    import json
import curses
import time
import atexit
import sys
import traceback
from collections import defaultdict

need_reset = True
screen = None

import locale
locale.setlocale(locale.LC_NUMERIC, "")


def format_num(num):
    """Format a number according to given places.
    Adds commas, etc. Will truncate floats into ints!"""

    try:
        inum = int(num)
        return locale.format("%.*f", (0, inum), True)

    except (ValueError, TypeError):
        return str(num)


def get_max_width(table, index):
    """Get the maximum width of the given column index"""
    return max([len(format_num(row[index])) for row in table])


def human_size(n):
    # G
    if n >= (1024 * 1024 * 1024):
        return "%.1fG" % (n / (1024 * 1024 * 1024))
    # M
    if n >= (1024 * 1024):
        return "%.1fM" % (n / (1024 * 1024))
    # K
    if n >= 1024:
        return "%.1fK" % (n / 1024)
    return "%d" % n


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


def exc_hook(type, value, tb):
    global need_reset, screen
    need_reset = False
    if screen:
        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:
    freq = int(sys.argv[2])
except:
    pass

screen = curses.initscr()
curses.start_color()

try:
    # busy
    curses.init_pair(1, curses.COLOR_GREEN, curses.COLOR_BLACK)
    # cheap
    curses.init_pair(2, curses.COLOR_MAGENTA, curses.COLOR_BLACK)
    # pause
    curses.init_pair(3, curses.COLOR_RED, curses.COLOR_BLACK)
    # sig
    curses.init_pair(4, curses.COLOR_YELLOW, curses.COLOR_BLACK)
except curses.error:
    # the terminal doesn't support colors
    pass

screen.timeout(freq * 1000)
atexit.register(game_over)

try:
    curses.curs_set(0)
except:
    pass
screen.clear()


def calc_percent(req, tot):
    if tot == 0:
        return 0.0
    return "%.2f" % ((100 * float(req)) / float(tot))

try:
    # busy
    curses.init_pair(1, curses.COLOR_WHITE, curses.COLOR_GREEN)
    # cheap
    curses.init_pair(2, curses.COLOR_WHITE, curses.COLOR_MAGENTA)
    # pause
    curses.init_pair(3, curses.COLOR_WHITE, curses.COLOR_RED)
    # sig
    curses.init_pair(4, curses.COLOR_WHITE, curses.COLOR_YELLOW)
except curses.error:
    # the terminal doesn't support colors
    pass

# RPS calculation
last_tot_time = time.time()
last_reqnumber = 0

last_hits_per_cache = defaultdict(int)
last_miss_per_cache = defaultdict(int)

while True:
    screen.clear()
    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.decode('utf8')
    except:
        raise Exception("unable to get uWSGI statistics")

    dd = json.loads(js)

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

    if not 'listen_queue' in dd:
        dd['listen_queue'] = 0

    cwd = ""
    if 'cwd' in dd:
        cwd = "- cwd: %s" % dd['cwd']

    uid = ""
    if 'uid' in dd:
        uid = "- uid: %d" % dd['uid']

    gid = ""
    if 'gid' in dd:
        gid = "- gid: %d" % dd['gid']

    masterpid = ""
    if 'pid' in dd:
        masterpid = "- masterpid: %d" % dd['pid']

    screen.addstr(1, 0, "node: %s %s %s %s %s" % (socket.gethostname(), cwd, uid, gid, masterpid))
    tot = sum([worker['requests'] for worker in dd['workers']])
    dt = time.time() - last_tot_time
    total_rps = (tot - last_reqnumber) / dt
    last_reqnumber = tot

    tx = human_size(sum([worker['tx'] for worker in dd['workers']]))
    screen.addstr(0, 0, "uwsgi%s - %s - req: %d - RPS: %d - lq: %d - tx: %s" % (uversion, time.ctime(), tot, int(round(total_rps)), dd['listen_queue'], tx))
    rows = [("NAME", "HITS", "MISS", "HITS/s", "MISS/s", "ITEMS", "MAXITEMS", "H/M RAT", "FULL")]
    for cache in dd['caches']:
        ratio = calc_percent(cache['hits'], cache['hits'] + cache['miss'])
        hits_per_s = (cache['hits'] - last_hits_per_cache[cache['name']]) / dt
        miss_per_s = (cache['miss'] - last_miss_per_cache[cache['name']]) / dt
        last_hits_per_cache[cache['name']] = cache['hits']
        last_miss_per_cache[cache['name']] = cache['miss']
        rows.append((cache['name'], cache['hits'], cache['miss'], hits_per_s, miss_per_s, cache['items'], cache['max_items'], ratio, cache['full']))

    last_tot_time = time.time()

    col_paddings = []
    for i in range(len(rows[0])):
        col_paddings.append(get_max_width(rows, i))

    pos = 2
    # heading
    heading = rows.pop(0)
    screen.addstr(pos, 0, reduce(lambda x, y: x + str(heading[y]).ljust(col_paddings[y] + 2), range(len(heading)), ""), curses.A_BOLD)
    pos = 3

    for row in rows:
        screen.addstr(pos, 0, "")
        for i in range(len(row)):
            screen.addstr(format_num(row[i]).ljust(col_paddings[i] + 2))
        pos += 1

    screen.refresh()
    s.close()

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