/* -*- 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 <pthread.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>
#include "mpitest.h"

#include "connectstuff.h"

int main(int argc, char **argv)
{
    MPI_Comm tmp, comm, startComm;
    char *fname;
    char *actualFname = NULL;
    char *globalFname = NULL;
    int totalSize, expectedRank, size, cachedRank;
    char portName[MPI_MAX_PORT_NAME];
    int rankToAccept = 1;

    /* Debug - print out where we picked up the MPICH build from */
#ifdef MPICHLIBSTR
    msg("MPICH library taken from: %s\n", MPICHLIBSTR);
#endif

    if (argc != 4) {
        printf("Usage: %s <fname> <totalSize> <idx-1-based>\n", argv[0]);
        exit(1);
    }

    /* This is the base name of the file into which we write the port */
    fname = argv[1];
    /* This is the total number of processes launched */
    totalSize = atoi(argv[2]);
    /* Each process knows its expected rank */
    expectedRank = atoi(argv[3]) - 1;

    /* Start a watchdog thread which will abort after 120 seconds, and will
     * print stack traces using GDB every 5 seconds if you don't call
     * strokeWatchdog() */
    startWatchdog(120);

    /* Print a debug header */
    msg("Waiting for: %d - my rank is %d\n", totalSize, expectedRank);

    /* Singleton init */
    MPI_Init(0, 0);

    /* Duplicate from MPI_COMM_SELF the starting point */
    MPI_Comm_dup(MPI_COMM_SELF, &startComm);


    if (expectedRank == 0) {
        /* This process opens the port, and writes the information to the file */
        MPI_Open_port(MPI_INFO_NULL, portName);

        /* Write the port to fname.<rank> so that the connecting processes can
         * wait their turn by checking for the correct file to show up */
        actualFname = writePortToFile(portName, "%s.%d", fname, rankToAccept++);

        /* The wrapper script I'm using checks for the existance of "fname", so
         * create that - even though it isn't used  */
        globalFname = writePortToFile(portName, fname);
        installExitHandler(globalFname);

        comm = startComm;
    }
    else {
        char *readPort;
        readPort = getPortFromFile("%s.%d", fname, expectedRank);
        strncpy(portName, readPort, MPI_MAX_PORT_NAME);
        free(readPort);
        msg("Read port <%s>\n", portName);

        MPI_Comm_connect(portName, MPI_INFO_NULL, 0, startComm, &comm);
        MPI_Intercomm_merge(comm, 1, &tmp);
        comm = tmp;
        MPI_Comm_size(comm, &size);
        msg("After my first merge, size is now: %d\n", size);
    }
    while (size < totalSize) {
        /* Make sure we don't print a stack until we stall */
        strokeWatchdog();

        /* Accept the connection */
        MPI_Comm_accept(portName, MPI_INFO_NULL, 0, comm, &tmp);

        /* Merge into intracomm */
        MPI_Intercomm_merge(tmp, 0, &comm);

        /* Free the intercomm */
        MPI_Comm_free(&tmp);

        /* See where we're up to */
        MPI_Comm_rank(comm, &cachedRank);
        MPI_Comm_size(comm, &size);

        if (expectedRank == 0) {
            msg("Up to size: %d\n", size);

            /* Delete the old file, create the new one */
            unlink(actualFname);
            free(actualFname);

            /* Allow the next rank to connect */
            actualFname = writePortToFile(portName, "%s.%d", fname, rankToAccept++);
        }
    }
    MPI_Comm_rank(comm, &cachedRank);

    msg("All done - I got rank: %d.\n", cachedRank);

    MPI_Barrier(comm);

    if (expectedRank == 0) {

        /* Cleanup on rank zero - delete some files */
        MTestSleep(4);
        unlink(actualFname);
        free(actualFname);
        unlink(globalFname);
        free(globalFname);

        /* This lets my wrapper script know that we did everything correctly */
        indicateConnectSucceeded();
    }
    MPI_Finalize();

    return 0;
}