#!/usr/bin/python

import subprocess
import sys
import re
import pprint

def humanCapasity(b):
    b = float(b)
    G = ['%.0f B', '%.2f KB', '%.2f MB', '%.2f GB', '%.2f TB']
    for p in G:
        if b < 1024:
            return p % b
        else:
            b /= 1024
    raise 'Too large hard drive ;-)'


class Disk(object):
    # fdisk -G /dev/rdsk/c0t0d0p0
    # * Physical geometry for device /dev/rdsk/c0t0d0p0
    # * PCYL     NCYL     ACYL     BCYL     NHEAD NSECT SECSIZ
    #   2088     2088     0        0        255   63    512
    _fdiskG = re.compile('\s*(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s*')

    # fdisk -W - /dev/rdsk/c0t0d0p0
    #* Id    Act  Bhead  Bsect  Bcyl    Ehead  Esect  Ecyl    Rsect      Numsect
    #  191   128  65     2      0       168    44     380     4096       6111232   
    #  0     0    0      0      0       0      0      0       0          0         
    #  0     0    0      0      0       0      0      0       0          0         
    #  0     0    0      0      0       0      0      0       0          0         
    _fdiskW = re.compile('\s*(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s*')

    _name = None
    _desc = None
    _geometry = {}
    _capacity = 0
    _capacity_human = None
    _partitions = []
    _raw_device = None

    @property
    def name(self):
        return self._name

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, desc):
        self._desc = desc
    
    def read_partitions():
        self._partitions = {}
        out = subprocess.Popen(['fdisk', '-W', '-', self._raw_device],
                stdout=subprocess.PIPE).stdout
        for line in out:
            m = self._fdiskW.match(line)
            if m:
                self._partitions.append({
                    'id'      : int(m.group(1)),
                    'rsect'   : int(m.group(9)),
                    'numsect' : int(m.group(10)),
                    'cap'     : humanCapasity(int(m.group(10))*self._geometry['secsiz']),
                    })

    @name.setter
    def name(self, name):
        """ Set disk name and update disk's info: capacity, geometry etc."""
        if name != self._name:
            self._raw_device = '/dev/rdsk/{}p0'.format(name)
            self._name = name
            out = subprocess.Popen(['fdisk', '-G', self._raw_device],
                    stdout=subprocess.PIPE).stdout
            for line in out:
                m = self._fdiskG.match(line)
                if m:
                    self._geometry = {
                            'pcyl'   : int(m.group(1)),
                            'ncyl'   : int(m.group(2)),
                            'acyl'   : int(m.group(3)),
                            'bcyl'   : int(m.group(4)),
                            'nhead'  : int(m.group(5)),
                            'nsect'  : int(m.group(6)),
                            'secsiz' : int(m.group(7)),
                            }
            self._capacity = ( self._geometry['ncyl'] * self._geometry['nhead'] *
                             self._geometry['nsect'] * self._geometry['secsiz'] )
            assert self._capacity != 0
            self._capacity_human = humanCapasity(self._capacity)

    @property
    def geometry(self):
        return self._geometry

    @property
    def partitions(self):
        return self._partitions

    @property
    def capacity(self):
        """ Capacity in bytes """
        return self._capacity

    @property
    def cap(self):
        """ Human readable capacity string, e. g. 30 GB """
        return self._capacity_human

    def __init__(self, name, desc=None):
        self.name = name
        self.desc = desc
        pass