#!/usr/local/bin/perl
#
#Script to list accounts known to sacctmgr for a given cluster
use strict;
use warnings;

use lib "/cell_root/software/HPCC/lib/site_perl";
#use lib "/afs/.glue.umd.edu/software/HPCC/lib/site_perl";

use Slurm::Sacctmgr;
use Slurm::Sacctmgr::Association;

my $SAAssoc = 'Slurm::Sacctmgr::Association';

use Getopt::Long;
my $sacctmgr;

sub usage()
{	print STDERR <<EOF;
$0: List accounts known to Slurm/sacctmgr

Usage:
$0 [-verbose] [-cluster CLUSTER] [ -user USER ] 

Displays the list of accounts that Slurm sacctmgr knows about.
This _should_ correlate to what is in the project DB, but never know.

If CLUSTER or USER given, only list accounts in given CLUSTER and/or
with given USER as a member.  

In verbose mode, print out extra info about accounts

EOF
}


#-------------------------------------------------------------------
#		Main
#-------------------------------------------------------------------

my $help;
my $cluster;
my $users = [];
my $verbose;

my $res = GetOptions(
	'help!' => \$help,
	'verbose!' => \$verbose,
	'cluster=s' => \$cluster,
	'user=s' => $users,
);

unless ( $res )
{	usage;
	die "\nError parsing arguments.\n";
}
if ( $help )
{	usage;
	exit 0;
}

$sacctmgr = Slurm::Sacctmgr->new;

my $all_assocs = $SAAssoc->sacctmgr_list($sacctmgr);
die "Error listing associations: $all_assocs\n" 
	unless $all_assocs && ref($all_assocs) eq 'ARRAY';

#If we specified a cluster, filter on it
if ( $cluster )
{	$all_assocs = [ grep { $_->cluster eq $cluster } @$all_assocs ];
}

my @accounts = ();
my @tmp;

#If one or more user specified, only get accounts with those users
if ( scalar(@$users) )
{	USER: foreach my $user (@$users)
	{	@tmp = grep { $_->user && $_->user eq $user } @$all_assocs;
		@tmp = map { $_->account } @tmp;
		push @accounts, @tmp;
	}
} else
{	@accounts = map { $_->account } @$all_assocs;
}
#Remove duplicates
my %temp = map { $_ => $_ } @accounts;
@accounts = sort (values %temp);

if ( $verbose )
{	#Do verbose output
	unless ( @accounts )
	{	print "No matching accounts found.\n";
		exit 0;
	}
	my $fmt = "%7s %20s %10s %10s\n";
	printf $fmt, "Cluster",	"Account","GrpCPUMins", "kSU";
	ACCOUNT: foreach my $account (@accounts)
	{	#Get clusters account is in (filtered on $cluster)
		@tmp = grep { $_->account eq $account } @$all_assocs;
		@tmp = map { $_->cluster } @tmp;
		%temp = map { $_ => $_ } @tmp;
		my @clusters = sort ( values %temp );
		die "No clusters found " unless @clusters;

		CLUSTER: foreach my $clus (@clusters)
		{	#Get CPUMins
			@tmp = grep { $_->account eq $account &&
				$_->cluster eq $clus && ! $_->user }
				@$all_assocs;
			die "Multiple records for $account, $clus user=null"
				if @tmp > 1;

			my $rec = $tmp[0];
			my ($cpumin, $ksu);
			$cpumin = $rec->grpcpumins if $rec && ref($rec);
			$ksu = sprintf("%.3f", $cpumin/60/1000) if $cpumin;
			$cpumin = 'unlimited' unless $cpumin;
			$ksu = 'unlimited' unless $ksu;
			printf $fmt, $clus,	$account, $cpumin, $ksu;
		}
	}
} else
{	my $tmp = join "\n", @accounts;
	print "$tmp\n";
}


