#!/usr/local/bin/perl

use strict;
use warnings;

use lib '/cell_root/software/HPCC/lib/site_perl';
use Slurm::Scontrol;
use Slurm::Squeue;
use Slurm::Sinfo;
use Slurm::Sinfo::Node;
use Slurm::Sinfo::CombinedNode;
use Data::Dumper;
use Getopt::Long;
use POSIX qw(strftime);

#Globals
my $sinfo = new Slurm::Sinfo;
my $squeue = new Slurm::Squeue;
my $scontrol = new Slurm::Scontrol;

sub usage()
{	print STDERR <<EOF;
$0: Displays information about node usage/availability.

Usage

$0 [ options ]

Without any options, displays all nodes on system along with their
status (allocated, idle, down, etc).  

For allocated nodes, will print job numbers on node as well as owner of job.  
For idle nodes, will print the list of partitions node is part of.
For other states, will print the reason flag.

Options can be used to control output:

--mode MODE: MODE can be one of
	free (or idle): only display free/idle nodes
	down: only display down or problematic nodes
	used: only display nodes in use.  Possibly further restricted
		by other options.
	node: only display specified nodes
	all: display infomation for all nodes.
	Default is 'all'
--free/--idle: Same as --mode free
--used: Same as --mode used
--down: Same as --mode down
--all: Same as --mode all
--stats: Prints a summary of usage stats, only valid with mode=all

In 'node' mode, the nodes are specified with the
--node NODENAME option.  NODENAME can be in standard slurm abbreviated
form (e.g. compue-f10-[1-5]), and this option can also be repeated for
multiple NODENAMEs.  Use of the --node option forces mode to 'node'.

In 'used' mode, one can further restrict output to nodes associated
with specific jobs or jobs belonging to specific users.
--user USERNAME: only list nodes where user USERNAME has a job running
--job JOBID: only list nodes running job JOBID
--user, and --job, may be combined and repeated; the result is an ORing
of criteria.  Use of --user or --job options forces mode to 'used'.

EOF
}

sub nodename_sort_function
{	my ($computea, $racka, $numa);
	my ($computeb, $rackb, $numb);

	if ( $a =~ /^compute-/ )
	{	($computea, $racka, $numa) = split /-/, $a,3;
	} else
	{	($racka, $numa) = split /-/, $a, 2;
	}

	if ( $b =~ /^compute-/ )
	{	($computeb, $rackb, $numb) = split /-/, $b,3;
	} else
	{	($rackb, $numb) = split /-/, $b, 2;
	}

	return ( $racka cmp $rackb ) || ( $numa <=> $numb);
}

sub get_free_nodes($)
#Takes ref to $nodehash and returns nodenames of nodes in idle state
{	my $nodehash = shift || {};

	my @list = ();
	LOOP: foreach my $node (values %$nodehash)
	{	next unless $node->node_idle;
		next if $node->node_mixed;
		push @list, $node->nodename;
	}
	return [ sort nodename_sort_function @list ];
}

sub get_down_nodes($)
#Takes ref to $nodehash and returns nodenames of nodes in 'other' state
{	my $nodehash = shift || {};

	my @list = ();
	LOOP: foreach my $node (values %$nodehash)
	{	next unless $node->node_other;
		push @list, $node->nodename;
	}
	return [ sort nodename_sort_function @list ];
}

sub shorten_nodename($)
#Abbreviates nodename.  Basically, strip off domain, compute- prefix
{	my $nodename = shift;
	my $short = $nodename;
	$short =~ s/\.*$//;
	$short =~ s/^compute-//;
	return $short;
}

sub dump_idle_node($;$)
#Given nodehash and nodename for an idle node, dump it
{	my $nodehash = shift;
	my $nodename = shift;

	unless ($nodename)
	{ 	#Dump header
		print "Node: State: partitions\n";
		return;
	}

	my $short = shorten_nodename($nodename);
	my $node = $nodehash->{$nodename};
	my $state = $node->state;
	my $partitions = $node->partitions;
	my $partstr = join ", ", @$partitions;

	print "$short: $state: partitions=[$partstr]\n";
}

sub dump_down_node($;$)
#Given nodehash and nodename for an down node, dump it
{	my $nodehash = shift;
	my $nodename = shift;

	unless ($nodename)
	{ 	#Dump header
		print "Node: State: reason [time]\n";
		return;
	}
	my $short = shorten_nodename($nodename);
	my $node = $nodehash->{$nodename};
	my $state = $node->state;
	my $reason = $node->reason;
	my $tstamp = $node->reason_timestamp;
	print "$short: $state: $reason [$tstamp]\n";
}

sub dump_used_node($$;$)
#Given nodehash and jobhash and nodename for used node, dump it
{	my $nodehash = shift;
	my $jobhash = shift;
	my $nodename = shift;

	unless ( $nodename )
	{	#Dump header
		print "Node: State: User[JobNumbers]\n";
		return;
	}

	my $short = shorten_nodename($nodename);
	my $node = $nodehash->{$nodename};
	my $state = $node->state;
	my $nhash = $jobhash->{$nodename} || {};
	my @jobids = sort (keys %$nhash);
	my %jobs_by_user = ();
	my ($user, $tmp);
	foreach my $jobid (@jobids)
	{	my $job = $nhash->{$jobid};
		$user = $job->{user};
		$tmp = $jobs_by_user{$user} || [];
		push @$tmp, $jobid;
		$jobs_by_user{$user} = $tmp;
	}

	my $jobstr = "";
	foreach $user (sort keys %jobs_by_user)
	{	$jobstr .= ", " if $jobstr;
		$tmp = $jobs_by_user{$user} || [];
		$tmp = join ",", @$tmp;
		$jobstr .= "$user\[$tmp\]";
	}

	print  "$short: $state: $jobstr\n";
}

sub dump_node($$;$)
#Given nodehash and jobhash and nodename for arbitrary node, dump it
#Calls one of the specific dump_*_node methods above
{	my $nodehash = shift;
	my $jobhash = shift;
	my $nodename = shift;

	unless ( $nodename )
	{	#Dump header
		print "Node: State: Jobs/Partitions/Reason/etc\n";
		return;
	}

	my $node = $nodehash->{$nodename};
	if ( $node->node_other )
	{	dump_down_node($nodehash, $nodename);
	} elsif ( $node->node_allocated || $node->node_mixed )
	{	dump_used_node($nodehash, $jobhash, $nodename);
	} elsif ( $node->node_idle )
	{	dump_idle_node($nodehash, $nodename);
	} else
	{	die "Should not reach here.\n";
	}
	return;
}

sub get_all_running_jobs()
#Returns  list ref of records for all running jobs
{	my $jobs = $squeue->run_squeue_cmd("-t", "R");
	unless ($jobs && ref($jobs) eq 'ARRAY' )
	{	die "Error looking up running jobs: $jobs at ";
	}
	return $jobs;
}

sub get_jobs_on_node($)
#Given a nodename (of range expansion, etc), returns
#a list ref of all job records for all jobs running on any
#of those nodes.
{	my $nodelist = shift;

	my $nlist = join ",", @$nodelist;
	my $jobs = $squeue->run_squeue_cmd("-t", "R", "-w", $nlist);
	unless ($jobs && ref($jobs) eq 'ARRAY' )
	{	die "Error looking up jobs on node(s) $nodelist: $jobs at ";
	}
	return $jobs;
}

sub get_job($)
#Returns record for specified job id
{	my $jobid = shift;
	my $jobs = $squeue->run_squeue_cmd("-t", "R", "-j", $jobid) || [];
	unless ( $jobs && ref($jobs) eq 'ARRAY' )
	{	die "Error looking up job with id $jobid: $jobs at ";
	}
	return $jobs;
}

sub get_jobs_for_user($)
#Returns records for all running jobs owned by specified user
{	my $user = shift;
	my $jobs = $squeue->run_squeue_cmd("-t", "R", "-u", $user) || [];
	unless ( $jobs && ref($jobs) eq 'ARRAY' )
	{	die "Error looking up jobs for user $user: $jobs at ";
	}
	return $jobs;
}

sub update_jobs_hash($$)
#Given a list ref of job records, update the JobsByNode hash
{	my $jobhash = shift;
	my $jobs = shift || [];
	my $me = 'update_jobs_hash';

	die "$me: Invalid jobhash $jobhash at " unless
		$jobhash && ref($jobhash) eq 'HASH';

	JOB: foreach my $job (@$jobs)
	{	my $raw_nodelist = $job->nodelist;
		my $nodelist = $scontrol->expand_hostnames($raw_nodelist);
		my $jobid = $job->jobid;
		my $uid = $job->user;
		if ( $uid =~ /\D/ )
		{	print STDERR "job=$job\n", Dumper($job), "\n\n";
		}

		my $user = getpwuid($uid);
		NODE: foreach my $node (@$nodelist)
		{	my $nhash = $jobhash->{$node} || {};
			my $jhash = $nhash->{$jobid};
			next NODE if $jhash && ref($jhash) eq 'HASH';
			$jhash = {	jobid => $jobid,
					uid =>  $uid,
					user => $user,
				};
			$nhash->{$jobid} = $jhash;
			$jobhash->{$node} = $nhash;
		}
	}

	return $jobhash;
}
		
	
sub generate_stats($$)
# Generate stats about the cluster's usage
# Returns to a hast ref, with keys
# 	total_nodes: total number of nodes
# 	total_cores: total number of cores
# 	running_jobs: number of running jobs
# 	used_nodes: nodes with one or more jobs on them
# 	used_cores: cores assigned to jobs
# 	down_nodes: nodes in down state
# 	down_cores: cores belonging to nodes in down state
# 	free_nodes: nodes not assigned to any job
# 	free_cores: cores not assigned to any job
# 	pused_nodes: used_nodes/total_nodes*100
# 	pused_cores: used_cores/total_cores*100
{	my $nodehash = shift;
	my $jobs = shift; #Not really used?

	my $running_jobs = scalar(@$jobs);

	my $total_nodes = 0;
	my $total_cores = 0;
	my $used_nodes = 0;
	my $used_cores = 0;
	my $down_nodes = 0;
	my $down_cores = 0;
	my $free_nodes = 0;
	my $free_cores = 0;

	NODE: foreach my $node (values %$nodehash)
	{	#Update totals
		#use Data::Dumper;
		#print "node=$node\n", Dumper($node), "\n\n\n";
		#exit 1;
		$total_nodes += 1;
		$total_cores += $node->cpus;

		#Update categorized totals
		$down_nodes += $node->node_other;
		$down_cores += $node->cpus_other;

		$used_nodes += $node->node_allocated;
		$used_cores += $node->cpus_allocated;

		$free_nodes += $node->node_idle;
		$free_cores += $node->cpus_idle;
	}

	my $pused_nodes = 0;
	$pused_nodes = 100 * $used_nodes/$total_nodes if ( $total_nodes > 0 );
	my $pused_cores = 0;
	$pused_cores = 100 * $used_cores/$total_cores if ( $total_cores > 0 );

	my $stats = 
	{	running_jobs => $running_jobs,

		total_nodes => $total_nodes,
		down_nodes => $down_nodes,
		free_nodes => $free_nodes,
		used_nodes => $used_nodes,

		total_cores => $total_cores,
		down_cores => $down_cores,
		free_cores => $free_cores,
		used_cores => $used_cores,

		pused_nodes => $pused_nodes,
		pused_cores => $pused_cores,
	};

	return $stats;
}

sub dump_stats($)
{	my $stats = shift || {};

	my $running_jobs = $stats->{running_jobs} || 0;
	my $total_nodes = $stats->{total_nodes} || 'unknown';
	my $down_nodes = $stats->{down_nodes} || 0;
	my $free_nodes = $stats->{free_nodes} || 0;
	my $used_nodes = $stats->{used_nodes} || 0;
	my $total_cores = $stats->{total_cores} || 'unknown';
	my $down_cores = $stats->{down_cores} || 0;
	my $free_cores = $stats->{free_cores} || 0;
	my $used_cores = $stats->{used_cores} || 0;
	my $pused_nodes = sprintf("%0.2f%%", ($stats->{pused_nodes} || 0));
	my $pused_cores = sprintf("%0.2f%%", ($stats->{pused_cores} || 0));

	print <<EOF;

---------------------------------------------------------------------
			Statistics
---------------------------------------------------------------------
Running Jobs Count: $running_jobs
Nodes: Total = $total_nodes, Alloc=$used_nodes, Idle=$free_nodes, Other=$down_nodes 
Cores: Total = $total_cores, Alloc=$used_cores, Idle=$free_cores, Other=$down_cores

Percent Used: 
	By Nodes: $pused_nodes
	By Cores: $pused_cores

EOF
}

	
#Main
my $help;
my $mode;
my $jobs = [];
my $users = [];
my $nodes = [];
my $header = 1;
my $dostats = 0;

my $res = GetOptions(
	'help!' => \$help,
	'header!' => \$header,
	'stats|statistics!' => \$dostats,

	'mode=s' => \$mode,
	'free|idle' => sub { $mode = 'free' },
	'down' => sub { $mode = 'down' },
	'all' => sub { $mode = 'all' },
	'used' => sub { $mode = 'used' },

	'jobs=s' => $jobs,
	'users=s' => $users,
	'nodes=s' => $nodes,
);

unless ( $res )
{	usage;
	die "\nError parsing options.\n";
}

if ( $help )
{	usage;
	exit 0;
}

if ( scalar(@$nodes) ) 
{ 	if ( scalar(@$jobs) || scalar(@$users) )
	{	die "--node is incompatible with --user or -job options.\n";
	}
	warn "Use of --node forces mode => 'node' from $mode at "
		if defined $mode;
	$mode = 'node';
}

if ( scalar(@$jobs) || scalar(@$users) )
{	if ( scalar(@$nodes) )
	{	die "--node is incompatible with --user or -job options.\n";
	}
	warn "Use of --node forces mode => 'used' from $mode at "
		if defined $mode;
	$mode = 'used';
}

$mode = 'all' unless $mode;
$mode = lc $mode;
$mode = 'free' if $mode eq 'idle';

unless ( $mode eq 'free' || $mode eq 'all' || $mode eq 'down' ||
	$mode eq 'used' || $mode eq 'node' )
{	usage;
	die "Illegal value $mode for --mode: want one of free/all/down/used/node\n";
}

if ( $dostats )
{	unless ( $mode eq 'all' )
	{	$dostats = 0;
		warn "Ignoring --stats flag since mode != 'all'\n";
	}
}

#Get hash of all nodes
my $nodelist = $sinfo->run_sinfo_cmd();
my $nodehash = Slurm::Sinfo::CombinedNode->make_combined_node_hash($nodelist);

my %JobsByNode = ();
my @NodesToDisplay = ();
my ($tmp, $nodename, $joblist, );

if ( $mode eq 'free' )
{	$tmp = get_free_nodes($nodehash);
	@NodesToDisplay = @$tmp;
	if ( @NodesToDisplay )
	{	dump_idle_node($nodehash) if $header;
		foreach $nodename (@NodesToDisplay)
		{	dump_idle_node($nodehash, $nodename);
		}
	} else
	{	print "No nodes in state 'idle'\n";
	}
} elsif ( $mode eq 'down' )
{	$tmp = get_down_nodes($nodehash);
	@NodesToDisplay = @$tmp;
	if ( @NodesToDisplay )
	{	dump_down_node($nodehash) if $header;
		foreach $nodename (@NodesToDisplay)
		{	dump_down_node($nodehash, $nodename);
		}
	} else
	{	print "No nodes in state 'down'\n";
	}
} elsif ( $mode eq 'node' )
{	$tmp = $scontrol->expand_hostnames($nodes);
	@NodesToDisplay = @$tmp;

	$jobs = get_jobs_on_node($nodes);
	update_jobs_hash(\%JobsByNode, $jobs);

	dump_node($nodehash, \%JobsByNode);
	foreach $nodename (@NodesToDisplay)
	{	dump_node($nodehash, \%JobsByNode, $nodename);
	}
} elsif ( $mode eq 'used' )
{ 	if ( scalar(@$jobs) || scalar(@$users) )
	{   	foreach my $job (@$jobs)
		{	$tmp = get_job($job);
			update_jobs_hash(\%JobsByNode, $tmp);
		}
		foreach my $user (@$users)
		{	$jobs = get_jobs_for_user($user);
			update_jobs_hash(\%JobsByNode, $jobs);
		}
	} else
	{	#Did not restrict to specific jobs/users
		$jobs = get_all_running_jobs();
		update_jobs_hash(\%JobsByNode, $jobs);
	}

	@NodesToDisplay = sort (keys %JobsByNode);
	dump_used_node($nodehash, \%JobsByNode);
	foreach $nodename (@NodesToDisplay)
	{	dump_used_node($nodehash, \%JobsByNode, $nodename);
	}
	
	

} elsif ( $mode eq 'all' )
{ 	@NodesToDisplay = sort nodename_sort_function (keys %$nodehash);
	$jobs = get_all_running_jobs();
	update_jobs_hash(\%JobsByNode, $jobs);
	
	dump_node($nodehash, \%JobsByNode);
	foreach $nodename (@NodesToDisplay)
	{	dump_node($nodehash, \%JobsByNode, $nodename);
	}

	if ( $dostats )
	{ 	my $stats = generate_stats($nodehash, $jobs);
		dump_stats($stats);
	}
	
}
