#! /usr/bin/python -tt

import os
import sys

import yum
from urlgrabber.progress import format_number

# Decent (UK/US English only) number formatting.
import locale
locale.setlocale(locale.LC_ALL, '') 

def loc_num(x):
    """ Return a string of a number in the readable "locale" format. """
    return locale.format("%d", int(x), True)


verbose = False
if len(sys.argv) >= 2 and sys.argv[1] == '-v':
    verbose = True
    sys.argv.pop(1)
if len(sys.argv) < 2 or sys.argv[1][0] != '/':
    print >>sys.stderr, " Usage: yum-du <absolute-path>"

prefix = sys.argv[1]

yb = yum.YumBase()


pkg_files = {}
pkg_ghost = {}
for pkg in yb.rpmdb:
    for fname in pkg.filelist + pkg.dirlist:
        if not fname.startswith(prefix):
            continue
        pkg_files[fname] = pkg
    for fname in pkg.ghostlist:
        if not fname.startswith(prefix):
            continue
        pkg_ghost[fname] = pkg

def _dir_prefixes(path):
    while path != '/':
        path = os.path.dirname(path)
        yield path

pkgs_size = {}
gkgs_size = {}
pres_size = {}
data_size = {}
pkgs_count = 0
gkgs_count = 0
data_count = 0
num        = 0
for root, dirs, files in os.walk(prefix):
    for fname in files:
        num += 1
        fpath = root + '/' + fname
        if os.path.islink(fpath):
            continue
        size = os.path.getsize(fpath)
        if fpath in pkg_files:
            pkgs_count += size
            if pkg_files[fpath] not in pkgs_size:
                pkgs_size[pkg_files[fpath]] = 0
            pkgs_size[pkg_files[fpath]] += size
        elif fpath in pkg_ghost:
            gkgs_count += size
            if pkg_ghost[fpath] not in gkgs_size:
                gkgs_size[pkg_ghost[fpath]] = 0
            gkgs_size[pkg_ghost[fpath]] += size
        else:
            fpre_path = fpath
            for fpre_path in _dir_prefixes(fpath):
                if fpre_path in pkg_files:
                    if pkg_files[fpre_path] not in pres_size:
                        pres_size[pkg_files[fpre_path]] = 0
                    pres_size[pkg_files[fpre_path]] += size
                    break

            data_count += size
            data_size[fpath] = size

print "Files            :", loc_num(num)
tot = pkgs_count + gkgs_count + data_count
print "Total size       :", format_number(tot)
num = pkgs_count
if not verbose:
    num += gkgs_count
print "       Pkgs size :", "%-5s" % format_number(num),
print "(%3.0f%%)" % ((num * 100.0) / tot)
if verbose:
    print " Ghost pkgs size :", "%-5s" % format_number(gkgs_count),
    print "(%3.0f%%)" % ((gkgs_count * 100.0) / tot)
print "       Data size :", "%-5s" % format_number(data_count),
print "(%3.0f%%)" % ((data_count * 100.0) / tot)
if verbose:
    print ''
    print "Pkgs       :", loc_num(len(pkgs_size))
    print "Ghost Pkgs :", loc_num(len(gkgs_size))

def _pkgs(p_size, msg):
    tot = min(4, len(p_size))
    if tot:
        print ''
        print msg % tot
    num = 0
    for pkg in sorted(p_size, key=lambda x: p_size[x], reverse=True):
        num += 1
        print "%2d. %60s %-5s" % (num, pkg, format_number(p_size[pkg]))
        if num >= tot:
            break
_pkgs(pkgs_size, 'Top %d packages:')
if verbose:
    _pkgs(gkgs_size, 'Top %d ghost packages:')
_pkgs(pres_size, 'Top %d prefix packages:')

if verbose:
    print ''
    tot = min(4, len(data_size))
    if tot:
        print 'Top %d files:' % tot
    num = 0
    for fname in sorted(data_size, key=lambda x: data_size[x], reverse=True):
        num += 1
        print "%2d. %60s %-5s" % (num, fname, format_number(data_size[fname]))
        if num >= tot:
            break