#!/usr/bin/env perl

use File::Basename;
use Cwd;
$cwd = dirname Cwd::abs_path($0);

$i = 0;
while ($i < @ARGV) {
    $arg = $ARGV[$i];
    if ($arg eq "-n") {
        $i++;
        $pnum = $ARGV[$i];
        $i++;
    } elsif ($arg eq "-f") {
        $i++;
        $file = $ARGV[$i];
        $i++;
    } elsif ($arg eq "-m") {
        $i++;
        $master = $ARGV[$i];
        $i++;
    } elsif ($arg eq "-h") {
        print_help();
        exit 0;
    } elsif ($arg =~ /^-/) {
        print_help();
        exit 1;
    } else {
        $pgms = "";
        $exe = $ARGV[$i];
        while ($i < @ARGV) {
            $pgms .= " ".$ARGV[$i];
            $i++;
        }
    }
}

$master = 0 unless defined $master;
$pnum = -1 unless defined $pnum;
print_help() and exit 1 unless defined $file;
print_help() and exit 1 if $pgms eq "";

open FH, "< $file" or die "cannot open '$file'.\n";
while ($line = <FH>) {
    chomp $line;
    if ($line =~ /^[\s\t]*(.*)[\s\t]*(.*)$/) {
        push @nodes, $1;
        push @options, $2;
    } else {
        die "'$file' format is wrong. (You can see help by 'dmimw -h'.)\n";
    }
}
close FH;
$pnum = @nodes if $pnum == -1;
;die "'$file' includes too few nodes.\n" if $pnum > @nodes;

$hostname = (`hostname`)[0];
$tmpfile = "dmimw_tmp_$hostname_$$.sh";
open FH, "> $tmpfile" or die "cannot open '$tmpfile'.\n";
print FH "#!/bin/bash\n\n";
print FH "cd `pwd`\n";
$orig_cmd = "$cwd/dmirun $options[$master] $pgms";
print FH "$orig_cmd &\n";
print FH "pid=\$!\n";
print FH "sleep 2\n";
$imax = $master >= $pnum ? $pnum - 1 : $pnum;
for ($i = 0; $i < $imax; $i++) {
    next if $i == $master;
    $cmd = "$cwd/dmirun -i $nodes[$master] $options[$i] $pgms";
    $cmd =~ s/\s*$//;
    push @cmds, $cmd;
    #print FH "ssh -f $nodes[$i] \"cd `pwd` && $cmd 2> hosei_result/${exe}_$nodes[$i].dat\"\n";
    print FH "ssh -f $nodes[$i] \"cd `pwd` && $cmd\"\n";
}
print FH "\n";
print FH "function handler()\n";
print FH "{\n";
print FH "    count=`expr \$count + 1`;\n";
print FH "    if [ \$count -gt 3 ]; then\n";
print FH "        exit\n";
print FH "    fi\n";
for ($i = 0; $i < $imax; $i++) {
    next if $i == $master;
    print FH "    ssh -f $nodes[$i] \"ps aux | grep '$pgms' | grep -v grep | awk '{print \\\$2}' | xargs kill -INT 2> /dev/null\"\n"
}
print FH "    kill -INT \$pid 2> /dev/null\n";
print FH "}\n";
print FH "\n";
print FH "sleep 6\n";
for ($i = 0; $i < $imax; $i++) {
    next if $i == $master;
    print FH "ssh -f $nodes[$i] \"ps aux | grep '$pgms' | grep -v grep | awk '{print \\\$2}' | xargs kill -INT 2> /dev/null\"\n"
}
print FH "kill -INT \$pid 2> /dev/null\n";
print FH "count=0\n";
print FH "trap \"handler\" 2\n";
print FH "for i in `seq 1 2`; do\n";
print FH "    wait \$pid && break 2> /dev/null\n";
print FH "done\n";
close FH;

$SIG{INT} = sub {
    system "kill -INT $pid 2> /dev/null\n";
};

$pid = fork;
if ($pid) {
    waitpid $pid, 0;
} else {
    exec "sh $tmpfile";
}

unlink $tmpfile or die "cannot delete '$tmpfile'.\n";

sub print_help {
    print STDERR 
        "dmimw : execute DMI processes in a master-worker manner

====== usage ======
dmimw   -f node_file
        [-n number_of_nodes]
        [-m index of master_node]
        [-h]
        program arguments

====== options ======
-f node_file : specify a node file which describes possible nodes on which DMI processes run
        The format of the node file is as follows : 
        
            node000
            node001
            node002 -s 10240000
            node003 -s 20480000
        
        You can describe one node with some options in each line.
        These options are used when a DMI process is spawned on each node.
        See 'dmirun -h' for more details of these options.
        DMI processes are spawned on the nodes specified in the i-th line (0 <= i < n), 
        where n is specified by -n option.
        If you omit -n option, DMI processes are spawned on all nodes in the node file.
        The m-th node becomes a master node and the rest become worker nodes, 
        where m is specified by -m option.
        If you omit -m option, the first node in the node file becomes a master node.
        
-n number_of_nodes : the number of nodes on which DMI processes run
-m master_node : specify which node becomes a master node
-h : print this message

See http://haraken.info/dmi/ for more details.

" 
}
