#! /usr/bin/python

import sys, os, string, types, commands
from optparse import OptionParser

from math import log

from statvfs import *

def get_terminal_width():
    """
    this does not work:
    if os.environ.has_key('COLUMNS'):
        return int(os.environ['COLUMNS'])
    else:
        return 80
    """
    c = commands.getoutput('resize').split('\n')
    c = [x for x in c if x.startswith('COLUMNS=')]
    if c:
        c = c[0]
        dummy, c = c.split('=', 1)
        if c[-1]==';':
            c = c[:-1]
    if c:
        return int(c)
    else:
        return 80


#some default definitions
colours = {    
            'none'       :    "",
            'default'    :    "\033[0m",
            'bold'       :    "\033[1m",
            'underline'  :    "\033[4m",
            'blink'      :    "\033[5m",
            'reverse'    :    "\033[7m",
            'concealed'  :    "\033[8m",

            'black'      :    "\033[30m", 
            'red'        :    "\033[31m",
            'green'      :    "\033[32m",
            'yellow'     :    "\033[33m",
            'blue'       :    "\033[34m",
            'magenta'    :    "\033[35m",
            'cyan'       :    "\033[36m",
            'white'      :    "\033[37m",

            'on_black'   :    "\033[40m", 
            'on_red'     :    "\033[41m",
            'on_green'   :    "\033[42m",
            'on_yellow'  :    "\033[43m",
            'on_blue'    :    "\033[44m",
            'on_magenta' :    "\033[45m",
            'on_cyan'    :    "\033[46m",
            'on_white'   :    "\033[47m",

            'beep'       :    "\007"
            }


normal_colour = 'default'
header_colour = 'yellow'
local_fs_colour = 'default'
remote_fs_colour = 'green'
special_fs_colour = 'blue'

filled_fs_colour = 'red'
full_fs_colour = 'magenta'

sizeformat = "-h"

FILL_THRESH = 95.0
FULL_THRESH = 99.0


format = [    
            ('fs', 0.2, "l"), ('size', 0.1, "r"), 
            ('used', 0.1, "r"), ('avail', 0.1, "r"), ('perc', 0.1, "r"),
            ('bar', 0.1, "l"), ('on', 0.3, "l")
         ]


barchar = '#'

mountfile = ['/etc/mtab', '/etc/mnttab', '/proc/mounts']


#end of default definitions

# read configuration file
for i in ["/etc/pydfrc", os.environ['HOME']+"/.pydfrc"]:
    if os.path.isfile(i):
        execfile(i)


header =    {        
            'fs'     :    "Filesystem",
            'size'   :    "Size",
            'used'   :    "Used",
            'avail'  :    "Avail",
            'on'     :    "Mounted on",
            'fstype' :    "Type",
            'perc'   :    "Use%",
            'bar'    :    ""
            }

def out(s):
    sys.stdout.write(s)

def hfnum(size, base):
    "human readable number"
    if size == 0:
        return "0"
    units = ["B", "k", "M", "G", "T", "P", "Z", "Y"]
    power = int(log(size)/log(base))
    if power<0:
        power = 0
    if power>=len(units):
        power = len(units)-1
    nsize = int(round(1.*size/(base**power)))
    if nsize<10 and power>=1:
        power -=1
        nsize = int(round(1.*size/(base**power)))
    r = str(nsize)+units[power]
    return r


def myformat(number, sizeformat, fs_blocksize):
    "format number as file size. fs_blocksize here is a filesysem blocksize"
    size = long(number)*fs_blocksize
    if blocksize: # that is, blocksize was explicitly set up
       sn = round(1.*size/blocksize)
       sn = int(sn)
       return str(sn)

    if sizeformat == "-k":
        sn = round(size/1024.)
        sn = int(sn)
        return str(sn)
    elif sizeformat == "-m":
        sn = round(size/(1024.*1024))
        sn = int(sn)
        return str(sn)
    elif sizeformat == "-g":
        sn = round(size/(1024.*1024*1024))
        sn = int(sn)
        return str(sn)
    elif sizeformat == "-h":
        return hfnum(size, 1024)
    elif sizeformat == "-H":
        return hfnum(size, 1000)
    elif sizeformat == "--blocks":
        return str(number)
    else: # this should not happen
        raise "Impossible error, contact the author, sizeformat="+`sizeformat`

def myatof(s):
    "like float(), but be friendly to non-numerical values"
    try:
        return float(s)
    except ValueError:
        return 0


def manglestring(s, l, pos):
    "cut string to fit exactly into l chars"
    if pos == "r":
        ns = string.rjust(s, l)
    elif pos == "l":
        ns = string.ljust(s, l)
    elif pos == "c":
        ns = string.center(s, l)
    else:
        raise ValueError, 'Error in manglestring'
    if len(ns) > l:
        ns = ns[:l/2] + "~" + ns[-l/2+1:]
    return ns

def makecolour(clist):
    "take list (or tuple or just one name) of colour names and return string of ANSI definitions"
    s = ""
    if type(clist) == types.StringType:
        lclist = [clist]
    else:
        lclist = clist
    for i in lclist:
        s = s + colours[i]
    return s

def version():
    print "pydf 0.9.8.4"
    sys.exit()

def get_all_mountpoints():
    "return all mountpoints in fs"

    # fallabck when nothing else works
    dummy_result = {'/': ('/', '')}
    
    if isinstance(mountfile, str):
        f = open(mountfile,"r")
    else:
        for i in mountfile:
            if os.path.exists(i):
                f = open(i,"r")
                break
        else:
            # fallback, first try to parse mount output
            status, mout = commands.getstatusoutput('mount')
            if status !=0:
                return dummy_result
            mlines = mout.split('\n')
            r = {}
            for line in mlines:
                if not ' on ' in line:
                    continue
                device, on = line.split(' on ', 1)
                device = device.split()[0]
                on = on.split()[0]
                r[on] = (device, '')
            if r:
                return r
            else:
                return dummy_result
                    
            
    mountlines = f.readlines()
    r = {}
    for l in mountlines:
        spl = l.split()
        if len(spl)<3:
            print "Error in", mountfile
            print `l`
            continue
        device, mp, typ = spl[0:3]
        r[mp] = (device, typ)
    return r
    
def display_mp(mp):
    "display line for mountpoint mp, if mp==None display just the header"
    if mp:
        device, fstype = mountpoints[mp]

        try:
            status = os.statvfs(mp)
        except:
            status = 10*[0]

        fs_blocksize = status[F_BSIZE]
        if fs_blocksize == 0:
            fs_blocksize = status[F_FRSIZE]
        free = status[F_BFREE]
        used = long(status[F_BLOCKS]-free)
        size = status[F_BLOCKS]
        if size==0 and not allfss:
            return
        avail = status[F_BAVAIL]
        size_f = myformat(size, sizeformat, fs_blocksize)
        used_f = myformat(used, sizeformat, fs_blocksize)
        avail_f = myformat(avail, sizeformat, fs_blocksize)
    
        try:
            perc = round(100.*used/(used + status[F_BAVAIL]))
            perc = round(100.*used/size, 1)
        except ZeroDivisionError:
            perc = 0
        perc_f = str(perc)


        info = {
            'fs'     :    device,
            'size'   :    size_f,
            'used'   :    used_f,
            'avail'  :    avail_f,
            'on'     :    mp,
            'fstype' :    fstype,
            'perc'   :    perc_f,
                }
            
        current_colour = local_fs_colour
        if is_remote_fs(fstype):
            current_colour = remote_fs_colour
        elif size == 0:
            current_colour = special_fs_colour
    else: # header
        current_colour = header_colour


    for j in format:
            width = j[1]
            if 0<width<1: # i.e. percentage
                width = int(width*terminal_width)-1
            out( makecolour(current_colour) )
            if mp:
                if j[0] in ['perc', 'avail', 'bar'] and fstype not in  ("iso9660", "udf"):
                    if perc > FILL_THRESH:
                        out( makecolour(filled_fs_colour) )
                    if perc> FULL_THRESH:
                        out( makecolour(full_fs_colour) )
                if j[0] == 'bar':
                    out('[')
                    barsize = int(round((width-2)*perc/100.))
                    out( manglestring(barsize*barchar, width-2, j[2]) )
                    out(']')
                else:
                    out( manglestring(info[j[0]], width, j[2]) )
            else:
                out(manglestring(header[j[0]], width, j[2]))
            out(' ')
    out( colours['default'] )
    out('\n')

def is_remote_fs(fs):
    "test if fs (as type) is a remote one"
    return fs.lower() in [ "nfs", "smbfs", "cifs", "ncpfs", "afs", "coda", "ftpfs", "mfs", "sshfs" ]


# the fun begins here
    


parser = OptionParser(usage="usage: %prog [options] arg", add_help_option=False)
parser.add_option("", "--help", action="help", help="show this help message")
  
parser.add_option("-a", "--all",
      action="store_true", dest="show_all", default=False,
      help="include filesystems having 0 blocks")
parser.add_option("-h", "--human-readable",
      action="store_const", const='-h', dest="sizeformat",
      help="print sizes in human readable format (e.g., 1K 234M 2G)")
parser.add_option("-H", "--si",
      action="store_const", const='-H', dest="sizeformat",
      help="likewise, but use powers of 1000 not 1024")
parser.add_option("-b", "--blocks-size",
      action="store", dest="blocksize", default=0, type="int",
      help="use SIZE-byte blocks")
parser.add_option("-l", "--local",
      action="store_true", dest="local_only", default=False,
      help="limit listing to local filesystems")
parser.add_option("-k", "--kilobytes",
      action="store_const", const='-k', dest="sizeformat",
      help="like --block-size=1024")
parser.add_option("-m", "--megabytes",
      action="store_const", const='-m', dest="sizeformat",
      help="like --block-size=1048576")
parser.add_option("-g", "--gigabytes",
      action="store_const", const='-g', dest="sizeformat",
      help="like --block-size=1073741824")
parser.add_option("", "--blocks",
      action="store_const", const='--blocks', dest="sizeformat",
      help="use filesystem native block size")
parser.add_option("", "--bw",
      action="store_true", dest="b_w", default=False,
      help="do not use colours")
parser.add_option("", "--mounts",
      action="store", dest="mounts_file", type="string",
      help="""File to get mount information from.
On normal linux system, only /etc/mtab or proc/mounts make sense.
Some other unices use /etc/mnttab.
Use /proc/mounts when /etc/mtab is corrupted or inaccesable 
(the output looks a bit weird in this case).

""")

(options, args) = parser.parse_args()

blocksize = options.blocksize
allfss = options.show_all
localonly = options.local_only

if options.sizeformat:
    sizeformat = options.sizeformat

if options.b_w:
    normal_colour = header_colour = local_fs_colour = remote_fs_colour = special_fs_colour = filled_fs_colour = full_fs_colour = 'none'
if options.mounts_file:
    mountfile = options.mounts_file

terminal_width = get_terminal_width()

mountpoints = get_all_mountpoints()

if args:
    mp_to_display = args
else:
    mp_to_display = mountpoints.keys()
    if localonly:
        mp_to_display = [x for x in mp_to_display if not is_remote_fs(mountpoints[x][1])]
    mp_to_display.sort()

display_mp(None) #header
for mp in mp_to_display:
    display_mp(mp)
    
out( colours['default'] )

