#include <stdio.h>
#include <unistd.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/resource.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <errno.h>
#include <mpi.h>

#include <NTL/ZZ.h>

#include "controller.h"
#include "worker.h"
#include "dpnode.h"

NTL_CLIENT

static int mpi_size;

static int controllerfds[2];

// Close fds 3 and up, except for the one given (pass -1 to close them
// all)
static void close_highfds_except(rlim_t exceptfd)
{
    // Find the max fd number
    struct rlimit limit;
    getrlimit(RLIMIT_NOFILE, &limit);
    for (rlim_t fd = 3; fd < limit.rlim_cur; ++fd) {
	if (fd != exceptfd) {
	    // There's no ill effect from closing a non-open fd, so just
	    // do it
	    close(fd);
	}
    }
}

static void boundcb(const char *boundaddr, unsigned short boundport)
{
    // Write the port and addr to the pipe
    write(controllerfds[1], &boundport, 2);
    write(controllerfds[1], boundaddr, strlen(boundaddr));
    close(controllerfds[1]);
}

void desired_resources(const ZZ &order, unsigned short &desired_dpnodes,
    unsigned int &max_workers, unsigned int &dpfreq)
{
    // How many DPnodes should we use for a problem of this size?
    desired_dpnodes = 1;
    // How many workers would we like to use?
    ZZ sorder = SqrRoot(order >> 46);
    if (NumBits(sorder) > 30) {
	// Just use all the workers we can find
	max_workers = 4294967295U;  // 2^32 - 1
    } else {
	max_workers = trunc_long(sorder,31) + 1;
    }

    // By default, 1 in 1000 points are distinguihed points.  The
    // number in the next line is 2^32/1000
    dpfreq = 4294967;
    if (order < 1000) {
	// Just make every point a DP
	dpfreq = 4294967295U;
    } else if (NumBits(order) < 27) {
	// The frequency of DPs should be 10/sqrt(order) to avoid
	// a DP-free cycle, so dpfreq = (10*2^32)/sqrt(order)
	ZZ f = (to_ZZ(10) << 32) / SqrRoot(order);
	dpfreq = trunc_long(f, 31);
    }
}

void redirect_output(const char* dir, const char* filename_base, int rank_index, int process_index, const char* filename_ext){
    // redirect stdout and stderr to a file
    char buf[512];
    snprintf(buf, sizeof(buf), "%s/%s_%d_%d.%s", dir, filename_base, rank_index, process_index, filename_ext);

    int fd = open(buf, O_RDWR|O_CREAT|O_TRUNC, S_IRUSR|S_IWUSR);
    dup2(fd, 1);
    dup2(fd, 2);
    close(fd);
}

static pid_t fork_and_remember(vector<pid_t> &children)
{
    pid_t pid = fork();
    if (pid > 0) {
	children.push_back(pid);
    }

    return pid;
}

int main(int argc, char **argv)
{
    // Init MPI
    MPI_Init(&argc, &argv);

    if (argc < 4) {
	printf("You must pass at least three arguments:\n");
	printf(" 1) number of workers to start per node/rank\n");
	printf(" 2) number of dpnodes to start per node/rank\n");
	printf(" 3) directory to store output files (stdout+stderr)\n");
	printf("The remaining arguments will be passed to the controller.\n");

	MPI_Finalize();
	exit(1);
    }

    char hostname[257];
    gethostname(hostname, 256);

    vector<pid_t> children;

    int rank = 0;

    int ret = 0;

    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);

    char boundportaddr[259];

    unsigned int num_workers = strtol(argv[1], NULL, 10);
    unsigned int num_dpnodes = strtol(argv[2], NULL, 10);

    if (rank == 0) {
	// Start the controller
	pipe(controllerfds);
	char fdenv[30];
	sprintf(fdenv, "%d", controllerfds[1]);
	setenv("CONTROLLER_BOUNDCB_FD", fdenv, 1);
	if (fork_and_remember(children) == 0) {
	    // Child; close the read half of the pipe and all other fds
	    close_highfds_except(controllerfds[1]);
	    char name[] = "controller";
	    argv[3] = name;
	    execv("./controller", argv+3);
	    return 1;
	} else {
	    // Parent; close the write half of the pipe
	    close(controllerfds[1]);
	    int res;

	    memset(boundportaddr, '\0', 259);
	    res = read(controllerfds[0], boundportaddr, 2);
	    if (res < 2) return 1;
	    res = read(controllerfds[0], boundportaddr+2, 256);
	    if (res < 1) return 1;
	    close(controllerfds[0]);

	    // std::cerr << "Parent reports controller bound to " << boundaddr << ":" << boundport << "\n";
	}
    }

    MPI_Bcast(boundportaddr, 259, MPI_CHAR, 0, MPI_COMM_WORLD);

    unsigned short boundport;
    memmove(&boundport, boundportaddr, 2);
    char portstr[10];
    sprintf(portstr, "%hu", boundport);
    const char *boundaddr = boundportaddr + 2;

    // The child will spawn its own children to be the workers
    for (unsigned int i=0; i<num_workers; i++) {
	if (fork_and_remember(children) == 0) {
	    close_highfds_except(-1);

	    redirect_output(argv[3], "worker", rank, i, "out");

	    char gpu_index[10];
	    snprintf(gpu_index, 10, "%d", i);
	    execl("./worker", "./worker", boundaddr, portstr, gpu_index, NULL);
	    return 1;
	}
    }

    // And dpnodes
    for (unsigned int i=0; i<num_dpnodes; i++) {
	if (fork_and_remember(children) == 0) {
	    close_highfds_except(-1);

	    redirect_output(argv[3], "dpnode", rank, i, "out");

	    execl("./dpnode", "./dpnode", boundaddr, portstr, NULL);
	    return 1;
	}
    }

    // Now wait for all the children
    vector<pid_t>::iterator pidit;
    for (pidit = children.begin(); pidit != children.end(); ++pidit) {
	int status;
	waitpid(*pidit, &status, 0);
    }

    MPI_Finalize();

    return ret;
}