#!/usr/bin/env python

from qt import *
from kdecore import *
import sys, os

def getLabel(blocks):
    """ Translates blocksize into human readable labels, such as 17.3 Gb, 2.1 Mb. """

    try:
        blocks = int(blocks) # 1K blocks now.
    except ValueError:
        return i18n("n/a")
    if blocks<1024:
        return i18n("%1 Kb").arg(blocks)
    if blocks<1024*1024:
        return i18n("%1 Mb").arg(round(float(blocks)/1024.0,1))
    blocks /= 1024
    if blocks<1024*1024:
        return i18n("%1 Gb").arg(round(float(blocks)/1024.0,1))
    blocks /= 1024
    return i18n("%1 Tb").arg(round(float(blocks)/1024.0,1))


class SizeViewApplication(QApplication):
    """ Boilerplate """
    def __init__(self,devicename,devicepath,args=[]):
        QApplication.__init__(self,args)

        self.maindialog = SizeView(None,devicename,devicepath)
        self.setMainWidget(self.maindialog)
        self.maindialog.show()
        self.exec_loop()     

class SizeView(QDialog):
    """ A SizeView represents a horizontal list of PartitionGroupWidgets.
        It supplies the code to read the sizes and the values that have 
        to be filled in, using the /proc filesystem and the program "df".
    """

    dev_path = "/dev/"      # Where to look for the partitions
    devicename = ""         # Such as hda1
    partitions = {}         # List of partitions
    sizes = {}              # Maps devicenames to blocksizes
    mountpoints = {}        # Maps devicenames to mountpoints
    used = {}               # Blocks used on a partition
    available = {}          # Blocks available
    part_types = {}         # Maps devicenames to partitiontypes
    partitionwidgets = []   # Holds a list of the PartitionGroup widgets

    def __init__(self,parent,devicename,devicepath=None):
        self.partitionwidgets = []
        QDialog.__init__(self,None,None,0,0)
        self.dialogtitle = i18n("Diskspace & Partitions")
        self.setCaption(self.dialogtitle)
        self.devicename = devicename
        if devicepath:
            self.dev_path = devicepath

        # Retrieve all information from the system.
        self.readMounts()
        self.readSize()
        self.readSwaps()

        partitions = self.partitions.keys()
        partitions.sort()

        number=1
        for part in partitions:
            try:
                fill = self.sizes[part]
                mountpoint = self.mountpoints[part]
                used = self.used[part]
                available = self.available[part]
            except KeyError:
                # Handles empty or not-mounted partitions
                fill = None
                mountpoint = i18n("n/a")
                used = str(i18n("n/a"))
                available = str(i18n("n/a"))

            pwidg = PartitionGroup(part,self,fill,number,self.part_types,self.dev_path)
            pwidg.setSize(self.partitions[part])
            pwidg.setMountPoint(mountpoint)
            pwidg.setUsed(used)
            pwidg.setAvailable(available)
            pwidg.setSizePolicy(QSizePolicy(QSizePolicy.MinimumExpanding,QSizePolicy.MinimumExpanding,0,0,
                pwidg.sizePolicy().hasHeightForWidth()))
            self.partitionwidgets.append(pwidg)
            number += 1

        n = len(partitions)
        r = 2
        c = 0
        cols = 1

        # Compute number of rows needed for partitions.
        if n%cols > 0:
            rows = int(n/cols)+1
        else:
            rows = int(n/cols)
        if n is 1: rows = 2

        # Build main Gridlayout.
        total_rows = rows+2
        self.grid = QGridLayout(self,total_rows,2,5)
        #self.setSizeGripEnabled(1)

        self.buttonCancel = QPushButton(i18n("Close"),self,"buttonCancel")
        self.buttonCancel.setAutoDefault(1)
        self.buttonCancel.setFixedWidth(80)
        self.grid.addWidget(self.buttonCancel,total_rows-1,1,Qt.AlignRight)

        self.grid.setRowStretch(0,0)
        self.grid.setRowStretch(total_rows-1,0)

        # Stretch all but first and last rows.
        for row in range(1,total_rows-1):
            self.grid.setRowStretch(row,5)

        self.clearWState(Qt.WState_Polished)
        self.connect(self.buttonCancel,SIGNAL("clicked()"),self.hide)

        #self.mainlabel = QLabel("<font size=+2><b>"+self.dialogtitle+"</b></font>",self)
        #self.grid.addWidget(self.mainlabel,0,0)

        self.diskgroup = DiskGroup(self,self.devicename,self.dev_path,self.partitions,self.totalsize,self.mountpoints)
        self.grid.addMultiCellWidget(self.diskgroup,1,1,0,1)

        for pw in self.partitionwidgets:
            self.grid.addWidget(pw,r,c)
            if c is cols:
                r += 1
                c = 0
            else:
                c += 1

    def readSize(self):
        fhandle = open("/proc/partitions","r")
        self.partitions = {}
        self.totalsize = 0
        for line in fhandle.readlines():
            try:
                major,minor,blocks,name = line.split()
                if name == self.devicename:
                    self.totalsize = blocks
                if name[:len(self.devicename)] == self.devicename and len(name) > len(self.devicename):
                    self.partitions[name] = blocks
            except ValueError:
                pass
        fhandle.close()

    def readMounts(self):
        fhandle = os.popen("/bin/df")
        for l in fhandle.readlines():
            v = l.split()
            try:
                p,s =  v[0].split("/")[2],v[4][:-1]
                self.sizes[p] = s
                self.mountpoints[p] = v[5]
                self.used[p] = v[2]
                self.available[p] = v[3]
                self.part_types[p] = "filesystem"
            except IndexError:
                pass
        fhandle.close()

    def readSwaps(self):
        fhandle = open("/proc/swaps")
        for line in fhandle.readlines():
            try:
                device,type,size,used,priority = line.split()
                device = device[len(self.dev_path):]
                self.used[device] = used
                self.sizes[device] = round(float(used)/float(size)*100 ,1)
                self.available[device] = str(int(size)-int(used))
                self.mountpoints[device] = "swap"
                self.part_types[device] = "swap"
            except:
                pass
        fhandle.close()

    """
    def __show__(self):
        print self.partitions
        print "Device", self.devicename, self.totalsize
        for p in self.partitions:
            print p, self.partitions[p], self.partitions[p]
    """
class DiskGroup(QGroupBox):
    """ Shows an overview of the physical layout of the disks, with the different partitions on it. """

    def __init__(self,parent,device,dev_path,partitions,totalsize,mountpoints):

        QGroupBox.__init__(self,parent,"DiskViewGroup")
        self.setTitle(i18n("Disk %1%2").arg(dev_path).arg(device))
        self.mountpoints = mountpoints
        self.partitions = partitions
        self.totalsize = totalsize

        self.setColumnLayout(0,Qt.Vertical)
        self.layout().setSpacing(6)
        self.layout().setMargin(11)
        DiskViewGroupLayout = QVBoxLayout(self.layout())
        DiskViewGroupLayout.setAlignment(Qt.AlignTop)
        colors = ["dark orange","dodger blue","gold","green","firebrick","navy","darkorange","darkblue"]

        self.diskview = DiskView(self,self.percentages(),colors)
        self.diskview.setScaledContents(1)
        DiskViewGroupLayout.addWidget(self.diskview)

        parts = self.partitions.keys()
        parts.sort()
        self.percentages()

        cols = 3    # Number of columns to use for colorlabels.
        rows = len(parts)/cols
        mod = len(parts)%cols
        if  mod > 0:
            rows += cols-mod

        # We multiply the number of cols by 3, first for the colorlabel, second for the name, third for spacing.
        cols = cols*3
        DiskViewPartitionListLayout = QGridLayout(DiskViewGroupLayout,rows,cols)

        i = cl = r = c = 0
        ps = ls = {}
        for dev in parts:
            ps[i] = LegendLabel(self,colors[cl])
            DiskViewPartitionListLayout.addWidget(ps[i],r,c)
            try:
                lbl = self.mountpoints[dev]
            except KeyError:
                lbl = "not mounted"
            ls[i] = QLabel(self,lbl+'<br /> ('+dev_path+dev+')',self)
            DiskViewPartitionListLayout.addWidget(ls[i],r,c+1)
            cl += 1
            if cl == len(colors):
                cl = 0
            i += 1
            if c is cols:
                c = 0
                r += 1
            else:
                c += 3

    def percentages(self):

        p_t = 0
        for p in self.partitions.values():
            p_t += int(p)

        self.perc = {}
        for p in self.partitions.keys():
            self.perc[p] = float(float(self.partitions[p])/float(p_t))
        return self.perc


class PartitionGroup(QGroupBox):
    """ Represents a groupbox with the filled bar and a couple of labels with
        information about the partition in it."""

    blocksize = 0
    title = str(i18n("Partition"))

    def __init__(self,device,parent,fill_percent,number,part_types,dev_path):
        QGroupBox.__init__(self,parent)
        self.part_types = part_types
        self.dev_path = dev_path
        self.setGeometry(QRect(110,100,370,203))
        self.setColumnLayout(0,Qt.Vertical)
        self.layout().setSpacing(3)
        self.layout().setMargin(5)
        self.setMinimumSize(280,120)

        partitiongroup_layout = QGridLayout(self.layout())
        partitiongroup_layout.setAlignment(Qt.AlignTop)
        self.available = QLabel(i18n("available"),self)
        partitiongroup_layout.addWidget(self.available,3,4)

        self.device = QLabel(i18n("device"),self)
        partitiongroup_layout.addMultiCellWidget(self.device,1,1,3,4)

        self.partpixmap = PartitionView(self,fill_percent,self.part_types,device)
        self.partpixmap.setScaledContents(1)

        partitiongroup_layout.addMultiCellWidget(self.partpixmap,0,0,0,4)
        self.textLabel1_3 = QLabel("textLabel1_3",self)
        partitiongroup_layout.addWidget(self.textLabel1_3,3,0)
        self.totalsize = QLabel("totalsize",self)
        partitiongroup_layout.addWidget(self.totalsize,2,1)
        self.textLabel1_2 = QLabel(self,"textLabel1_2")
        partitiongroup_layout.addWidget(self.textLabel1_2,2,0)
        self.textLabel1 = QLabel(self,"textLabel1")
        partitiongroup_layout.addWidget(self.textLabel1,1,0)
        self.textLabel3_2 = QLabel(self,"textLabel3_2")
        partitiongroup_layout.addMultiCellWidget(self.textLabel3_2,2,2,2,3)
        self.percentfilled = QLabel(self,"percentfree")
        partitiongroup_layout.addWidget(self.percentfilled,2,4)
        self.textLabel3_3 = QLabel(self,"textLabel3_3")
        partitiongroup_layout.addWidget(self.textLabel3_3,3,2)
        self.textLabel3 = QLabel(self,"textLabel3")
        partitiongroup_layout.addWidget(self.textLabel3,1,2)
        self.used = QLabel(self,"used")
        partitiongroup_layout.addWidget(self.used,3,1)
        self.mountpoint = QLabel(self,"mountpoint")
        self.mountpoint.setMinimumSize(QSize(60,0))
        partitiongroup_layout.addWidget(self.mountpoint,1,1)
        self.clearWState(Qt.WState_Polished)

        self.setTitle(i18n("%1. Partition").arg(number))
        self.textLabel1_3.setText(i18n("Used:"))
        self.textLabel1_2.setText(i18n("Total Size:"))
        self.textLabel1.setText(i18n("Mountpoint:"))
        self.textLabel3_2.setText(i18n("% Used:"))
        self.textLabel3_3.setText(i18n("Available:"))
        self.textLabel3.setText(i18n("Device:"))

        self.setDevice(self.dev_path+device)
        self.setFillPercentage(fill_percent)

    def setSize(self,label):
        self.totalsize.setText(getLabel(label))

    def setDevice(self,device):
        self.device.setText(device)

    def setMountPoint(self,mountpoint):
        self.mountpoint.setText(mountpoint)
        self.setTitle(i18n("Partition %1").arg(mountpoint))

    def setTotalSize(self,totalsize):
        self.totalsize.setText(getLabel(totalsize))

    def setFillPercentage(self,fill_percent):
        self.fill_percent = self.partpixmap.fill_percent = fill_percent
        if fill_percent is not None:
            self.percentfilled.setText("%s%%" % fill_percent)
        else:
            self.percentfilled.setText(i18n("Unknown"))

    def setUsed(self,used):
        self.used.setText(getLabel(used))

    def setAvailable(self,available):
        self.available.setText(getLabel(available))

class LegendLabel(QLabel):
    """ Show some color in the DiskView legend """

    def __init__(self,parent,color="green",style=QBrush.SolidPattern):
        QLabel.__init__(self,parent,"bla")
        self.w = 40
        self.h = 20
        self.pmsize = QSize(self.w,self.h)
        self.pm = QPixmap(self.pmsize)
        self.linewidth = 2
        self.color = QColor(color)
        self.style = style
        self.framecolor = QColor("black")
        self.paintMe()
        self.setPixmap(self.pm)
        self.setScaledContents(1)
        self.setSizePolicy(QSizePolicy(QSizePolicy.Fixed,QSizePolicy.Fixed,0,0,
            self.sizePolicy().hasHeightForWidth()))

    def paintMe(self):
        p = QPainter(self.pm)
        p.fillRect(0,0,self.w,self.h,QBrush(self.color,self.style))
        p.setPen(QPen(QColor("black"),self.linewidth))
        p.drawRect(self.linewidth/2,self.linewidth/2,self.w-self.linewidth/2,self.h-self.linewidth/2)
        p.end()

class PartitionView(QLabel):
    """ PartitionView is a label carryig a pixmap. This class's main purpose is handlig layout
        of the underlying pixmap."""
    w = 250
    h = 35
    def __init__(self,parent,fill_percent,part_types,device):
        self.part_types = part_types
        self.fill_percent = fill_percent
        QLabel.__init__(self,parent,"pview")
        self.setSizePolicy(QSizePolicy(QSizePolicy.Expanding,QSizePolicy.Expanding,0,0,
            self.sizePolicy().hasHeightForWidth()))
        self.setMinimumSize(QSize(self.w,self.h))        
        self.setPixmap(PartitionPixmap(QSize(self.w,self.h),self.fill_percent,self.part_types,device))
        self.setScaledContents(1)
        self.setAlignment(QLabel.AlignCenter)

class DiskView(PartitionView):
    """ PartitionView is a label carryig a pixmap. This class's main purpose is handlig layout
        of the underlying pixmap."""

    w = 540
    h = 50
    linewidth = 2

    def __init__(self,parent,percents,colors):
        QLabel.__init__(self,parent)

        self.setSizePolicy(QSizePolicy(QSizePolicy.Expanding,QSizePolicy.Expanding,0,0,
            self.sizePolicy().hasHeightForWidth()))
        self.setPixmap(DiskPixmap(percents,colors,(self.w,self.h)))
        self.setScaledContents(1)
        self.setAlignment(QLabel.AlignCenter)

class DiskPixmap(QPixmap):

    linewidth = 2 # Width of surrounding frame

    def __init__(self,percents,colors,(w,h)):
        self.percents = percents
        self.w,self.h = w,h
        self.colors = colors
        QPixmap.__init__(self,w,h)
        self.paintMe()

    def paintMe(self):
        p = QPainter(self)
        w,h = self.w,self.h
        i = 0
        x0 = 0
        y = 0

        # Paint background, this is interesting for empty partitions.
        p.fillRect(0,0,w,h,QBrush(QColor("white")))

        parts = self.percents.keys()
        parts.sort()
        xa = wa = 0
        for part in parts:
            W = (w * self.percents[part])
            # We need to adjust a little to avoid to get wholes.
            if x0>0: xa = 2
            if W < self.w: wa = 2
            p.fillRect(x0-xa,0,W+wa,h,QBrush(QColor(self.colors[i])))
            i += 1
            x0 += W

        # Paint Frame around it.
        p.setPen(QPen(QColor("black"),self.linewidth))
        p.drawRect(self.linewidth/2,self.linewidth/2,self.width()-self.linewidth/2,self.height()-self.linewidth/2)
        p.end()


class PartitionPixmap(QPixmap):
    """ A PartitionPixmap is a two colored bar with a black frame. The first color represents the 
        percentage that's used, the second one the free percentage."""
    linewidth = 2 # Width of surrounding frame

    def __init__(self,pmsize,fill_percent,part_types,device):
        QPixmap.__init__(self,pmsize)

        self.pmsize = pmsize # Size of the pixmap
        self.part_types = part_types # Array to look up the type of the partition
        self.fill_percent = fill_percent
        self.device = device # Device name of the partition

        self.w = self.pmsize.width()
        self.h = self.pmsize.height()
        self.paintMe()

    def paintMe(self):
        p = QPainter(self)
        try:
            fill_percent = int(self.fill_percent)
            if self.part_types[self.device] == "swap":
                # Swap partitions get blueish colors.
                color_used = QColor("blue")
                color_free = QColor("lightblue")
            else:
                # Regular partitions get a red / green color combo.
                color_used = QColor("red")
                color_free = QColor("forest green")
        except (KeyError,TypeError):
            # Partition has no fillsize, might be empty or not mounted partition
            p.fillRect(0,0,self.w,self.h,QBrush(QColor("darkgrey")))
            p.setPen(QPen(QColor("black"),self.linewidth))
            p.drawRect(self.linewidth/2,self.linewidth/2,self.w-self.linewidth/2,self.h-self.linewidth/2)
            p.end()
            return
        # Total width of the pixmap
        W,H = float(self.w),float(self.h)

        # Paint filled == red part of the bar.
        x = y = 0
        w = W - (W*(1-(fill_percent/100.00)))
        h = H
        p.fillRect(x,y,w,h,QBrush(color_used))

        # Paint green part == space left
        x = w
        w = W - w
        p.fillRect(x,y,w,h,QBrush(color_free))

        # Paint Frame around it.
        p.setPen(QPen(QColor("black"),self.linewidth))
        p.drawRect(self.linewidth/2,self.linewidth/2,W-self.linewidth/2,H-self.linewidth/2)

        p.end()

if __name__ == "__main__":
    device = "sdc"
    app = SizeViewApplication(device,None,sys.argv)