import os, time
import whatap.util.process_util as process_util

proc_diskstats='/proc/diskstats'

filesystems = { "btrfs": True, "ext2": True, "ext3": True, "ext4": True, "reiser": True, "xfs": True,
        "ffs": True, "ufs": True, "jfs": True, "jfs2": True, "vxfs": True, "hfs": True, "ntfs": True, "fat32": True,
        "zfs": True, "refs": True, "nfs": True, "nfs2": True, "nfs3": True, "nfs4": True, "cifs": True, "ocfs2": True }

class _DiskPerf:
    def __init__(self):
        self.mountPoint = None
        self.deviceId = None
        self.fileSystem = None
        self.totalSpace = None
        self.usedSpace = None
        self.usedPercent = None
        self.freeSpace = None
        self.freePercent = None
        self.blksize = None
        self.readIops = 0
        self.writeIops = 0
        self.readBps = 0
        self.writeBps = 0
        self.ioPercent = 0
        self.queueLength = 0


def getDisk():
    if not os.path.exists('/etc/mtab'):
        return None
    f = open('/etc/mtab','r')
    if not f:
        return None
    try:
        diskQuotas = []
        for l in f.readlines():
            words = l.split()
            if words[2] in filesystems:
                deviceId = parseDeviceIdByUUID(words[0])
                p = _DiskPerf()
                p.mountPoint = words[1]
                p.deviceId = deviceId
                p.fileSystem = words[2]
                f_bsize, f_frsize, f_blocks, f_bfree, f_bavail, f_files, f_ffree, f_favail, f_flag, f_namemax = os.statvfs(p.mountPoint)
                p.totalSpace = int(f_bsize * f_blocks)
                p.usedSpace = int(f_bsize * (f_blocks - f_bfree))
                p.usedPercent = float(100.0) - float(100*f_bavail) / float(f_blocks - f_bfree + f_bavail)
                p.freeSpace = int(f_bsize * f_bavail )
                p.freePercent = float(100.0) * float(f_bavail) / float(f_blocks - f_bfree + f_bavail)

                p.blksize = int(f_bsize)
                diskQuotas.append(p)
    finally:
        f.close()

    return parseDiskIO(diskQuotas)        

disk_sector_size_lookup = {}
def getDiskSectorSize(logicalDisk):
    global disk_sector_size_lookup
    sector_size = 0
    if logicalDisk not in disk_sector_size_lookup:
        realDisk = getMappedDeviceId(logicalDisk)
        if realDisk == logicalDisk:
            realDisk = os.path.join('/dev',realDisk)
        else:
            realDisk = os.path.join('/dev/mapper', realDisk)
        try:

            #print 'getDiskSectorSize step -1 ',"blockdev --getbsz %s"%(realDisk)
            sector_size = int(process_util.executeCmd("blockdev --getbsz %s"%(realDisk)))
        except:
            pass
        disk_sector_size_lookup[logicalDisk] = sector_size

    return disk_sector_size_lookup[logicalDisk]

def parseDeviceIdByUUID(deviceID):
    try:
        return os.readlink(deviceID)
    except:
        return deviceID

    return deviceID


def getMappedDeviceId(deviceId):
    sysblock_filepath = os.path.join('/sys/block',deviceId,'dev')
    if os.path.exists(sysblock_filepath):
        f = open(sysblock_filepath)
        deviceContent = f.read()
        #print 'opened ', sysblock_filepath,deviceContent

        deviceMajor, deviceMinor = deviceContent.split(':')
        f.close()
    else:
        return deviceId

    searchDir = '/dev/mapper'
    for filename in os.listdir(searchDir):
        linkpath = os.path.join(searchDir, filename)
        try:
            st = os.stat(linkpath)
            targetMajor = os.major(st.st_rdev)
            targetMinor = os.minor(st.st_rdev)
            #print deviceId, filename, int(deviceMajor) == targetMajor and int(deviceMinor) == targetMinor
            if int(deviceMajor) == targetMajor and int(deviceMinor) == targetMinor:
                return filename
        except Exception, e:
            pass
            #print e
    return deviceId

def getCpuGuestJiff():
    f = open('/proc/uptime','r')
    try:
        for l in f.readlines():
            words = l.split()
            numbers = words[0].split('.')
            seconds = int(numbers[0])
            millis = int(numbers[1])

            return seconds*100 + millis
    finally:
        f.close()

class _DiskIORaw:
    def __init__(self):
        self.deviceID= None
        self.readIoCount= None
        self.readIoByteCount= None
        self.writeIoCount= None
        self.writeIoByteCount= None
        self.timestamp= None
        self.ioMillis= None
        self.avgQSize= None
        self.guestjiff= None

oldDiskIOMap = None
def parseDiskIO(diskQuotas):
    global oldDiskIOMap
    diskPerfMap = {}
    for p in diskQuotas:
        diskPerfMap[p.deviceId.split('/')[-1]] = p
    diskIOMap = {}
    timestamp = int(time.time() * 1000000000)
    guestjiff = getCpuGuestJiff()

    if not os.path.exists(proc_diskstats):
        return diskQuotas
    f = open(proc_diskstats,'r')
    if not f:
        return diskQuotas
    try:
        for l in f.readlines():
            words = l.split()
            if len(words) < 14:
                continue
            deviceID = words[2]
            sectorSize = getDiskSectorSize(deviceID)
            if deviceID not in diskPerfMap:
                deviceID= getMappedDeviceId(deviceID)
                deviceID= deviceID.split('/')[-1]
            #print 'disk_linux.parseDiskIO step -1', deviceID, deviceID in diskPerfMap

            if deviceID in diskPerfMap:
                d = _DiskIORaw()
                d.deviceID= deviceID
                d.readIoCount = int(words[3])
                d.readIoByteCount = int(words[5]) * sectorSize
                d.writeIoCount = int(words[7])
                d.writeIoByteCount = int(words[9]) * sectorSize
                d.ioMillis = int(words[12])
                d.avgQSize = int(words[13])
                d.timestamp = timestamp
                d.guestjiff = guestjiff
                diskIOMap[deviceID] = d

                if oldDiskIOMap and deviceID in oldDiskIOMap:

                    oldd = oldDiskIOMap[deviceID]
                    p = diskPerfMap[deviceID]
                    p.readIops = float( d.readIoCount - oldd.readIoCount ) / float( d.timestamp - oldd.timestamp ) * 1000000000.0
                    p.writeIops = float( d.writeIoCount - oldd.writeIoCount ) / float( d.timestamp - oldd.timestamp ) * 1000000000.0
                    p.readBps = float( d.readIoByteCount - oldd.readIoByteCount ) / float( d.timestamp - oldd.timestamp ) * 1000000000.0
                    p.writeBps = float( d.writeIoByteCount - oldd.writeIoByteCount ) / float( d.timestamp - oldd.timestamp ) * 1000000000.0
                    p.ioPercent = float( d.ioMillis - oldd.ioMillis ) / float( d.timestamp - oldd.timestamp ) * 1000000.0
                    p.queueLength = float( d.avgQSize - oldd.avgQSize ) / float( d.guestjiff - oldd.guestjiff ) * 100.0 / 1000.0

                    #print 'disk_linux.parseDiskIO step -2', deviceID, 'read', p.readBps ,  d.readIoByteCount - oldd.readIoByteCount, 'write',p.writeBps, d.writeIoByteCount - oldd.writeIoByteCount
        oldDiskIOMap = diskIOMap
        return diskQuotas
    finally:
        f.close()

def test():
    getDisk()
    time.sleep(5)
    for diskperf in getDisk():
        print [item for item in vars(diskperf).items()]

if __name__ == '__main__':
    test()
