#!/bin/sh -e
#
#    disk_io: calculate the disk io rate
#    Copyright (C) 2010 Canonical Ltd.
#
#    Authors: Dustin Kirkland <kirkland@canonical.com>
#
#    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 Free Software Foundation, version 3 of the License.
#
#    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.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

PKG="byobu"
DATA="$HOME/.$PKG"
color 2>/dev/null || color() { true; }
[ "$UTF8" = "1" ] && ICON_RD="◀" || ICON_RD="<"
[ "$UTF8" = "1" ] && ICON_WR="▶" || ICON_WR=">"

# Default to disk providing /, but let users override with MONITORED_DISK
[ -z "$MONITORED_DISK" ] && MP="/" ||  MP="$MONITORED_DISK"
case $MP in
	/dev/*) disk="$MP" ;;
	*) disk=$(grep -m 1 " $MP " /etc/mtab | sed -e "s: .*$::") ;;
esac
which greadlink 2>/dev/null && READLINK="greadlink" || READLINK="readlink"
disk=$($READLINK -f "$disk" | sed -e "s: .*$::" -e "s:^.*/::" -e "s:\([hsv]d[a-z]\)[0-9]*$:\1:")

if [ "$1" = "--detail" ]; then
	if which iostat >/dev/null; then
		iostat -d -m -h
	else
		echo "Please install iostat if you want detailed information on your disk throughput"
	fi
	exit 0
fi

[ -d "/var/run/screen/S-$USER" ] && DIR="/var/run/screen/S-$USER" || DIR="$DATA"
t2=`date +%s`
for i in "read" "write"; do
	cache="$DIR/$PKG.disk_$i"
	t1=`stat -c %Y "$cache"` 2>/dev/null || t1=0
	unit="kB/s"
	if [ $t2 -le $t1 ]; then
		rate=0
	else
		x1=`cat "$cache"` 2>/dev/null || tx1=0
		if [ "$i" = "read" ]; then
			symbol="$ICON_RD"
			x2=`awk '{print $3}' /sys/block/$disk/stat`
		else
			symbol="$ICON_WR"
			x2=`awk '{print $7}' /sys/block/$disk/stat`
		fi
		echo "$x2" > "$cache"
		rate=`echo "$t1" "$t2" "$x1" "$x2" | awk '{printf "%.0f", ($4 - $3) / ($2 - $1) * 512 / 1024 }'`
		if [ "$rate" -lt 0 ]; then
			rate=0
		elif [ "$rate" -gt 1024 ]; then
			rate=`echo "$rate" | awk '{printf "%.1f", $1/1024}'`
			unit="MB/s"
		fi
	fi
	[ "$rate" != "0" ] || continue
	printf "$(color b M W)$symbol$rate$(color -)$(color M W)$unit$(color -) "
done
