#!/usr/bin/perl -w

##################################################################
## Current revision 	: 9					##
## Current version date	: 28 Jan 2004				##
## Last Modifier 	: Murali Ramsunder			##
##--------------------------------------------------------------##
## Originally created 	: 23 Jan 2004				##
## Original creation by	: Murali Ramsunder			##
##################################################################

##################################################################
## Function	: printUsage					##
##--------------------------------------------------------------##
## Input	: None	 					##
## Return	: None						##
##--------------------------------------------------------------##
## Action	: Print usage and exit				##
##		  -h, -help, -site, -jobmanager, -range, -queue ##
##		  -env, -localdir, -remotedir, -rsl		##
##################################################################
sub printUsage() {
	print STDERR "dagSubmit.pl [-h]";
	print STDERR " [OPTION VALUE] cmd-line\n";

	print STDERR "\n";
	print STDERR "\tincremental environment variable DAG_JOB_ID will be set \n";
	print STDERR "\tfor each job submitted\n";
	print STDERR "\n";

	## Help
	$line=sprintf("\t-%-"."s\t\t%s\n",'h','prints this');
	print STDERR $line;
	## Help
	$line=sprintf("\t-%-"."s\t%s\n",'help','prints this');
	print STDERR $line;

	print STDERR "\n";
	$line=sprintf("\t%s"."\t%s\n",'OPTION','VALUE');
	print STDERR $line;

	## Job Submit Site
	$line=sprintf("\t-%-"."s\t%s\n",'site','job submit site');
	print STDERR $line;
	$line=sprintf("\t\t\t\t%s\n",'DEFAULT: ligo-grid.aset.psu.edu');
	print STDERR $line;

	## JobManager
	$line=sprintf("\t-%-"."s\t%s\n",'jobmanager','fork | condor | pbs | ...');
	print STDERR $line;
	$line=sprintf("\t\t\t\t%s\n",'DEFAULT: pbs');
	print STDERR $line;

	## Number of Nodes
	$line=sprintf("\t-%-"."s\t%s\n",'range','range of nodes to submit the job');
	print STDERR $line;
	$line=sprintf("\t\t\t\t%s\n",'DEFAULT: 1; FORMAT: 1:10:2');
	print STDERR $line;

	## Queue
	$line=sprintf("\t-%-"."s\t%s\n",'queue','Job submit queue [OPTIONAL]');
	print STDERR $line;

	## Environment
	$line=sprintf("\t-%-"."s\t\t%s\n",'env','list of environment variable [OPTIONAL]');
	print STDERR $line;
	$line=sprintf("\t\t\t\t%s\n",'FORMAT: "VAR1=VAL1;VAR2=VAL2,..."');
	print STDERR $line;

	## Local Directory
	$line=sprintf("\t-%-"."s\t%s\n",'localdir','where condor files will reside [OPTIONAL]');
	print STDERR $line;
	$line=sprintf("\t\t\t\t%s\n",'DEFAULT: current working directory');
	print STDERR $line;

	## Remote Directory
	$line=sprintf("\t-%-"."s\t%s\n",'remotedir','where job will execute on remote machine [OPTIONAL]');
	print STDERR $line;

	## RSL params
	$line=sprintf("\t-%-"."s\t\t%s\n",'rsl','Additional Globus RSL params [OPTIONAL]');
	print STDERR $line;
	$line=sprintf("\t\t\t\t%s\n",'FORMAT: "(VAR1=VAL1)(VAR2=VAL2)..."');
	print STDERR $line;

	## Example Usage
	print STDERR "\n";
	print STDERR "dagSubmit.pl -site ligo-grid.aset.psu.edu -range 1-4:1 -jobmanager PBS \\ \n";
	print STDERR "\t-env \"TEST1=test1,NAME=dagSubmit\" /bin/hostname\n";
	print STDERR "\n";
	exit;
}

##################################################################
## Function	: changeDir					##
##--------------------------------------------------------------##
## Input	: List -> $dir 					##
## Return	: None						##
##--------------------------------------------------------------##
## Action	: Create directory as requested by user		##
##		  if $dir doesn't exist, prompt user to create 	##
##		  if create failes, prompt user to continue	##
##		  Also make OUTPUT, ERROR and LOGS within 	##
##		  submit directory				##
##################################################################
sub changeDir {
	my ($dir) = @_;
	## Change directory to store files
	my $ok = chdir "$dir";
	if(! $ok) {
		print STDERR "Directory does not exist\n";
		print STDERR "Would you like to create the directory \"$dir\"? [y/n] : ";
		my $response = <STDIN>;
		chomp $response;
		if($response eq 'y' || $response eq 'yes') {
			`mkdir $dir`;
			sleep 1;
			chdir "$dir" or die "ERROR: Unable to change directory, dying... \n";
		}
		else {
			print STDERR "Would you like to continue by using the current directory? [y/n] : ";
			$response = <STDIN>;
			chomp $response;
			print STDERR "Response: $response\n";
			if($response eq "y" || $response eq "yes") {
			}
			else {
				print STDERR "Aborting the program... \n";
				exit;
			}
		}
	}
	`mkdir OUTPUT` if not (-d "OUTPUT");
	`mkdir ERROR` if not (-d "ERROR");
	`mkdir LOGS` if not (-d "LOGS");
	if(!-d "OUTPUT" || !-d "ERROR" || !-d "LOGS") {
		print "Unable to create individual directories for output, error and logs\n";
	}
	print "Changing directory to \"$dir\"... \tDONE\n";
}

##################################################################
## Function	: generateCondorSubmitFiles			##
##--------------------------------------------------------------##
## Input	: List -> $exe, $args, $site, $jobmanager, $env,##
## 		  $queue, $rsl, $remotedir, $start, $end, $hop	##
## Return	: String -> $uniqueString			##
##--------------------------------------------------------------##
## Action	: Generate $nodes number of Condor Job Submit   ##
##		  files and the Condor DAG file if $nodes > 1 	##
##################################################################
sub generateCondorSubmitFiles {
	my ($exe, $args, $site, $jobmanager, $env, $queue, $rsl, $remotedir, $start, $end, $hop) = @_;

	select STDOUT;
	$|=1;
	
	## Unique string to create unique filenames
	my $uniqueString = (time^$$^unpack "%32L*", `ps -eaf | gzip`);
	$uniqueString="1234" if not defined $uniqueString;
	
	my $globusrsl;
	$globusrsl .= "(queue=$queue)" if defined $queue;
	$globusrsl .= "(directory=$remotedir)" if defined $remotedir;
	$globusrsl .= "$rsl" if defined $rsl;

	## If separate directories exist for OUTPUT, ERROR and LOGS
	## 	separate the files accordingly
	my $logfile="log.\$(cluster).$uniqueString" if not (-d "LOGS");
	my $outfile="out.\$(cluster).$uniqueString" if not (-d "OUTPUT");
	my $errfile="err.\$(cluster).$uniqueString" if not (-d "ERROR");

	$logfile="LOGS/log.\$(cluster).$uniqueString"   if (-d "LOGS");
	$outfile="OUTPUT/out.\$(cluster).$uniqueString" if (-d "OUTPUT");
	$errfile="ERROR/err.\$(cluster).$uniqueString"  if (-d "ERROR");
	
	print "Generating Condor Job Submit Files.";
	my $i=$start;
	while($i <= $end) {
		open OUTFILE, "> J$i.submit.$uniqueString" or die "\nERROR: Can't create file : $!";
		# print OUTFILE "TransferExecutable=false\n";
		print OUTFILE "executable=$exe\n";
		print OUTFILE "arguments=$args\n" if length($args);
		print OUTFILE "environment= DAG_JOB_ID=$i";
		print OUTFILE ";$env" if defined $env;
		print OUTFILE "\n";
		print OUTFILE "universe=globus\n";
		print OUTFILE "globusscheduler=$site/jobmanager-" . lc($jobmanager) . "\n";
		print OUTFILE "globusrsl=$globusrsl\n" if defined $globusrsl;
		print OUTFILE "log=$logfile\n";
		print OUTFILE "output=$outfile\n";
		print OUTFILE "error=$errfile\n";
		print OUTFILE "notification=NEVER\n";
		print OUTFILE "queue\n";
		close OUTFILE;
		print ".";
		$i += $hop;
	}
	print " \tDONE\n";

	## If we have only one node, we shouldn't have to do a DAG submit
	# return $uniqueString if $nodes == 1;
	if($start == 1 && $end == 1) {
		return $uniqueString;
	}

	$i=$start;
	open OUTFILE, "> submit.dag.$uniqueString" or die "\nERROR: Can't create DAG file : $!";
	print "Generating Condor DAG File...";
	while($i <= $end) {
		print OUTFILE "Job J$i J$i.submit.$uniqueString\n";
		$i += $hop;
	}
	close OUTFILE;
	print " \t\tDONE\n";
	
	return $uniqueString;
}

##################################################################
## Function	: gridProxyTest					##
##--------------------------------------------------------------##
## Input	: None						##
## Return	: None						##
##--------------------------------------------------------------##
## Action	: Check if valid proxy exists with minimum 	##
##		  5 minutes validity				##
## 		  Exits program if proxy is not valid		##
##################################################################
sub gridProxyTest() {
	my @gridResult = `grid-proxy-info -timeleft`;
	foreach (@gridResult) {
		if($_ =~ /ERROR/) {
			print STDERR "ERROR: Unable to locate valid user proxy\n";
			print STDERR "ERROR: Please run grid-proxy-init\n";
			exit;
		}
	}
	my $timeLeft = pop @gridResult;
	chomp $timeLeft;
	if($timeLeft < 300 || !defined($timeLeft)) {
		print STDERR "ERROR: Proxy is valid for only $timeLeft seconds [must be valid for atleast 5 minutes]\n";
		print STDERR "ERROR: Please run grid-proxy-init\n";
		exit;
	}
}

##################################################################
## Function	: checkCondor					##
##--------------------------------------------------------------##
## Input	: None						##
## Return	: None						##
##--------------------------------------------------------------##
## Action	: Check if the Condor executables exist		##
## 		  Exits program if executables are not present	##
##################################################################
sub checkCondor() {
	my @condorExe=`which condor_submit 2>\&1`;
	foreach(@condorExe) {
		chomp;
		if($_ =~ /which:/ || $_ =~ /no condor_submit/) {
			print STDERR "ERROR: Unable to find \"condor_submit\" in PATH\n";
			exit();
		}
	}
	@condorExe=`which condor_submit_dag 2>\&1`;
	foreach(@condorExe) {
		chomp;
		if($_ =~ /which:/ || $_ =~ /no condor_submit_dag/) {
			print STDERR "ERROR: Unable to find \"condor_submit_dag\" in PATH\n";
			exit();
		}
	}
}

################################
### Main body of the program ###
################################

use Getopt::Long;
my $oHelp;
my $oDir;
my $oSite = 'ligo-grid.aset.psu.edu';
my $oJobManager = 'pbs';
my $oRange;
my $oEnv;
my $oRsl;
my $oQueue;
my $cmd;
my $oRemoteDir;

## Parse the input command line parameters
GetOptions(
	"h"=>\$oHelp,
	"help"=>\$oHelp,
	"localdir=s"=>\$oDir,
	"remotedir=s"=>\$oRemoteDir,
	"rsl=s"=>\$oRsl,
	"queue=s"=>\$oQueue,
	"q=s"=>\$oQueue,
	"site=s"=>\$oSite,
	"jobmanager=s"=>\$oJobManager,
	"range=s"=>\$oRange,
	"env:s"=>\$oEnv
);

printUsage() if $Getopt::Long::error;
printUsage() if defined $oHelp;
	
## Construct the command line
foreach (@ARGV) {
	$cmd .= "$_ ";
}
printUsage() if not defined $cmd;

## Parse the range of nodes param
my $oStart = 1;
my $oEnd = 1;
my $oHop = 1;
if(defined($oRange)) {
	my $start;
	my $end;
	my $hop;
	($start,$end,$hop) = split(/:/,$oRange);
	$oStart = $start if defined $start;
	$oEnd = $end if defined $end;
	$oHop = $hop if defined $hop;
	if($oEnd < $oStart || !(defined($oStart))  || !(defined($oEnd)) || !(defined($oHop))) {
		print STDERR "Invalid range specified: $oRange\n";
		print STDERR "For usage, try dagSubmit.pl -h\n";
		exit;
	}
}

## Test if a valid user Grid proxy exists before proceeding
gridProxyTest();

## Test if Condor executables are present in PATH
checkCondor();

## Split the command line to executable and arguments
@cmd = split(/\s+/,$cmd);
my $exe=shift @cmd;
my $args=join " ", @cmd if defined @cmd;

## Change Directory as requested by the user
changeDir($oDir) if defined $oDir;

## Generate all Condor Job files
$fileString = generateCondorSubmitFiles($exe, $args, $oSite, $oJobManager, $oEnv, $oQueue, $oRsl, $oRemoteDir, $oStart, $oEnd, $oHop);

## If requested nodes = 1; simple Condor job; else Condor DAG job
if($oEnd > $oStart) {
	print "Submitting Condor DAG Job..."
}
else {
	print "Submitting Condor Job...";
}

## Print info for user
print "\n********************************\n";
print "Job related files are located in \t: ";
print `pwd`;
print "with filename extension \t\t: $fileString\n";

my $a;
my $b;
my @result;
if($oEnd > $oStart) {
	@result = `condor_submit_dag submit.dag.$fileString`;
}
else {
	@result = `condor_submit J$oStart.submit.$fileString`;
}
## Parse for Condor JobID
foreach(@result) {
	if ($_ =~ /submitted to cluster/) {
		($a,$b) = split(/cluster/,$_);
		chomp $b;
		print "Condor Job ID \t\t\t\t:$b";
	}
}
print "\n********************************\n";

1;