#!/usr/bin/env perl

##**************************************************************
##
## Copyright (C) 1990-2007, Condor Team, Computer Sciences Department,
## University of Wisconsin-Madison, WI.
## 
## Licensed under the Apache License, Version 2.0 (the "License"); you
## may not use this file except in compliance with the License.  You may
## obtain a copy of the License at
## 
##    http://www.apache.org/licenses/LICENSE-2.0
## 
## Unless required by applicable law or agreed to in writing, software
## distributed under the License is distributed on an "AS IS" BASIS,
## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
## See the License for the specific language governing permissions and
## limitations under the License.
##
##**************************************************************



use strict;

scalar @ARGV > 0 or exiterror("Error, zero arguments were passed to $0\n");

# Receive the name of the job ad file.
my $jobadfile = shift @ARGV;
my $node = 0;
my $rshargs = "";

if( @ARGV ) {
    # We are being called to start up a comrade node
    $node = shift @ARGV;
    $rshargs = join " ", @ARGV;
}

open(JOBADS, "< $jobadfile") or exiterror("Can't open $jobadfile: $!\n");
my @jobads = <JOBADS>;
close JOBADS;

# Put the contents of the job ad file into our data structure.
my @jobadarray;
my $adnum = 0;
for( @jobads ) {
    if( /^\s*(\S+)\s*=\s*(\"?)(.*?)\"?\s*$/ ) {
	$jobadarray[$adnum]{$1}[0] = $3;
	$jobadarray[$adnum]{$1}[1] = $2;
    } elsif ( /^\s*...\s*$/ ) {
	$adnum++;
    } else {
	exiterror("Invalid line in ad file: $_\n");
    }
}


###################
# MPICH specific
# begins here
###################


my $clust;
my $proc;
if( defined $jobadarray[0]{'ClusterId'}[0] ) {
    $clust = $jobadarray[0]{'ClusterId'}[0];
} else {
    exiterror("Error, $jobadfile doesn't have a ClusterId\n");
}
if( defined $jobadarray[0]{'ProcId'}[0] ) {
    $proc = $jobadarray[0]{'ProcId'}[0];
} else {
    exiterror("Error, $jobadfile doesn't have a ProcId\n");
}

unless( $node ) {
    # We are doing the initial setup for the master node

    # Create our procgroup file.
    open(PROCGR,"> procgroup.$clust.$proc") or
	exiterror("Can't open procgroup.$clust.$proc: $!");
    my $currentnode = 0;
    while( $currentnode < scalar @jobadarray ) {
	$jobadarray[$currentnode]{'ThisHost'}[0] or
	    exiterror("Error, ad number $currentnode in $jobadfile " .
		"doesn't have a ThisHost\n");
	print PROCGR "$jobadarray[$currentnode]{'ThisHost'}[0] " .
	    ($currentnode ? "1" : "0") .
		" condor_exec\n";
	$currentnode++;
    }
    close PROCGR;

    # Add the procgroup file to the list of files to transfer
    if( $jobadarray[0]{'TransferInput'}[0] ) {
	if( $jobadarray[0]{'TransferInput'}[0] ne "procgroup.$clust.$proc" ) {
	    $jobadarray[0]{'TransferInput'}[0] .= ", procgroup.$clust.$proc";
	}
    } else {
	$jobadarray[0]{'TransferInput'}[0] = "procgroup.$clust.$proc";
    }
    $jobadarray[0]{'TransferInput'}[1] = '"';

    # Add -p4pg procgroup.x to the master ad
    $jobadarray[0]{'Args'}[0] = "-p4pg procgroup.$clust.$proc " . 
	$jobadarray[0]{'Args'}[0];
    $jobadarray[0]{'Args'}[1] = '"';

} else {
    # We are being called for a comrade node

    # Args given to rsh will look like this:
    # firebird.cs.wisc.edu -l bethenco -n condor_exec firebird.cs.wisc.edu
    # 35153  -p4amslave -p4yourname firebird.cs.wisc.edu

    $rshargs =~ s/.*?condor_exec\s+//;
    $jobadarray[$node]{'Args'}[0] = $rshargs ." $jobadarray[$node]{'Args'}[0]";
    $jobadarray[$node]{'Args'}[1] = '"';

    # Get rid of the old procgroup file
    unlink "procgroup.$clust.$proc"
	if( $node == 1 and -e "procgroup.$clust.$proc" );
}


###################
# MPICH specific
# ends here
###################


# Output the modified job ads
open(JOBADS,"> $jobadfile") or exiterror("Can't open > $jobadfile: $!");

my $currentnode = 0;
while( $currentnode < scalar @jobadarray ) {
    # Print MyType and TargetType first
    if( defined $jobadarray[$currentnode]{'MyType'}[0] ) {
	my $quotes = $jobadarray[$currentnode]{'MyType'}[1];
	print JOBADS "MyType = " .
	    "$quotes$jobadarray[$currentnode]{'MyType'}[0]$quotes\n";
    }
    if( defined $jobadarray[$currentnode]{'TargetType'}[0] ) {
	my $quotes = $jobadarray[$currentnode]{'TargetType'}[1];
	print JOBADS "TargetType = " .
	    "$quotes$jobadarray[$currentnode]{'TargetType'}[0]$quotes\n";
    }
    for( keys %{$jobadarray[$currentnode]} ) {
	next if( /MyType/ or /TargetType/ );
	my $quotes = $jobadarray[$currentnode]{$_}[1];
	print JOBADS "$_ = $quotes$jobadarray[$currentnode]{$_}[0]$quotes\n";
    }
    print JOBADS "...\n";
    $currentnode++;
}
close JOBADS;

sub exiterror {
    open(OBIT, "> scriptobit.$$"); # No point in error checking here ...
    print OBIT @_;
    close OBIT;
    die;
}
