#!/usr/local/bin/perl
#
#Script to list users 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] [ -account ACCOUNT ]

Displays the list of users that Slurm sacctmgr knows about.
This _should_ correlate to what is in the project DB, but never know.
Actually, users not always created for account on both clusters, so
can differ significantly.

If CLUSTER or ACCOUNT given, only list users in given CLUSTER and/or
belonging to given ACCOUNT.

In verbose mode, print out extra info.

EOF
}


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

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

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

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 @users = ();
my @tmp;

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

my %accounts = map { $_ => undef } @$accounts;

if ( $verbose )
{	#Do verbose output
	unless ( @users )
	{	print "No matching users found.\n";
		exit 0;
	}
	my $fmt = "%7s %10s %20s\n";
	printf $fmt, "Cluster",	"User", "Account";
	USER: foreach my $user (@users)
	{	#Get clusters user is in (filtered on $cluster)
		@tmp = grep { $_->user && $_->user eq $user } @$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 accounts user belongs to in cluster
			@tmp = grep { $_->user && $_->user eq $user &&
				$_->cluster eq $clus }
				@$all_assocs;
			my @accounts = map { $_->account } @tmp;
			#Restrict to specified accounts (if specified!)
			@accounts = grep { exists $accounts{$_} } @accounts
				if scalar(%accounts);
			%temp = map { $_ => $_ } @accounts;
			@accounts = sort ( values %temp );

			ACCOUNT: foreach my $acc (@accounts)
			{	printf $fmt, $clus, $user, $acc;
			}	
		}
	}
} else
{	my $tmp = join "\n", @users;
	print "$tmp\n";
}


