#!/usr/bin/python
"""
Collect and generate RRD graphs.
"""

from pyrrd.rrd import DataSource, RRA, RRD
from pyrrd.graph import DEF, CDEF
from pyrrd.graph import AREA
from pyrrd.graph import ColorAttributes, Graph
from pyrrd.meta import version as PYRRD_VERSION
PYRRD_VERSION = tuple([int(_) for _ in PYRRD_VERSION.split('.')])
if PYRRD_VERSION < (0, 1, 0):
	RIGID_ARG = '--rigid'  # workaround for broken bool param handling
else:
	#from pyrrd.backend import bindings as backend
	RIGID_ARG = True
import libvirt
import time
import sys
import os

STEP = 10 # collect data every X seconds

MINUTE = 60
HOUR = 60 * MINUTE
DAY = 24 * HOUR
WEEK = 7 * DAY
MONTH = 30 * DAY
YEAR = 365 * DAY
HALF = YEAR / 2
QUATER = YEAR / 4

FILENAME = "uvmm.rrd"
GRAPHFILE = "uvmm.png"
KEYS = ('kernel', 'idle', 'user', 'iowait',
	'cached', 'total', 'buffers', 'free')


def collect():
	"""
	Collect data from libvirt.
	"""
	dss = [
		DataSource(dsName="kernel", dsType="COUNTER", heartbeat="600"),
		DataSource(dsName="idle", dsType="COUNTER", heartbeat="600"),
		DataSource(dsName="user", dsType="COUNTER", heartbeat="600"),
		DataSource(dsName="iowait", dsType="COUNTER", heartbeat="600"),
		DataSource(dsName="cached", dsType="GAUGE", heartbeat="600"),
		DataSource(dsName="total", dsType="GAUGE", heartbeat="600"),
		DataSource(dsName="buffers", dsType="GAUGE", heartbeat="600"),
		DataSource(dsName="free", dsType="GAUGE", heartbeat="600"),
		]
	rras = [
		RRA(cf='LAST', xff=0.5, steps=1, rows=5 * MINUTE / STEP),
		RRA(cf='AVERAGE', xff=0.5, steps=MINUTE / STEP, rows=HOUR / MINUTE),
		RRA(cf='AVERAGE', xff=0.5, steps=HOUR / STEP, rows=DAY / HOUR),
		RRA(cf='AVERAGE', xff=0.5, steps=DAY / STEP, rows=WEEK / DAY),
		]
	my_rrd = RRD(FILENAME, step=STEP, ds=dss, rra=rras)
	# , backend=backend
	if not os.path.exists(FILENAME):
		my_rrd.create()

	conn = libvirt.open('qemu:///system')
	#model, memory, cpus, mhz, nodes, sockets, cores, threads = conn.getInfo()
	while True:
		now = int(time.time())

		cpu = conn.getCPUStats(libvirt.VIR_NODE_CPU_STATS_ALL_CPUS, 0)
		mem = conn.getMemoryStats(libvirt.VIR_NODE_MEMORY_STATS_ALL_CELLS, 0)

		values = [cpu[key] for key in KEYS if key in cpu] + \
			[mem[key] for key in KEYS if key in mem]
		my_rrd.bufferValue(now, *values)

		my_rrd.update()
		time.sleep(STEP)


def graph_cpu(now, width=HOUR):
	"""
	Draw RRD cpu graph.
	"""
	my_rrd = RRD(FILENAME, mode='r') # , backend=backend

	def1 = DEF(rrdfile=my_rrd.filename, vname="vkernel", dsName="kernel")
	def2 = DEF(rrdfile=my_rrd.filename, vname="viowait", dsName="iowait")
	def3 = DEF(rrdfile=my_rrd.filename, vname="vuser", dsName="user")
	def4 = DEF(rrdfile=my_rrd.filename, vname="vidle", dsName="idle")
	cdef1 = CDEF(vname="ckernel",
			rpn="vkernel,vkernel,viowait,+,vuser,+,vidle,+,/,100,*")
	cdef2 = CDEF(vname="ciowait",
			rpn="viowait,vkernel,viowait,+,vuser,+,vidle,+,/,100,*")
	cdef3 = CDEF(vname="cuser",
			rpn="vuser,vkernel,viowait,+,vuser,+,vidle,+,/,100,*")
	cdef4 = CDEF(vname="cidle",
			rpn="vidle,vkernel,viowait,+,vuser,+,vidle,+,/,100,*")
	area1 = AREA(defObj=cdef1, color="#555555", legend="Kernel")
	area2 = AREA(defObj=cdef2, color="#FF0000", legend="IO wait", stack=True)
	area3 = AREA(defObj=cdef3, color="#0000FF", legend="user", stack=True)
	area4 = AREA(defObj=cdef4, color="#00FF00", legend="idle", stack=True)

	cattr = ColorAttributes()
	cattr.back = '#333333'
	cattr.canvas = '#333333'
	cattr.shadea = '#000000'
	cattr.shadeb = '#111111'
	cattr.mgrid = '#CCCCCC'
	cattr.axis = '#FFFFFF'
	cattr.frame = '#AAAAAA'
	cattr.font = '#FFFFFF'
	cattr.arrow = '#FFFFFF'

	graph = Graph(GRAPHFILE, start=now - width, end=now,
			lower_limit=0, upper_limit=100, rigid=RIGID_ARG,
			vertical_label="'CPU Usage (%)'", color=cattr)
	graph.data.extend([def1, def2, def3, def4,
		cdef1, cdef2, cdef3, cdef4,
		area1, area2, area3, area4])
	graph.write(debug=True)


def graph_mem(now, width=HOUR):
	"""
	Draw RRD memory graph.
	"""
	my_rrd = RRD(FILENAME, mode='r') # , backend=backend

	def1 = DEF(rrdfile=my_rrd.filename, vname="vtotal", dsName="total")
	def2 = DEF(rrdfile=my_rrd.filename, vname="vbuffers", dsName="buffers")
	def3 = DEF(rrdfile=my_rrd.filename, vname="vcached", dsName="cached")
	def4 = DEF(rrdfile=my_rrd.filename, vname="vfree", dsName="free")
	cdef1 = CDEF(vname="ctotal", rpn="vtotal,1024,*")
	cdef2 = CDEF(vname="cused", rpn="vtotal,vfree,-,1024,*")
	cdef3 = CDEF(vname="cbuffers", rpn="vbuffers,1024,*")
	cdef4 = CDEF(vname="ccached", rpn="vcached,1024,*")
	area1 = AREA(defObj=cdef1, color="#00FF00", legend="Free")
	area2 = AREA(defObj=cdef2, color="#FF0000", legend="Used")
	area3 = AREA(defObj=cdef3, color="#00FFFF", legend="Buffers", stack=True)
	area4 = AREA(defObj=cdef4, color="#0000FF", legend="Cached", stack=True)

	cattr = ColorAttributes()
	cattr.back = '#333333'
	cattr.canvas = '#333333'
	cattr.shadea = '#000000'
	cattr.shadeb = '#111111'
	cattr.mgrid = '#CCCCCC'
	cattr.axis = '#FFFFFF'
	cattr.frame = '#AAAAAA'
	cattr.font = '#FFFFFF'
	cattr.arrow = '#FFFFFF'

	graph = Graph(GRAPHFILE, start=now - width, end=now,
			vertical_label="'Memory usage'", color=cattr, base=1024)
	graph.data.extend([def1, def2, def3, def4,
		cdef1, cdef2, cdef3, cdef4,
		area1, area2, area3, area4])
	graph.write()


def main():
	"""
	Collect and generate RRD graphs.
	"""
	try:
		mode = sys.argv[1]
	except LookupError:
		print >> sys.stderr, '%s <mode>' % (sys.argv[0],)
		sys.exit(2)

	now = int(time.time()) - 1
	#now = 1374518800 # FIXME

	if mode == 'c':
		try:
			collect()
		except KeyboardInterrupt:
			pass
	elif mode == 'gc':
		graph_cpu(now, HOUR)
	elif mode == 'gm':
		graph_mem(now, HOUR)
	else:
		print >> sys.stderr, 'Unknown mode: %s' % (mode,)


if __name__ == '__main__':
	main()