/*
 * Copyright (C) by Argonne National Laboratory
 *     See COPYRIGHT in top-level directory
 */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include "process.h"
#include "ioloop.h"

struct iofds {
    int readOut[2], readErr[2];
};

typedef struct {
    char label[32];
    int lastNL;
} IOLabel;

int mypreamble(void *);
int myprefork(void *, void *, ProcessState *);
int mypostamble(void *, void *, ProcessState *);
int labeler(int, int, void *);

/*
   Test the argument parsing and processing routines.  See the Makefile
   for typical test choices
*/
int main(int argc, char *argv[], char *envp[])
{
    int rc;
    struct iofds newfds;

    MPIE_ProcessInit();
    MPIE_Args(argc, argv, &pUniv, 0, 0);
    MPIE_PrintProcessUniverse(stdout, &pUniv);
    MPIE_ForkProcesses(&pUniv.worlds[0], envp, mypreamble, &newfds, myprefork, 0, mypostamble, 0);
    MPIE_IOLoop(10);

    rc = MPIE_ProcessGetExitStatus();

    return rc;
}

/* Redirect stdout and stderr to a handler */
int mypreamble(void *data)
{
    struct iofds *newfds = (struct iofds *) data;
    /* Each pipe call creates 2 fds, 0 for reading, 1 for writing */
    if (pipe(newfds->readOut))
        return -1;
    if (pipe(newfds->readErr))
        return -1;
}

/* Close one side of each pipe pair and replace stdout/err with the pipes */
int myprefork(void *predata, void *data, ProcessState * pState)
{
    struct iofds *newfds = (struct iofds *) predata;

    close(newfds->readOut[0]);
    close(1);
    dup2(newfds->readOut[1], 1);
    close(newfds->readErr[0]);
    close(2);
    dup2(newfds->readErr[1], 2);
}

#include <fcntl.h>
/* Close one side of the pipe pair and register a handler for the I/O */
int mypostamble(void *predata, void *data, ProcessState * pState)
{
    struct iofds *newfds = (struct iofds *) predata;
    IOLabel *leader, *leadererr;
    static int wRank = 0;       /* TEMP */

    close(newfds->readOut[1]);
    close(newfds->readErr[1]);

    /* We need dedicated storage for the private data */
    leader = (IOLabel *) malloc(sizeof(IOLabel));
    snprintf(leader->label, sizeof(leader->label), "%d>", pState->wRank);
    leader->lastNL = 1;
    leadererr = (IOLabel *) malloc(sizeof(IOLabel));
    snprintf(leadererr->label, sizeof(leadererr->label), "err%d>", pState->wRank);
    leadererr->lastNL = 1;
    MPIE_IORegister(newfds->readOut[0], IO_READ, labeler, leader);
    MPIE_IORegister(newfds->readErr[0], IO_READ, labeler, leadererr);

    /* subsequent forks should not get these fds */
    fcntl(newfds->readOut[0], F_SETFD, FD_CLOEXEC);
    fcntl(newfds->readErr[0], F_SETFD, FD_CLOEXEC);
}

int labeler(int fd, int rdwr, void *data)
{
    IOLabel *label = (IOLabel *) data;
    char buf[1024], *p;
    int n;

    n = read(fd, buf, 1024);
    if (n == 0) {
        printf("Closing fd %d\n", fd);
        return 1;       /* ? EOF */
    }

    p = buf;
    while (n > 0) {
        int c;
        if (label->lastNL) {
            printf("%s", label->label);
            label->lastNL = 0;
        }
        c = *p++;
        n--;
        putchar(c);
        label->lastNL = c == '\n';
    }
    return 0;
}

void mpiexec_usage(const char *str)
{
    fprintf(stderr, "Usage: %s\n", str);
    return 0;
}
