#!/usr/bin/env python

import subprocess
import sys
import os
import json


def pages2mb(pages):
    return pages * 4096 / 1024 / 1024


def main():
    vzlist = "/usr/sbin/vzlist --json"

    if not os.geteuid() == 0:
        sys.exit("This program can only be run under root.")

    if len(sys.argv) > 1 and sys.argv[1] == "-a":
        all = "-a"
    else:
        all = ""

    proc = subprocess.Popen("%s %s" % (vzlist, all), shell=True, stdout=subprocess.PIPE)
    data = json.loads(proc.communicate()[0])

    for i in open("/proc/meminfo"):
        if i.startswith("MemTotal"):
            node_ram = int(i.split()[1]) / 1024
    node_la = open("/proc/loadavg").read().split()
    node_la = "%.2f/%.2f/%.2f" % (float(node_la[0]), float(node_la[1]), float(node_la[2]))
    node_cpu = 0
    for i in open("/proc/cpuinfo"):
        if i.startswith("processor"):
            node_cpu += 1

    used_cpu = 0
    used_ram = 0

    print "%-8s %-25s %8s   / %4s %17s  %5s %10s" % ("CTID", "HOSTNAME", "USED", "TOTAL",
                                                     "LOADAVG\t", "CPUS", "CPUUNITS")

    for ct in data:
        ctid = ct["ctid"]
        hostname = ct["hostname"]
        ram_limit = pages2mb(ct["physpages"]["limit"])
        used_ram += ram_limit
        ram_current = pages2mb(ct["physpages"]["held"])
        cpus = ct["cpus"]
        used_cpu += cpus
        cpuunits = ct["cpuunits"]
        if ct["status"] != "running":
            hostname = "(stopped) %s" % hostname
            loadavg = None
        else:
            loadavg = []
            for la in ct["laverage"]:
                loadavg.append("%0.02f" % la)
            loadavg = "/".join(loadavg)
        if len(hostname) > 25:
            hostname = hostname[:24]
        print "%-8s [%-25s] %8d / %4d MB %17s %5s %10s" % (ctid, hostname, ram_current, ram_limit,
                                                           loadavg, cpus, cpuunits)

    print "\nNode loadavg: %s" % node_la
    if used_cpu > node_cpu:
        overcommit = used_cpu - node_cpu
        node_cpu = "%d (overcommitment: %d)" % (node_cpu, overcommit)
    print "Allocated cores: %d / %s" % (used_cpu, node_cpu)
    if used_ram > node_ram:
        overcommit = used_ram - node_ram
        node_ram = "%d MB (overcommitment: %d MB)" % (node_ram, overcommit)
    print "Allocated memory: %d / %s" % (used_ram, node_ram)

if __name__ == '__main__':
    main()