/* -*- Mode: C; c-basic-offset:4 ; indent-tabs-mode:nil ; -*- */
/*
 *
 *  (C) 2009 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;
    int can_spawn;

    MTest_Init(&argc, &argv);

    errs += MTestSpawnPossible(&can_spawn);

    if (can_spawn) {
        /* 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 MTestReturnValue(errs);
        }

        /* Note that the MTest_Finalize get errs only over COMM_WORLD */
        if (parentcomm == MPI_COMM_NULL) {
            MTest_Finalize(errs);
        } else {
            MPI_Finalize();
        }
        /* free the argument vectors */
        for (i = 0; i < MAX_ARGV; i++) {
            free(inargv[i]);
            free(outargv[i]);
        }

    } else {
        MTest_Finalize(errs);
    }

    return MTestReturnValue(errs);
}

/* 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;
}