#!/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: Lists nodes with some information about their resources

Usage

$0 [ options ]

Allowed options are
--machine-readable: Print machine readable (CSV-like) output
--[no]state : Print state information
--[no]cpus  : Print CPU/core information
--[no]memory: Print memory information
--[no]features: Print features of node
--[no]gres  : Print generic resources for node

Default is --state --cpus --memory --features --gres

EOF
}

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

my $default_options = 
{	show_cores => 1,
	show_memory => 1,
	show_features => 1,
	show_gres => 1,
	show_state => 1,

	machine_readable => 0,
};

my %summary_info = ();

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

	my $node;
	$node = $nodehash->{$nodename} if $nodename;
	if ( $nodename )
	{	die "Nodename $nodename not in nodehash, aborting.\n" 
			unless $node && ref($node);
	}

	my $mach = $optionhash->{machine_readable};
	my $human = $mach?0:1;

	my $format = ($human?"%-6s ":"%s|") ; #nodename
	my @data = ();
	my $val;
	$val = ($node)?shorten_nodename($nodename):'Node';
	push @data, $val;

	if ( $optionhash->{show_state} )
	{ 	$val = ($node)?$node->state:'State';
		push @data, $val;
		$format .= $human?'%-9s ':'%s|';
	}

	if ( $optionhash->{show_cores} )
	{	$val = ($node)?($node->sockets . ':' . $node->cores .
			':' . $node->threads): 'S:C:T';
		push @data, $val;
		$format .= ($human? '%-8s ' : '%s|');

		$val = $node? $node->cpus : 'CPUs';
		push @data, $val;
		$format .= $human? '%-4s ' : '%s|';
	}

	if ( $optionhash->{show_memory} )
	{	$val = 'RAM(GB)';
		$val = sprintf "%d", $node->memory/1000 if $node;
		push @data, $val;
		$format .= $human? '%5s ' : '%s|';
	}

	if ( $optionhash->{show_features} )
	{	$val = 'Features';
		if ($node)
		{	my $feats = $node->features;
			$val = join ', ', @$feats;
		}
		push @data, $val;
		$format .= $human? '%12s ' : '%s|';
	}

	if ( $optionhash->{show_gres} )
	{	$val = 'Resources';
		if ($node)
		{	my $gres = $node->gres;
			my @keys = keys %$gres;
			@keys = grep { $_ ne '(null)' } @keys;
			@keys = grep { defined $gres->{$_} } @keys;
			my @tmp = map { "$_=$$gres{$_}" } @keys;
			$gres = join ", ", @tmp;
			$val = $gres;
			$val = '(null)' unless $val;
		}
		push @data, $val;
		$format .= $human? '%12s ' : '%s|';
	}


	$format =~ s/[ \|]$//;
	$format .= "\n";
	printf  "$format", @data;

	shift @data; #Remove nodename
	my $summ_ind = join '|', @data;
	$summary_info{$summ_ind}++ if $node;
}



sub dump_stats()
{	

	my @keys = keys %summary_info;
	@keys = sort @keys;
	
	print "\nSummary info: Node Counts by system properties\n";
	foreach my $key (@keys)
	{	my $cnt = $summary_info{$key};
		print "  $key :	$cnt\n";
	}
}

	
#Main
my $help;
my $options = $default_options;


my $res = GetOptions(
	'help!' => \$help,

	'machine-readable|machine_readable!' =>
		\$options->{machine_readable},
	'state!' => \$options->{show_state},
	'cpus!' => \$options->{show_cores},
	'memory!' => \$options->{show_memory},
	'features!' => \$options->{show_features},
	'gres|resources!' => \$options->{show_gres},
);

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

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

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

my $node_name_sorter = sub
{	my ( $junk1, $arack, $anum ) = split /-/, $a;
	my ( $junk2, $brack, $bnum ) = split /-/, $b;

	return ( $arack cmp $brack ) || ( $anum <=> $bnum );
};

@nodes = sort $node_name_sorter @nodes;

dump_node($nodehash, undef, $options);

foreach my $node (@nodes)
{	dump_node($nodehash, $node, $options);
}

dump_stats;
