# -*- coding: utf-8 -*-

#                _              _          _
#               (_)            | |        | |
#  _ __   _   _  _   ___   ___ | |_  __ _ | |_
# | '_ \ | | | || | / _ \ / __|| __|/ _` || __|
# | |_) || |_| || || (_) |\__ \| |_| (_| || |_
# | .__/  \__, ||_| \___/ |___/ \__|\__,_| \__|
# | |      __/ |
# |_|     |___/
#

"""
一个使用 Python 编写的 iostat 克隆版本
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

usage:

   >>> from pyiostat import *
   >>> compare_with_iostat('sda', 10, unit='kb')
"""

import time
import pickle
import commands
from datetime import datetime
from collections import namedtuple
from pprint import pprint

__title__ = 'pyiostat'
__version__ = '0.0.1'
__author__ = 'derekqdong'

PARTITIONS = "/proc/partitions"
DISKSTATS = "/proc/diskstats"

DiskStat = namedtuple('DiskStat', 'major, minor, dev, r_ios, r_merges, '
                                  'r_sec, r_ticks, w_ios, w_merges, '
                                  'w_sec, w_ticks, ios_pgr, tot_ticks, '
                                  'rq_ticks')


class IOStatError(Exception):
    """本模块的根异常"""


class UnitError(IOStatError):
    """测量单位出错，必须为 sector, kb, mb"""


class EmptyPreviousData(IOStatError):
    """模块第一次运行，没有相关数据，这是正常情况，抛出异常是为了通知调用模块这一信息"""


class IOStatus:
    is_firt_time = True
    pre_diskstats = None
    curr_diskstats = None

    def __init__(self, unit='mb', ignore_nbd=True):
        """根据分区信息获取磁盘列表
        curr 中保存本次磁盘信息，prev 中保存上次调用时保存的磁盘信息
        """

        # 貌似 vsQemuDiskMonitor 会把所有的磁盘都统计
        # self._devices = self.get_devices()

        if unit not in ('sector', 'kb', 'mb'):
            raise UnitError('unit must be sector, kb, mb')
        self._unit = unit
        self._ignore_nbd = ignore_nbd

        IOStatus.pre_diskstats = IOStatus.curr_diskstats
        IOStatus.curr_diskstats = self.get_disk_stats()

    def get_iostat(self, device=None):
        """获取 iostat -x 风格的磁盘 IO 数据
        根据两次的差值计算出 iostat
         
        Exception: 如果是第一次运行，IOStatus.pre_diskstats还没有数据，
        默认抛出异常，由调用者来决定如何处理。
        """
        if IOStatus.is_firt_time:
            IOStatus.is_firt_time = False
            raise EmptyPreviousData('First run, run it again to get iostatus.'
                                    ' See you later.')

        iostat = {}
        for dev in IOStatus.curr_diskstats.keys():
            if device and dev != device:
                continue
            try:
                res = self.compute_iostat(IOStatus.curr_diskstats[dev], IOStatus.pre_diskstats[dev])
            except KeyError as e:
                print 'device %s is missing, may be mounted or unmounted.' % e
            else:
                iostat[dev] = res

        return iostat

    @staticmethod
    def get_devices():
        """获取本地挂载的设备信息，只包含磁盘，去除分区"""
        with open(PARTITIONS) as f:
            partitions = f.readlines()[2:]
            return set(line.split()[-1] for line in partitions
                       if not line.strip()[-1].isdigit())

    def get_disk_stats(self):
        """从内核文件中读取磁盘的 io 信息"""

        def parse_diskstat_line(line):
            fields = line.split()[0:14]
            stats = DiskStat._make(fields)._asdict()

            for key, val in stats.items():
                stats[key] = IOStatus.to_num(val)
                stats['ts'] = time.time()
            return stats

        with open(DISKSTATS) as f:
            lines = f.readlines()
            all_stats = [parse_diskstat_line(line) for line in lines]
            if self._ignore_nbd:
                all_stats = filter(lambda x: not x['dev'].startswith("nbd"), all_stats)
            disk_stats = {}
            for stat in all_stats:
                # if stat['dev'] in self._devices:
                disk_stats[stat['dev']] = stat
            return disk_stats

    def compute_iostat(self, curr, prev):
        """计算 iostat 各项参数，由 time_delta 这段时间作为差值
        :rrqm/s: (rd_merges[1] - rd_merges[0])
        :wrqm/s: (wr_merges[1] - wr_merges[0])
        :r/s: (rd_ios[1] - rd_ios[0])
        :w/s: (wr_ios[1] - wr_ios[0])
        :rMB/s: (rd_sectors[1] - rd_sectors[0]) * sector_size
        :wMB/s: (wr_sectors[1] - wr_sectors[0]) * sector_size
        :avgrq-sz: ((rd_sectors[1] - rd_sectors[0]) + 
                    (wr_sectors[1] - wr_sectors[0]))
                    / (rd_ios[1] - rd_ios[0]) + (wr_ios[1] - wr_ios[0])
        :avgqu-sz : (rq_ticks[1] - rq_ticks[0]) / 1000
        :await: ((rd_ticks[1] - rd_ticks[0]) + (wr_ticks[1] - wr_ticks[0])) 
                    / (rd_ios[1] - rd_ios[0]) + (wr_ios[1] - wr_ios[0])
        :svctm: util / (rd_ios[1] - rd_ios[0]) + (wr_ios[1] - wr_ios[0])
        :util: (tot_ticks[1] - tot_ticks[0]) / 1000 * 100
        """

        SECTOR_SIZE = 512

        def diff(field):
            return (curr[field] - prev[field]) / (curr["ts"] - prev["ts"])

        stat = {}
        stat['rrqm/s'] = diff('r_merges')
        stat['wrqm/s'] = diff('w_merges')
        stat['r/s'] = diff('r_ios')
        stat['w/s'] = diff('w_ios')

        if self._unit == 'mb':
            stat['rMB/s'] = diff('r_sec') * SECTOR_SIZE / 2 ** 20
            stat['wMB/s'] = diff('w_sec') * SECTOR_SIZE / 2 ** 20
        elif self._unit == 'kb':
            stat['rkB/s'] = diff('r_sec') * SECTOR_SIZE / 1024
            stat['wkB/s'] = diff('w_sec') * SECTOR_SIZE / 1024
        else:
            stat['rsec/s'] = diff('r_sec')
            stat['wsec/s'] = diff('w_sec')

        stat['avgqu-sz'] = diff('rq_ticks') / 1000
        stat['util'] = diff('tot_ticks') / 10

        if diff('r_ios') + diff('w_ios') > 0:
            stat['avgrq-sz'] = (diff('r_sec') + diff('w_sec')) / (
                diff('r_ios') + diff('w_ios'))
            stat['await'] = (diff('r_ticks') + diff('w_ticks')) / (
                diff('r_ios') + diff('w_ios'))
            stat['svctm'] = diff('tot_ticks') / (diff('r_ios') + diff('w_ios'))
        else:
            stat['avgrq-sz'] = 0
            stat['await'] = 0
            stat['svctm'] = 0

        return stat

    @staticmethod
    def to_num(n):
        if n.isdigit():
            return int(n)
        return n