#!/usr/bin/perl

# Copyright (c) 2008-2012, Kurt D. Starsinic <kstarsinic@gmail.com>
# Patches welcome.
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the Artistic License 2.0.
#
# This program is distributed in the hope that it will be useful, but it is
# provided "as is" and without any express or implied warranties.

# For details, see the full text of the license at
# <http://www.perlfoundation.org/artistic_license_2_0>.

# This program is short on documentation.  This shall be remedied in the
# upcoming release.

# Basic usage:
# dus.pl [any arguments you would pass to du(1)]
# The "Page Up" and "Page Down" keys will do what
# you expect.  At any time you can press "q" to quit.

# TODO:
#   Allow an intelligent reload
#   Reformat for large disks
#   PageUp/PageDown should work after we're done loading.

use strict;
use warnings; no warnings 'numeric', 'once';

use Switch;


require Curses;
my @Colors;
my $W = begin_curses();


### %Du and @Du are the master data structures.
###
### There is one entry in %Du for every directory we've seen so far, and all
### of those directories' parent directories, all the way up to the top of
### the hierarchy we're doing a "du" on.  The keys of %Du are directory names,
### and the values are array references.
###
### The first element of the array is the directory's size -- either estimated
### or actual.
###
### If there are *any* following elements, they are the names of the
### directory's subdirectories that we've seen so far (and the first entry is
### an *estimated* size).
###
### If there are *no* following elements, then the first element is the
### *actual* size of the directory.
###
### There is one line in @Du for every line displayed on the screen.  Each
### element of @Du is a directory name.  @Du is always sorted in size order.
my (%Du, @Du);


@ARGV = '.' if not @ARGV;

# Remove trailing '/' from non-option args (excepting '/' itself):
s{/$}{} for grep { $_ ne '/' and not '^-' } @ARGV;

### Figure out which of our command line arguments are directory names.
### Create an entry in %Du for each of them, and tag them as "ROOT" nodes
### (by blessing them into the ROOT package).  These aren't objects;
### blessing is just a convenient way to hang some metadata on those entries.
foreach my $arg (@ARGV) {
    if ($arg !~ /^-/) {
        my ($norm) = ($arg =~ m:(.*?)/*$:);
        $Du{$norm} = [ 0 ];
        bless $Du{$norm}, 'ROOT';
    }
}


### The main loop:
### Call "du", passing along the arguments we were called with.  For each line
### of output from "du", update the data structures and re-paint the screen.
my $Page = 1;
open DU, "du @ARGV 2>/dev/null |";
    while (<DU>) {
        chomp;

        my ($n, $s) = (/^(\d+)\s+(.+)/);
        $s  = '' if $s eq '/';
        @{ $Du{$s} }    = ( $n );

        remove($s); # Take the entries for $s and its parents out of @Du
        bubble($s); # Put them back, in their new correct locations

        my $files_per_page = $Curses::LINES - 1;
        for my $y (0 .. $files_per_page - 1) {
            my $i = $y + (($Page - 1) * $files_per_page);

            if ($i <= $#Du) {
                my $name    = $Du[$i]; $name ='/' if $name eq '';
                my $d       = $Du[$i];

                if    (ref $Du{$d} eq 'ROOT')   { $W->attrset($Colors[1]) }
                elsif ($#{ $Du{$d} })           { $W->attrset($Colors[0]) }
                else                            { $W->attrset($Colors[2]) }

                $W->addstr($y, 0, sprintf("%8d %-200s", $Du{$d}[0], $name));
            }
        }
        $W->clrtobot;

        my $hdr = "Page $Page";
        $W->attrset($Colors[1]);
        $W->addstr(0, $Curses::COLS-length($hdr)-5, $hdr);

        process_input();
    }
close DU;
$W->attrset(Curses::COLOR_PAIR(2) | Curses::A_BOLD());
$W->addstr($Curses::LINES - 1, 0, "--- done ---");

$W->nodelay(0);
process_input() while 1;;
exit;


### Given the newly-arrived entry for $here, insert it into the proper slot in
### @Du.  Then re-calculate $here's parent directory's size, and recurse up
### the directory tree until we hit a ROOT node.
sub bubble {
    my ($here)  = @_;
    my $dotdot  = $here;    $dotdot =~ s:/[^/]*$::;

    # Update the parent directory's entry, unless we're at the root of a tree:
    if (ref $Du{$here} ne 'ROOT' and $here ne $dotdot) {
        # Set $dotdot's size to zero:
        $Du{$dotdot}[0] = 0;

        # Record $here as a subdirectory of $dotdot, . . .
        push(@{ $Du{$dotdot} }, $here)
            unless grep         # . . . being careful not to RE-record it.
                { $_ eq $here }
                @{ $Du{$dotdot} }[1 .. $#{ $Du{$dotdot} }];

        # Now calculate the estimated size of $dotdot:
        foreach my $subdir (@{ $Du{$dotdot} }[1 .. $#{ $Du{$dotdot} }]) {
            $Du{$dotdot}[0] += $Du{$subdir}[0];
        }
    }

    # Find the entry to insert $here after:
    my $size    = $Du{$here}[0];
    my $i;
    for ($i = $#Du; $i >= 0; $i--) {
        my $dui     = $Du[$i];
        my $test    = $Du{$dui}[0];

        last if $test > $size;
        last if ($test == $size) && (not index $here, $dui);
    }

    # Now insert $here.  If we made it all the way through the above loop
    # without finding a place for it, then $i is -1, and we'll insert $here
    # in the 0th position (i.e., at the beginning of the list):
    splice @Du, $i+1, 0, $here;

    # Keep @Du from growing without bound, because splice() doesn't scale.
    pop @Du if (@Du > 500 && @Du > ($Curses::LINES * ($Page + 1)));

    # Lather, rinse, repeat:
    bubble($dotdot) if (ref $Du{$here} ne 'ROOT') && $here ne $dotdot;
}


sub remove {
    my ($path)  = @_;

    # Remove $here from @Du, if present:
    @Du = grep { $_ ne $path } @Du;

    # Now remove its parent directory, unless we're at the top of our tree:
    if (ref $Du{$path} ne 'ROOT') {
        $path =~ s:/[^/]*$::;
        remove($path);
    }
}


sub process_input
{
    $W->refresh;

    switch ($W->getch) {
        case Curses::ERR()          { return }
        case Curses::KEY_NPAGE()    { $Page += 1 }
        case Curses::KEY_PPAGE()    { $Page -= 1 unless $Page == 1 }
        case 'q'                    { end_curses(); exit }
        else                        { Curses::beep() }
    }
}


sub begin_curses
{
    my $w = Curses->new;

    Curses::start_color();

    my $bg = Curses::COLOR_BLUE();
    Curses::init_pair(1, Curses::COLOR_CYAN(),   $bg);
    Curses::init_pair(2, Curses::COLOR_YELLOW(), $bg);
    Curses::init_pair(3, Curses::COLOR_WHITE(),  $bg);
    Curses::assume_default_colors(Curses::COLOR_WHITE(), $bg);

    @Colors = (
        Curses::COLOR_PAIR(1) | Curses::A_BOLD(),
        Curses::COLOR_PAIR(2) | Curses::A_BOLD(),
        Curses::COLOR_PAIR(3) | Curses::A_BOLD(),
    );

    $w->nodelay(1);
    $w->keypad(1);
    Curses::noecho();

    return $w;
}


sub end_curses { Curses::endwin() }