#include <stdio.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];

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 = 2;
    // 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);
    }
}

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

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

    int rank;

    int ret = 0;

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

    if (rank == 0) {
	// Start the controller
	pipe(controllerfds);
	if (fork() == 0) {
	    // Child; close the write half of the pipe
	    close(controllerfds[1]);
	    unsigned short boundport;
	    char boundaddr[257];
	    int res;

	    res = read(controllerfds[0], &boundport, 2);
	    if (res < 2) return 1;
	    res = read(controllerfds[0], boundaddr, 256);
	    if (res < 1) return 1;
	    close(controllerfds[0]);
	    boundaddr[res] = '\0';

	    std::cerr << "Child bound to " << boundaddr << ":" << boundport << "\n";

	    // The child will spawn two of its own children to be the
	    // workers
	    if (fork() == 0) {
		return worker_main(boundaddr, boundport);
	    } else if (fork() == 0) {
		return worker_main(boundaddr, boundport);
	    }

	    // And now become the dpnode
	    return dpnode_main(boundaddr, boundport);
	} else {
	    // Parent; close the read half of the pipe
	    close(controllerfds[0]);
	    unsigned short bindport;
	    Worklist worklist;

	    if (controller_parse_args(argc, argv, bindport, worklist)) {
		std::cerr << "Usage: " << argv[0] << " [-p listenport] N1 iter1 N2 iter2 ...\n";
		return 1;
	    }

	    ret = controller_main(worklist, bindport, boundcb);
	}
    }

    MPI_Finalize();

    return ret;
}