#!/usr/local/bin/perl -w
use English;
use Getopt::Std;

#
# nalloc - allocate nodes to an experiment. Takes a pid, and eid, and
# a list of nodes. Only nodes that are free can be allocated, and a
# special hack for sharks is included - allocating 'sh*' expands to
# allocation of the entire rack.
#

#
# Configure variables
#
my $TB       = "@prefix@";
use lib '@prefix@/lib';
use libdb;

my %opt = ();
getopts(\%opt,'v');

my $debug = 0;
if ($opt{v}) {
    $debug = 1;
}

if (@ARGV < 2) {
    die("Usage: nalloc [-v] <pid> <eid> <node> <node> <...>\n".
	 "		-v enables debugging output\n");
}

my $consetup="$TB/libexec/console_setup";
my $exportsetup="$TB/sbin/exports_setup";
my $error = 0;
my $pid = shift;
my $eid = shift;
my @node_names=@ARGV;
my @vals = ();
my @nodes= ();

if ($debug) { print "Expt '$eid', proj '$pid'\n"; }

#
# Make sure the user is a member of the correct project, and that the
# experiment exists
#
if (!ProjMember($pid)) {
    die "You are not a member of project $pid\n";
}
if (!ExpState($pid,$eid)) {
    die "There is no experiment $eid in project $pid\n";
}

######################################################################
# Step 1 - Make a list of nodes to reserve
#
# We prune nodes that are already reserved, etc., from the list, and
# do expansion of shark shelves
######################################################################

DBQueryFatal("lock tables nodes read, reserved write");

foreach my $n (@node_names) {
    # Shark hack
    if ($n =~ /(sh\d+)/ ) { $n = $1."-1"; }
    # End shark hack

    my ($r_pid, $r_eid);
    if (NodeidToExp($n,\$r_pid,\$r_eid)) {
	# Someone has already reserved this node
	if (($r_pid eq $pid) && ($r_eid eq $eid)) {
	    print "You have already reserved node '$n'.\n";
	    # Do not increment error code since that throws off tbprerun.
	} else {
	    print "Someone else has already reserved node '$n'.\n";
	    $error++;
	}
	next; # Go on to the next node if this one is reserved
    }

    my $result = DBQueryFatal("select * from nodes where node_id='$n'");
    if ($result->numrows() < 1) {
      print "Node '$n' does not exist.\n";
      $error++;
      next;
    } else {
	# No one has reserved it, and it exists, so add it to my list
	# Shark hack
	if ($n =~ /(sh\d+)/ ) {
	    # its a shark - do the whole shelf if its not done already.
	    my $shelf = $1;
	    if ( ! (join(",",@nodes) =~ /,$shelf-\d,/)) {
		# Shelf hasn't been done yet...
		foreach my $n ( 1 .. 8 ) {
		    push(@vals,"('$shelf-$n','$pid','$eid')");
		    push(@nodes,"$shelf-$n");
		}
	    }
	    # End shark hack
	} else {
	    # its not a shark - just add it in...
	    push(@vals,"('$n','$pid','$eid')");
	    push(@nodes,"$n");
	}
    }
}

if ($debug) { print "List Ready: @vals\nError=$error\n"; }

######################################################################
# Step 2 - Make the reservations in the database
#
# Uses the list built in step 1
######################################################################

if (!$error && @vals) {
  print "Reserving nodes...";
  my $cmd = "insert into reserved (node_id,pid,eid) values ".join(",",@vals);
  if (!DBQueryWarn($cmd)) { $error++; }
}

DBQueryFatal("unlock tables");

######################################################################
# Step 3 - Setup consoles
#
# Uses the list built in step 1
######################################################################

if (!$error && @nodes) {
    my @conlist=();
    my @sharks=();

    foreach $n ( @nodes ) { 
	if ($n =~ /(sh\d+)/) {
	    # Its a shark - do the shelf if it hasn't been done yet
	    my $shelf = $1;
	    if (!(join(",",@sharks) =~ /\b$shelf\b/)) {
		push(@sharks,$shelf);
		push(@conlist,$shelf);
	    }
	}
	else {
	    push(@conlist,$n);
	}
    }
    if ($debug) { print "Console setup list: @conlist\n"; }

    system("$consetup @conlist") == 0 or
	print STDERR "WARNING: $consetup @conlist failed!\n";
}

if ($debug) { print "Returning with value $error\n"; }
exit($error);