#!/usr/bin/perl
use charnames ':full'; # only necessary before v5.16
use strict;
use v5.24;
use warnings;
use autodie;

use DBIx::SimpleConnect;
use Encode qw(:fallbacks encode);
use Getopt::Long;
use Pod::Usage;
use POSIX qw(strftime);

my $db = "default";
my $dir = ".";

GetOptions(
    'dbname=s' => \$db,
    'directory=s' => \$dir,
) or pod2usage(1);

binmode STDOUT, ":encoding(UTF-8)";

{
    no autodie;
    unless (chdir($dir)) {
        # if at first we don't succeed ...
        mkdir $dir or die "cannot create $dir: $!";
        chdir $dir or die "cannot chdir to $dir: $!";
    }
}
my $errorfile = strftime("%Y-%m-%d %H:%M:%S%z.err", localtime);
open STDERR, ">:encoding(UTF-8)", $errorfile;

my $dbh = DBIx::SimpleConnect->connect($db, {RaiseError => 1, PrintError => 0});
# this is intended to run from cron, so limit the time it can run.
$dbh->do("set statement_timeout to 30000");

my $overview
    = $dbh->selectall_hashref("
                select schemaname, tablename,
                    sum(pg_total_relation_size(quote_ident(schemaname) || '.' || quote_ident(tablename))) as total_size,
                    sum(pg_table_size(quote_ident(schemaname) || '.' || quote_ident(tablename))) as table_size,
                    sum(pg_indexes_size(quote_ident(schemaname) || '.' || quote_ident(tablename))) as indexes_size
                from pg_tables
                group by rollup (schemaname, tablename)
                ",
                ['schemaname', 'tablename']
            );

open(my $fh, ">:encoding(UTF-8)", "index.html.$$");
print_header($fh);
say $fh "<h1>Disk usage for database $db</h1>";
say $fh "<div class='timestamp'>", strftime("%Y-%m-%d %H:%M:%S%z", localtime), "</div>";
say $fh "<table>";
say $fh "<tr class='header'>";
say $fh "<th class='name'>", "Schema", "</th>";
say $fh "<th class='name'>", "Table", "</th>";
say $fh "<th class='num'>", "Total size (bytes)", "</th>";
say $fh "<th class='num'>", "Table size (bytes)", "</th>";
say $fh "<th class='num'>", "Indexes size (bytes)", "</th>";
say $fh "</tr>";

for my $s (sort { $overview->{$b}{''}{total_size} <=> $overview->{$a}{''}{total_size} } keys $overview->%*) {
    
    say $fh "<tr class='schema'>";
    say $fh "<td class='name'>", escape($s), "</td>";
    say $fh "<td class='name'>", '', "</td>";
    say $fh "<td class='num'>", pretty($overview->{$s}{''}{total_size}), "</td>";
    say $fh "<td class='num'>", pretty($overview->{$s}{''}{table_size}), "</td>";
    say $fh "<td class='num'>", pretty($overview->{$s}{''}{indexes_size}), "</td>";
    say $fh "</tr>";

    for my $t (sort { $overview->{$s}{$b}{total_size} <=> $overview->{$s}{$a}{total_size} } keys $overview->{$s}->%*) {
        next unless $t; # schema - we already have that

        say $fh "<tr class='table'>";
        say $fh "<td class='name'>", escape($s), "</td>";
        say $fh "<td class='name'>", escape($t), "</td>";
        say $fh "<td class='num'>", pretty($overview->{$s}{$t}{total_size}), "</td>";
        say $fh "<td class='num'>", pretty($overview->{$s}{$t}{table_size}), "</td>";
        say $fh "<td class='num'>", pretty($overview->{$s}{$t}{indexes_size}), "</td>";
        say $fh "</tr>";
    }
}
say $fh "</table>";
close($fh);
rename "index.html.$$", "index.html";
unlink($errorfile); # if we get here, we are fine
exit(0);


sub pretty {
    my ($n) = @_;

    while ($n =~ /[0-9]{4}/) {
        $n =~ s/([0-9]+)([0-9]{3})/$1\N{NARROW NO-BREAK SPACE}$2/;
    }
    return $n;
}

sub escape {
    my ($s) = @_;

    $s =~ s/[<>&'"]/sprintf("&#%d;", ord($1))/eg;
    return $s;
}

sub print_header {
    my ($fh) = @_;

    say $fh "<meta charset='utf-8'/>";
    say $fh "<style>";
    say $fh "    .num {";
    say $fh "        text-align: right;";
    say $fh "    }";
    say $fh "    .schema {";
    say $fh "        background-color: #CCCCCC;";
    say $fh "    }";
    say $fh "    .table {";
    say $fh "        background-color: #E5E5E5;";
    say $fh "    }";
    say $fh "</style>";
}

# vim: sw=4 tw=132 expandtab