#!/usr/bin/env perl
# vim: ts=4 sw=4 noet
use List::Util qw(min);

sub csi {
	my ($format, $str) = @_;
	return (-t 1) ? "\e[${format}m${str}\e[m" : "$str";
}

sub fmt_size {
	my ($bytes, $si) = @_;
	return "0" if !$bytes;
	my @prefixes = qw(k M G T P E Z Y H);
	my $div = $si ? 1000 : 1024;
	my $exp = min(int(log($bytes) / log($div)), scalar @prefixes);
	sprintf "%.1f%s", ($bytes / $div ** $exp), ($exp ? $prefixes[$exp-1] : "");
}

# Get indent prefix for depth(path)
sub indent_for {
	my ($dir) = @_;
	my $depth = split(m!/!, $dir);
	my $indent = " " x 4;
	return $indent x $depth;
}

# Get depth-specific color for depth(path)
sub color_for {
	my ($dir) = @_;
	my $depth = split(m!/!, $dir);
	return 2 + $depth;
	#my @colors = (51, 39, 26, 21, 93, 165, 198);
	#return $colors[$idx] // $colors[$#colors];
}

sub fmt_dir {
	my ($dir, $darken) = @_;
	$dir =~ s{.*/}{};
	$dir = ($dir eq "" ? "." : $dir);
	return ($darken) ? csi("2", $dir) : $dir;
}

my $max_depth = undef;
my $show_files = 0;
my $sort_by_size = 0;
my $sort_reverse = 0;
my $decimal_size = 1;

my @cmd = ("annex", "find", "--format=\${key} \${file}\\n", "--include=*");

for (@ARGV) {
	if (/^-[av]$/) {
		$show_files = 1;
	} elsif (/^-d?([0-9]+)$/) {
		$max_depth = int $1;
	} elsif (/^-r$/) {
		$sort_reverse = 1;
	} elsif (/^-s$/) {
		$sort_by_size = 1;
	} elsif (/^-h/ || /^--iec$/) {
		$decimal_size = 0;
	} elsif (/^-H/ || /^--si$/) {
		$decimal_size = 1;
	} elsif (/^-[^-]/) {
		die "annex du: unknown option $_\n";
	} else {
		push @cmd, $_;
	}
}

# Default to infinite for -a, 2 levels otherwise
if (!defined $max_depth) {
	$max_depth = ($show_files) ? -1 : 2;
}

# Sizes per path; per directory; total
my %files;
my %dirs;
my $total = 0;

open(my $fh, "-|", @cmd);
while (<$fh>) {
	my ($key, $file, $size);

	chomp;
	($key, $file) = split(/ /, $_, 2);
	if ($key =~ /-s(\d+)--/) {
		$size = int $1;
	}

	if ($show_files) {
		$files{$file} += $size;
	}
	while ($file =~ s!/[^/]+$!!) {
		$dirs{$file} += $size;
	}
	$total += $size;
}
close($fh);

%dirs = (%dirs, %files);

my @dirs = sort keys %dirs;
if ($max_depth >= 0) {
	@dirs = grep {split(m!/!, $_) <= $max_depth} @dirs;
}
if ($sort_by_size) {
	@dirs = sort {$dirs{$b} <=> $dirs{$a}} @dirs;
}
if ($sort_reverse) {
	@dirs = reverse @dirs;
}

print
	indent_for(""),
	csi("38;5;".color_for(""),
	    sprintf("%8s",
	            fmt_size($total, $decimal_size) . ($decimal_size ? "" : "i"))),
	"  ",
	"TOTAL",
	"\n";

print
	indent_for($_),
	csi("38;5;".color_for($_),
	    sprintf("%8s",
	            fmt_size($dirs{$_}, $decimal_size) . ($decimal_size ? "" : "i"))),
	"  ",
	fmt_dir($_, exists($files{$_})),
	"\n"
	for @dirs;
