#!/usr/bin/env python2

import zmq, msgpack, time, psutil, yaml, os, subprocess
from collections import namedtuple

# Horrible hack to make check_output exist in 2.6
# http://stackoverflow.com/a/13160748/1332715
if "check_output" not in dir( subprocess ): # duck punch it in!
	def f(*popenargs, **kwargs):
		if 'stdout' in kwargs:
			raise ValueError('stdout argument not allowed, it will be overridden.')
		process = subprocess.Popen(stdout=subprocess.PIPE, *popenargs, **kwargs)
		output, unused_err = process.communicate()
		retcode = process.poll()
		if retcode:
			cmd = kwargs.get("args")
			if cmd is None:
				cmd = popenargs[0]
			raise subprocess.CalledProcessError(retcode, cmd)
		return output
	subprocess.check_output = f

ctx = zmq.Context()

sock = ctx.socket(zmq.PUSH)
sock.connect("ipc:///tmp/cstatsd")

with open("config/machine.yaml", "r") as cfile:
	config = yaml.safe_load(cfile)
	
interval = config["interval"]
old_net_data = {}

disk_map = {}
last_io_data = {}

if os.path.exists("/proc/user_beancounters") and not os.path.exists("/proc/vz/vestat"):
	openvz_burst = True
	FakeRam = namedtuple("FakeRam", ["total", "used", "available", "percent", "buffers", "cached"])
else:
	openvz_burst = False

for disk in psutil.disk_partitions():
	disk_map[disk.device] = disk
	
if len(disk_map) == 0:
	# We're probably on OpenVZ, so /proc/partitions doesn't exist. Fall back to 'df'.
	FakeDisk = namedtuple("FakeDisk", ["device", "mountpoint"])
	for line in subprocess.check_output(["df"]).splitlines()[1:]:
		device, _, _, _, _, mountpoint = line.split()
		disk_map[device] = FakeDisk(device, mountpoint)

while True:
	load_avgs = os.getloadavg()
	sock.send(msgpack.packb({
		"service": "machine",
		"msg_type": "value",
		"resource_type": "load_average",
		"unit": "",
		"values": {
			"1m": load_avgs[0],
			"5m": load_avgs[1],
			"15m": load_avgs[2]
		}
	}))
	
	cpu_loads = psutil.cpu_percent(percpu=True)
	
	for i in xrange(0, len(cpu_loads)):
		sock.send(msgpack.packb({
			"service": "machine",
			"msg_type": "value",
			"resource_type": "cpu",
			"unit": "core%d" % (i + 1),
			"values": {
				"load": cpu_loads[i]
			}
		}))
	
	try:
		io_counters = psutil.disk_io_counters(perdisk=True)
	except IOError, e:
		io_counters = {} # OpenVZ...
	
	for drive in config["drives"]:
		drive_data = psutil.disk_usage(disk_map[drive].mountpoint)
		io_data = None
		
		for diskname, data in io_counters.iteritems():
			if drive.endswith(diskname):
				io_data = data
				
		if io_data is None or drive not in last_io_data:
			read_bps = 0
			write_bps = 0
			read_iops = 0
			write_iops = 0
		else:
			read_bps = (io_data.read_bytes - last_io_data[drive].read_bytes) / interval
			write_bps = (io_data.write_bytes - last_io_data[drive].write_bytes) / interval
			read_iops = (io_data.read_count - last_io_data[drive].read_count) / interval
			write_iops = (io_data.write_count - last_io_data[drive].write_count) / interval
			
		if io_data is not None:
			last_io_data[drive] = io_data
			
		sock.send(msgpack.packb({
			"service": "machine",
			"msg_type": "value",
			"resource_type": "disk",
			"unit": drive,
			"values": {
				"total": drive_data.total,
				"used": drive_data.used,
				"free": drive_data.free,
				"used_percentage": drive_data.percent,
				"bps_read": read_bps,
				"bps_write": write_bps,
				"iops_read": read_iops,
				"iops_write": write_iops,
			}
		}))
		
	if openvz_burst:
		# Sigh, OpenVZ... let's use 'free', since that apparently -does- understand OpenVZ.
		lines = subprocess.check_output(["free", "-b"]).splitlines()
		_, ram_total, ram_used, ram_free, _, ram_buffers, ram_cached = lines[1].split()
		_, _, _, ram_available = lines[2].split()
		ram_total = int(ram_total)
		ram_free = int(ram_free)
		ram_buffers = int(ram_buffers)
		ram_cached = int(ram_cached)
		ram_available = int(ram_available)
		ram_used = int(ram_used)
		ram_percent = 1.0 * (ram_total - ram_available) / ram_total * 100
		ram_data = FakeRam(ram_total, ram_used, ram_available, ram_percent, ram_buffers, ram_cached)
	else:
		ram_data = psutil.virtual_memory()
		
	sock.send(msgpack.packb({
		"service": "machine",
		"msg_type": "value",
		"resource_type": "memory",
		"unit": "physical",
		"values": {
			"total": ram_data.total,
			"used": ram_data.used,
			"free": ram_data.available,
			"used_percentage": ram_data.percent,
			"buffers": ram_data.buffers,
			"cache": ram_data.cached
		}
	}))
		
	swap_data = psutil.swap_memory()
	sock.send(msgpack.packb({
		"service": "machine",
		"msg_type": "value",
		"resource_type": "memory",
		"unit": "swap",
		"values": {
			"total": swap_data.total,
			"used": swap_data.used,
			"free": swap_data.free,
			"used_percentage": swap_data.percent
		}
	}))
	
	net_data = psutil.net_io_counters(pernic=True)
	for nic, data in net_data.iteritems():
		try:
			old_in_b = old_net_data[nic].bytes_recv
			old_out_b = old_net_data[nic].bytes_sent
			old_in_p = old_net_data[nic].packets_recv
			old_out_p = old_net_data[nic].packets_sent
		except KeyError, e:
			# No old data yet, first run? Save and skip to next...
			old_net_data[nic] = data
			continue
		
		diff_in_b = data.bytes_recv - old_in_b
		diff_out_b = data.bytes_sent - old_out_b
		diff_in_p = data.packets_recv - old_in_p
		diff_out_p = data.packets_sent - old_out_p
		
		if diff_in_b < 0:
			diff_in_b = 0
		
		if diff_out_b < 0:
			diff_out_b = 0
		
		if diff_in_p < 0:
			diff_in_p = 0
		
		if diff_out_p < 0:
			diff_out_p = 0
			
		old_net_data[nic] = data
		
		sock.send(msgpack.packb({
			"service": "machine",
			"msg_type": "value",
			"resource_type": "network",
			"unit": nic,
			"values": {
				"bps_in": diff_in_b / interval,
				"bps_out": diff_out_b / interval,
				"pps_in": diff_in_p / interval,
				"pps_out": diff_out_p / interval
			}
		}))
		
	sock.send(msgpack.packb({
		"service": "machine",
		"msg_type": "value",
		"resource_type": "uptime",
		"unit": "",
		"values": {
			"uptime": time.time() - psutil.get_boot_time()
		}
	}))
	
	time.sleep(interval)