/* -*- Mode: C; c-basic-offset:4 ; indent-tabs-mode:nil ; -*- */
/*
 *
 *  (C) 2008 by Argonne National Laboratory.
 *      See COPYRIGHT in top-level directory.
 */
#include "mpi.h"
#include <stdio.h>
#include <stdlib.h>
#include "mpitest.h"
#include "mpitestconf.h"
#ifdef HAVE_STRING_H
#include <string.h>
#endif

/*
static char MTEST_Descrip[] = "A simple test of Comm_spawn, with many arguments";
*/

#define MAX_ARGV 1024
int worker( int argc, char *argv[], 
	    MPI_Comm intercomm, char *outargv[], int np );

int main( int argc, char *argv[] )
{
    int errs = 0, err;
    int rank, size, rsize, i;
    int np = 2;
    int errcodes[2];
    MPI_Comm      parentcomm, intercomm;
    char *inargv[MAX_ARGV];
    char *outargv[MAX_ARGV];
    int narg = 40;
    char *saveArgp = 0;

    MTest_Init( &argc, &argv );

    /* Initialize the argument vectors */
    for (i=0; i<MAX_ARGV; i++) {
      char *p;
      p = (char *)malloc( 9 );
      if (!p) {
	  fprintf( stderr, "Unable to allocated memory\n" );
	  MPI_Abort( MPI_COMM_WORLD, 1 );
      }
      strcpy( p, "01234567" );
      inargv[i] = p;
      p = (char *)malloc( 9 );
      if (!p) {
	  fprintf( stderr, "Unable to allocated memory\n" );
	  MPI_Abort( MPI_COMM_WORLD, 1 );
      }
      strcpy( p, "01234567" );
      outargv[i] = p;
    }

    MPI_Comm_get_parent( &parentcomm );

    if (parentcomm == MPI_COMM_NULL) {
	for (narg=16; narg < 512; narg += narg) {
	    /* Create 2 more processes */
	    /* Set a null at argument length narg */
	    saveArgp = inargv[narg];
	    inargv[narg] = 0;
	    /* ./ is unix specific .
	       The more generic approach would be to specify "spawnmanyarg" as 
	       the executable and pass an info with ("path", ".") */
	    MPI_Comm_spawn( (char*)"./spawnmanyarg", inargv, np,
			    MPI_INFO_NULL, 0, MPI_COMM_WORLD,
			    &intercomm, errcodes );
	    inargv[narg] = saveArgp;
	    /* We now have a valid intercomm */
	    
	    /* Master */
	    MPI_Comm_remote_size( intercomm, &rsize );
	    MPI_Comm_size( intercomm, &size );
	    MPI_Comm_rank( intercomm, &rank );
	    
	    if (rsize != np) {
		errs++;
		printf( "Did not create %d processes (got %d)\n", np, rsize );
	    }
	    /* Send the expected rank */
	    for (i=0; i<rsize; i++) {
		MPI_Send( &i, 1, MPI_INT, i, 0, intercomm );
	    }
	    /* Send the number of arguments */
	    for (i=0; i<rsize; i++) {
		MPI_Send( &narg, 1, MPI_INT, i, 0, intercomm );
	    }
	    /* We could use intercomm reduce to get the errors from the 
	       children, but we'll use a simpler loop to make sure that
	       we get valid data */
	    for (i=0; i<rsize; i++) {
		MPI_Recv( &err, 1, MPI_INT, i, 1, intercomm, MPI_STATUS_IGNORE );
		errs += err;
	    }
	    
	    /* Free the intercomm before the next round of spawns */
	    MPI_Comm_free( &intercomm );
	}
    }
    else {
	/* Note that worker also send errs to the parent */
	errs += worker( argc, argv, parentcomm, outargv, np );
	MPI_Comm_free( &parentcomm );
	MPI_Finalize();
	return 0;
    }

    /* Note that the MTest_Finalize get errs only over COMM_WORLD */
    if (parentcomm == MPI_COMM_NULL) {
	MTest_Finalize( errs );
    }

    MPI_Finalize();
    return 0;
}

/* Call this routine if this process is the spawned child */
int worker( int argc, char *argv[], 
	    MPI_Comm intercomm, char *outargv[], int np )
{
    int i, narg;
    int rsize, size, rank;
    int errs = 0;
    char *saveoutArgp = 0;
    MPI_Status status;

    MPI_Comm_remote_size( intercomm, &rsize );
    MPI_Comm_size( intercomm, &size );
    MPI_Comm_rank( intercomm, &rank );

    /* Child */
    /* FIXME: This assumes that stdout is handled for the children
       (the error count will still be reported to the parent) */
    if (size != np) {
	errs++;
	printf( "(Child) Did not create %d processes (got %d)\n", 
		np, size );
    }
    MPI_Recv( &i, 1, MPI_INT, 0, 0, intercomm, &status );
    if (i != rank) {
	errs++;
	printf( "Unexpected rank on child %d (%d)\n", rank, i );
    }
    MPI_Recv( &narg, 1, MPI_INT, 0, 0, intercomm, MPI_STATUS_IGNORE );
    saveoutArgp = outargv[narg];
    outargv[narg] = 0;
    /* Check the command line */
    for (i=1; i<argc; i++) {
	if (!outargv[i-1]) {
	    errs++;
	    printf( "Wrong number of arguments (%d)\n", argc );
	    break;
	}
	if (strcmp( argv[i], outargv[i-1] ) != 0) {
	    errs++;
	    printf( "Found arg %s but expected %s\n", argv[i],
		    outargv[i-1] );
	}
    }
    if (outargv[i-1]) {
	/* We had too few args in the spawned command */
	errs++;
	printf( "Too few arguments to spawned command (only %d)\n", i );
    }
    /* Restore the argument vector (not necessary in this case, since the 
       worker will exit) */
    outargv[narg] = saveoutArgp;
    /* Send the errs back to the master process */
    MPI_Ssend( &errs, 1, MPI_INT, 0, 1, intercomm );

   return errs;
}
