#!/usr/bin/python
#
# Copyright (C) 2014 Red Hat.
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Iostat Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
# for more details.
#
# pylint: disable=C0103,R0914,R0902
""" Display disk and device-mapper I/O statistics """

import sys
from pcp import pmapi, pmcc
from cpmapi import PM_TYPE_U64, PM_CONTEXT_ARCHIVE, PM_SPACE_KBYTE

IOSTAT_SD_METRICS = [ "disk.dev.read", "disk.dev.read_bytes",
	         "disk.dev.write",  "disk.dev.write_bytes",
		 "disk.dev.read_merge", "disk.dev.write_merge",
		 "disk.dev.blkread", "disk.dev.blkwrite",
		 "disk.dev.read_rawactive", "disk.dev.write_rawactive",
		 "disk.dev.avactive"]

IOSTAT_DM_METRICS = [ "disk.dm.read", "disk.dm.read_bytes",
	         "disk.dm.write",  "disk.dm.write_bytes",
		 "disk.dm.read_merge", "disk.dm.write_merge",
		 "disk.dm.blkread", "disk.dm.blkwrite",
		 "disk.dm.read_rawactive", "disk.dm.write_rawactive",
		 "disk.dm.avactive"]

class IostatReport(pmcc.MetricGroupPrinter):
    Hcount = 0
    def timeStampDelta(self, group):
	c = 1000000.0 * group.timestamp.tv_sec + group.timestamp.tv_usec
	p = 1000000.0 * group.prevTimestamp.tv_sec + group.prevTimestamp.tv_usec
	return (c - p) / 1000000.0

    def instlist(self, group, name):
	return dict(map(lambda x: (x[1], x[2]), group[name].netValues)).keys()

    def curVals(self, group, name):
	return dict(map(lambda x: (x[1], x[2]), group[name].netValues))

    def prevVals(self, group, name):
	return dict(map(lambda x: (x[1], x[2]), group[name].netPrevValues))

    def report(self, manager):
	if "dm" in IostatOptions.xflag:
	    subtree = "disk.dm"
	else:
	    subtree = "disk.dev"
	group = manager["iostat"]

	if group[subtree + ".read_merge"].netPrevValues == None:
	    # need two fetches to report rate converted counter metrics
	    return

	instlist = self.instlist(group, subtree + ".read")
	dt = self.timeStampDelta(group)
	timestamp = group.contextCache.pmCtime(long(group.timestamp)).rstrip()

	c_rrqm = self.curVals(group, subtree + ".read_merge")
	p_rrqm = self.prevVals(group, subtree + ".read_merge")

	c_wrqm = self.curVals(group, subtree + ".write_merge")
	p_wrqm = self.prevVals(group, subtree + ".write_merge")

	c_r = self.curVals(group, subtree + ".read")
	p_r = self.prevVals(group, subtree + ".read")

	c_w = self.curVals(group, subtree + ".write")
	p_w = self.prevVals(group, subtree + ".write")

	c_rkb = self.curVals(group, subtree + ".read_bytes")
	p_rkb = self.prevVals(group, subtree + ".read_bytes")

	c_wkb = self.curVals(group, subtree + ".write_bytes")
	p_wkb = self.prevVals(group, subtree + ".write_bytes")

	c_ractive = self.curVals(group, subtree + ".read_rawactive")
	p_ractive = self.prevVals(group, subtree + ".read_rawactive")

	c_wactive = self.curVals(group, subtree + ".write_rawactive")
	p_wactive = self.prevVals(group, subtree + ".write_rawactive")

	c_avactive = self.curVals(group, subtree + ".avactive")
	p_avactive = self.prevVals(group, subtree + ".avactive")

	# check availability
	if p_rrqm == {} or p_wrqm == {} or p_r == {} or p_w == {} or p_rkb == {} \
	   or p_wkb == {} or p_ractive == {} or p_wactive == {} or p_avactive == {}:
	    # no values (near start of archive?)
	    return

	if "h" not in IostatOptions.xflag:
	    self.Hcount += 1
	    if self.Hcount == 24:
		self.Hcount = 1
	    if self.Hcount == 1:
		if "t" in IostatOptions.xflag:
		    heading = ('# Timestamp', 'Device', 'rrqm/s', 'wrqm/s', 'r/s', 'w/s', 'rkB/s', 'wkB/s',
			       'avgrq-sz', 'avgqu-sz', 'await', 'r_await', 'w_await', '%util')
		    print "%-24s %-12s %7s %7s %6s %6s %8s %8s %8s %8s %7s %7s %7s %5s" % heading
		else:
		    heading = ('# Device', 'rrqm/s', 'wrqm/s', 'r/s', 'w/s', 'rkB/s', 'wkB/s',
			       'avgrq-sz', 'avgqu-sz', 'await', 'r_await', 'w_await', '%util')
		    print "%-12s %7s %7s %6s %6s %8s %8s %8s %8s %7s %7s %7s %5s" % heading

	for inst in instlist:
	    # basic stats
	    rrqm = (c_rrqm[inst] - p_rrqm[inst]) / dt
	    wrqm = (c_wrqm[inst] - p_wrqm[inst]) / dt
	    r = (c_r[inst] - p_r[inst]) / dt
	    w = (c_w[inst] - p_w[inst]) / dt
	    rkb = (c_rkb[inst] - p_rkb[inst]) / dt
	    wkb = (c_wkb[inst] - p_wkb[inst]) / dt

	    # totals
	    tot_rios = (float)(c_r[inst] - p_r[inst])
	    tot_wios = (float)(c_w[inst] - p_w[inst])
	    tot_ios = (float)(tot_rios + tot_wios)

	    # total active time in seconds (same units as dt)
	    tot_active = (float)(c_avactive[inst] - p_avactive[inst]) / 1000.0

	    avgrqsz = avgqsz = await = r_await = w_await = util = 0.0

	    # average request size units are KB (sysstat reports in units of sectors)
	    if tot_ios:
		avgrqsz = (float)((c_rkb[inst] - p_rkb[inst]) + (c_wkb[inst] - p_wkb[inst])) / tot_ios

	    # average queue length
	    avgqsz = (float)((c_ractive[inst] - p_ractive[inst]) + (c_wactive[inst] - p_wactive[inst])) / dt / 1000.0

	    # await, r_await, w_await
	    if tot_ios:
		await = ((c_ractive[inst] - p_ractive[inst]) + (c_wactive[inst] - p_wactive[inst])) / tot_ios

	    if tot_rios:
		r_await = (c_ractive[inst] - p_ractive[inst]) / tot_rios

	    if tot_wios:
		w_await = (c_wactive[inst] - p_wactive[inst]) / tot_wios

	    # device utilization (percentage of active time / interval)
	    if tot_active:
	    	util = 100.0 * tot_active / dt

	    if "t" in IostatOptions.xflag:
		print "%-24s %-12s %7.1f %7.1f %6.1f %6.1f %8.1f %8.1f %8.2f %8.2f %7.1f %7.1f %7.1f %5.1f" \
		% (timestamp, inst, rrqm, wrqm, r, w, rkb, wkb, avgrqsz, avgqsz, await, r_await, w_await, util)
	    else:
		print "%-12s %7.1f %7.1f %6.1f %6.1f %8.1f %8.1f %8.2f %8.2f %7.1f %7.1f %7.1f %5.1f" \
		% (inst, rrqm, wrqm, r, w, rkb, wkb, avgrqsz, avgqsz, await, r_await, w_await, util)


class IostatOptions(pmapi.pmOptions):
    # class attributes
    xflag = [] 

    def extraOptions(self, opt, optarg, index):
	if opt == "x":
	    IostatOptions.xflag = optarg.replace(',', ' ').split(' ')
	else:
	    print "Warning: option '", opt, "' not recognised"

    def __init__(self):
	pmapi.pmOptions.__init__(self, "A:a:D:h:O:S:s:T:t:VZ:z?x:")
	self.pmSetLongOptionHeader("General options:")
	self.pmSetLongOptionAlign()
	self.pmSetLongOptionArchive()
	self.pmSetLongOptionDebug()
	self.pmSetLongOptionHost()
	self.pmSetLongOptionOrigin()
	self.pmSetLongOptionStart()
	self.pmSetLongOptionSamples()
	self.pmSetLongOptionFinish()
	self.pmSetLongOptionInterval()
	self.pmSetLongOptionVersion()
	self.pmSetLongOptionTimeZone()
	self.pmSetLongOptionHostZone()
	self.pmSetLongOptionHelp()
	self.pmSetOptionCallback(self.extraOptions)
	self.pmSetLongOptionText("  -x comma separated extended options: [dm][,t][,h]")
	self.pmSetLongOptionText("     dm show device-mapper statistics (default is sd devices)")
	self.pmSetLongOptionText("     t precede every line with a timestamp in ctime format");
	self.pmSetLongOptionText("     h suppress headings");

if __name__ == '__main__':
    try:
	manager = pmcc.MetricGroupManager.builder(IostatOptions(), sys.argv)
	if "dm" in IostatOptions.xflag :
	    manager["iostat"] = IOSTAT_DM_METRICS
	else:
	    manager["iostat"] = IOSTAT_SD_METRICS
	manager.printer = IostatReport()
        sts = manager.run()
        sys.exit(sts)
    except pmapi.pmErr, error:
	print '%s: %s\n' % (error.progname(), error.message())
    except pmapi.pmUsageErr, usage:
	print usage.message()
        sys.exit(1)
    except KeyboardInterrupt:
	pass