#include <pthread.h>

#include "scheduler.h"
#include "dbinterface.h"
#include "BasicStrategy.h"
#include <stdlib.h>

#define MAXDATASIZE 100
#define MACHINE_PERMANENT_PORT 22

DbInterface dbi;
Scheduler *theScheduler;
int event_occored = 0;               // Event Variable.

// The CONV_RSH environment variable for the system.
char SERVER_RSH[256];
// Path for conv-host for CHARM jobs.
char CONV_HOST[1024];

void default_options(){
    printf("scheduler <port> <nproc> [<nodelist file>] \n");
}

// reg_unreg == 0 means register, 1 means unregister
int register_with_central_manager(int reg_unreg, int my_port, 
				  char *dest, int port){

    int sockfd, numbytes;
    char buf[MAXDATASIZE];
    struct hostent *he;
    struct sockaddr_in their_addr; /* connector's address information */
    int status, argcount;

    if ((he=gethostbyname(dest)) == NULL) {  /* get the host info */
        perror("CMM gethostbyname");
        return 0;
    }

    if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) == -1) {
        perror("CMM socket");
        return 0;
    }
    
    their_addr.sin_family = AF_INET;         /* host byte order */
    their_addr.sin_port = htons(port);       /* short, network byte order */
    their_addr.sin_addr = *((struct in_addr *)he->h_addr);
    bzero(&(their_addr.sin_zero), 8);        /* zero the rest of the struct */

    // SIGPIPE is received if the socket stream connection is broken.
    signal(SIGPIPE, SIG_IGN);

    if (connect(sockfd, (struct sockaddr *)&their_addr,
		sizeof(struct sockaddr)) == -1) {
        perror("CMM connect");
        return 0;
    }

    // Prepare and send the reg/unreg command to the scheduler.
    char my_name[MAXDATASIZE+1];
    gethostname(my_name, MAXDATASIZE);
    char str[MAXDATASIZE+1];
    sprintf(str,"%s %s %d\n",
            (reg_unreg==0)? "CMRegister": "CMUnRegister",
            my_name, my_port);
    printf("Sending string [%s] to %s:%d\n", str, dest, port);
    write(sockfd, str, strlen(str));
    sprintf(str,"q\n");
    write(sockfd, str, strlen(str));

    close(sockfd);

    // Go back to default handling of broken streams
    signal(SIGPIPE,SIG_DFL);

    return 1; // success
}

int main(int argc, char ** argv){

    if(argc < 3){
	default_options();
	exit(1);
    }

    //register_with_central_manager(0, atoi(argv[2]), "bogota.ks.uiuc.edu", 1999);

    // get CONV_RSH.
    if(getenv("CONV_RSH") != NULL)
	strcpy(SERVER_RSH, getenv("CONV_RSH"));
    else
	strcpy(SERVER_RSH, "rsh");

    // conv-host in current dir.
    sprintf(CONV_HOST, "%s/conv-host", getenv("PWD"));
    printf("CONV_HOST = %s\n", CONV_HOST);
    printf("Starting scheduler.  Use Control-C to exit.\n");

    dbi.set_nproc(atoi(argv[2]));
    dbi.connect("thrift.cs", "sameer", "test");

    Scheduler sch(atoi(argv[1]), atoi(argv[2]));
    theScheduler = &sch;

    if(argc >= 4) {
	sch.load_nodelist(argv[3]);
    } else {
        printf("Loading nodelist file: scheduler_nodelist.\n");
	sch.load_nodelist("scheduler_nodelist");
    }
    
    sch.start_scheduler();
    return 0;
}

Scheduler::Scheduler(int port, int nproc){

    int sockfd;

    this->port = port;

    this->nproc = nproc;
    struct sockaddr_in my_addr;

    if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) == -1) {
        perror("socket");
        exit(1);
    }

    this->sch_socfd = sockfd;

    my_addr.sin_family = AF_INET;
    my_addr.sin_port = htons(port);
    my_addr.sin_addr.s_addr = INADDR_ANY;
    bzero(&(my_addr.sin_zero), 8);

    //Make port reusable immediately after the server crashes.
    int yes = 1;
    if (setsockopt (sockfd, SOL_SOCKET, SO_REUSEADDR, (char*)&yes, 
                    sizeof (yes)) == -1) {
	perror ("setsockopt");
	exit(1);
    }
    
    if (bind(sockfd, (struct sockaddr *)&my_addr,
	     sizeof(struct sockaddr)) == -1) {
        perror("bind");
        close(sockfd);
        exit(1);
    }

    // INITIALIZATIONS FROM THE DATABASE.

    free_proc_vector = new char[nproc];
    num_free_proc = nproc;

    for(int i = 0; i < nproc; i++)
	free_proc_vector[i] = 1;

    runq = dbi.fetch_running_jobs();

    Job *jptr = runq;
    while(jptr != NULL){
	for(int i = 0; i < nproc; i++)
	    if(jptr->bit_map[i]){
		free_proc_vector[i] = 0;
		num_free_proc --;
	    }
	if(jptr->type == MCHARM)
	    jptr->connect();
	jptr = jptr->next;
    }	
    
    max_port = dbi.get_max_port() + 1;
    if(max_port < port + 1)
	max_port = port + 1;

    strategy = new BasicStrategy(nproc);
    strategy->free_proc_vector = free_proc_vector;
    strategy->num_free_proc = num_free_proc;

    proc_mutex = new pthread_mutex_t;
    pthread_mutex_init(proc_mutex, NULL);
    
    remove_job();
}

// Background thread that queues jobs in the database.
void* schedule_thread_handler(void *arg){
    theScheduler->wait_for_request();
    return NULL;
}

// Execute thread
void* execute_thread_handler(void *arg){

    while(1){	
	theScheduler->execute_job();
	while(!event_occored)
	    sleep(2);
    }
    return NULL;
}

// Background thread that waits for jobs to finish.
void* remove_thread_handler(void *){

    while(1){
	theScheduler->remove_job();
	
	sleep(2);
    }
    return NULL;
}

void* activelist_manager_handler(void *){
   
    while(1){
	theScheduler->activelist_manager();
	
	sleep(30);
    }
    return NULL;
}

void Scheduler::start_scheduler(){
    
    pthread_t * cur_thread = new pthread_t;

    // Create a thread to accept requests from remote clients.
    pthread_create(cur_thread, NULL, schedule_thread_handler, NULL);

    cur_thread = new pthread_t;

    // Create a thread to execute jobs stored in the database.
    pthread_create(cur_thread, NULL, execute_thread_handler, NULL);

    cur_thread = new pthread_t;

    // Create a thread to manage the nodelist.
    pthread_create(cur_thread, NULL, activelist_manager_handler, NULL);
    
    // Wait for jobs to exit.
    remove_thread_handler(NULL);
}

void Scheduler::wait_for_request(){
    int client_fd;
    socklen_t sin_size;
    struct sockaddr_in their_addr;

    FILE * client_fp;

    if (listen(sch_socfd, MAX_CONN) == -1) {
	perror("listen");
	exit(1);
    }

    while(1){ //Accept loop.
	
        printf("================\nAwaiting faucet request:\n");

	sin_size = sizeof(struct sockaddr_in);
	if ((client_fd = accept(sch_socfd, (struct sockaddr *)
				&their_addr, &sin_size))== -1) {
	    //	    perror("accept");
	    continue;
	}
        printf("Received faucet request:\n");

	handle_request(client_fd);
    }
}

void Scheduler::handle_request(int fd){
    
    printf("In handle Request \n");
    FILE * client_fp;
    client_fp = fdopen(fd, "r+");
    
    //!!CHECK
    char *cmd = new char[MAX_SIZE];
    fgets(cmd, MAX_SIZE, client_fp);
    
    //    To Be Enabled!!!!!!!
    
    if (strncmp(cmd, "1 GETSTATUS",11) == 0) {
	dump_stats(client_fp);
	fclose(client_fp);
	close(fd);
	return;
    }
    
    if(strncmp(cmd, "2 KILL", 6) == 0){
	char *cur = strtok(cmd, " "); // the 2
	cur = strtok(cmd, NULL);      // the KILL
	cur = strtok(cmd, NULL);      // the dbid
	int dbid = 0;
	
	if(cur != NULL)
	    dbid = atoi(cur);
	
	Job *jptr = runq;
	while(jptr != NULL) {
	    if(jptr->dbid == dbid){
		jptr->Kill();
		break;
	    }
            jptr = jptr->next;
        }
	return;
    }
	
    if(strncmp(cmd, "1 LISTJOBS", 10) == 0){
	list_jobs(client_fp);
	fclose(client_fp);
	close(fd);
        return;
    }

    Job j(nproc);            
    j.client_soc_fd = fd;
    j.port = max_port++;
    j.parse_query(cmd);
    
    //CHECK FOR QUERRY JOBS  INCOMPLETE!!!
    Job *waitq = dbi.fetch_queued_jobs();
    
    if (strncmp(j.name, "QUERY", 5) == 0) {
	if (strategy->is_available(&j, waitq, runq) != -1)
	    fputs("yes\n", client_fp);
	else
	    fputs("no\n", client_fp);
	fclose(client_fp);
	close(fd);
	return;
    }
    
    printf("before acquiring the lock\n");
    pthread_mutex_lock(proc_mutex);
    
    dbi.insert_job(j);
    event_occored = 1;
    j.~Job();
    pthread_mutex_unlock(proc_mutex);
    return;
}

void *start_func(void *arg){

    theScheduler->start_job((Job *)arg);
    return NULL;
}

void Scheduler::execute_job(){

    Job *waitq;
    printf("in execute job \n");
    printf("before locking\n");
    pthread_mutex_lock(proc_mutex);

    printf("before fetch\n");
    waitq = dbi.fetch_queued_jobs();

    if(waitq == NULL){
	printf("waitq = NULL\n");
    }

    printf("after fetch\n");

    strategy->num_free_proc = num_free_proc;
    strategy->allocate_processors(waitq, runq);
    num_free_proc = strategy->num_free_proc;
    
    printf("after allocate processors \n");
    if(waitq)
	printf("bitmap_changed = %d, next = %d\n", waitq->bitmap_changed, 
	       waitq->next);

    Job *jptr = runq;
    while(jptr != NULL){
	if(jptr->type == MCHARM)
	    jptr->set_bitmap();
	jptr = jptr->next;
    }
    
    printf("after setting bitmaps\n");

    jptr = waitq;
    Job *prev = NULL;

    while(jptr != NULL){
	if(jptr->bitmap_changed){
	    pthread_t * cur_thread = new pthread_t;

	    // Create the job starting thread.
	    pthread_create(cur_thread, NULL, start_func, (void *) jptr);
	    //start_job(jptr);
	 
	    // Assign waitq
	    if(prev)
		prev->next = jptr->next;
	    else
		waitq = jptr->next;
	
	    // Assign runq	    
	    jptr->next = runq;
	    runq = jptr;
	    
	    // Assign jptr
	    if(prev)
		jptr = prev->next;
	    else
		jptr = waitq;
	}
	else{
	    prev = jptr;
	    jptr = jptr->next;
	}
    }

    // Now we inform the migratable jobs that their bitmap might have
    // been changed.

    event_occored = 0;
    dbi.free_list(waitq);
    
    pthread_mutex_unlock(proc_mutex);
    printf("after execute\n");
}

// The handler function for the job startup wait thread.
void Scheduler::Connect(Job *new_job){    

    int connect_failed = 0;

    if(new_job->type != MCHARM)
	return;

    new_job->ip = 0;
    
    //Connecting to the job, sleep while the job starts.
    if(new_job->connect() != 0){
	connect_failed = 1;
    }
    
    // CHECK *******
    if(!(new_job->status == FINISHED) && (!connect_failed)){
	pthread_mutex_lock(proc_mutex);
	new_job->set_bitmap();
	pthread_mutex_unlock(proc_mutex);
    }
    
    return;
}

int Scheduler::start_job(Job *j){
    int childpid;

    printf("start job\n");
    j->init_arg();

    char *pgm_nodesfile = create_nodes_file(j);
    if(j->type == MPI)
	j->argv[4] = pgm_nodesfile;
    else if(j->type == NCHARM)
	j->argv[3] = pgm_nodesfile;
    else if(j->type == MCHARM)
	j->argv[6] = pgm_nodesfile;
    
    if(j->working_directory != NULL)
	printf("%s \n", j->working_directory);
       
    if(!(childpid = fork())){

	// Set CONV_RSH environment variable.
	char buf[512];
	sprintf(buf, "CONV_RSH=%s", SERVER_RSH);
	putenv(buf);
	printf("%s\n", getenv("CONV_RSH"));

        // This is the child of the fork, i.e. the new job.
        // Connect Stdin, Stdout, Stderr to the client that requested
        // the job.

	setpgrp();

	close(0);
	close(1);
	close(2);
	
	if(j->working_directory != NULL)
	    chdir(j->working_directory);

	// Pipe the Stdin Stdout and Stderr to the required destinations. 
	int cur_fd = 0;
	if(strcmp(j->Stdin, "socket") == 0)
	    dup(j->client_soc_fd);
	else if(strcmp(j->Stdin, "null") == 0)
	    cur_fd = open("/dev/null", O_RDONLY);
	else if(j->Stdin != NULL)
	    cur_fd = open(j->Stdin, O_RDONLY);
	
	if(cur_fd == -1)
	    dup(j->client_soc_fd);

	cur_fd = 0;
	if(strcmp(j->Stdout, "socket") == 0)
	    dup(j->client_soc_fd);
	else if(strcmp(j->Stdout, "null") == 0)
	    cur_fd = open("/dev/null", O_WRONLY);
	else if(j->Stdout != NULL) 
	    cur_fd = open(j->Stdout, O_CREAT | O_WRONLY);
	
	if(cur_fd == -1)
	    dup(j->client_soc_fd);

	cur_fd = 0;
	if(strcmp(j->Stderr, "socket") == 0)
	    dup(j->client_soc_fd);
	else if(strcmp(j->Stderr, "null") == 0)
	    cur_fd = open("/dev/null", O_WRONLY);
	else if(j->Stderr != NULL)
	    cur_fd = open(j->Stderr, O_CREAT | O_WRONLY);

	if(cur_fd == -1)
	    dup(j->client_soc_fd);
	
	if(j->type == MCHARM)
	    execv(CONV_HOST, j->argv);
	else if(j->type == NCHARM)
	    execv(CONV_HOST, j->argv);
	else if(j->type == MPI)
	    execv("/usr/local/bin/mpirun", j->argv);
    }

    // This is the parent of the fork, i.e. the scheduler.
    j->pid = childpid;
    j->started();
    close(j->client_soc_fd);

    Connect(j);
    // For suns, inefficient otherwise.
    int status;
    wait(&status);
    return 0;
}

void Scheduler::remove_job(){
    
    pthread_mutex_lock(proc_mutex);
    Job *jptr = runq;    
    Job *prev = NULL;

    while(jptr != NULL){
	if(!jptr->ping()){
	    printf("In remove Job\n");
	    num_free_proc += jptr->delete_all(free_proc_vector);
	    event_occored = 1;
	    
	    if(prev){
		prev->next = jptr->next;
		//delete jptr;
		jptr = prev->next;
	    }
	    else{
		runq = jptr->next;
		//delete jptr;
		jptr = runq;
	    }		
	}
	else{
	    prev = jptr;
	    jptr = jptr->next;
	}
    }
    
    pthread_mutex_unlock(proc_mutex);
}

void Scheduler::load_nodelist(char * node_listfile){
    FILE *fp;
    char buf[256];
    int num_nodes = 0;

    strcpy(nodesfile, node_listfile);

    freelist = NULL;
    activelist = NULL;
    fp = fopen(node_listfile, "r");

    fscanf(fp, "%s", buf);
    Node *nptr = NULL;
    while(1){
	if(strcmp(buf, "host") == 0){
	    if(fscanf(fp, "%s", buf) == 1){
		Node *cur_node = new Node;
		strcpy(cur_node->name, buf);
		
		cur_node->next = NULL;
		
		if(num_nodes == 0)
		    activelist = cur_node;		
		
		if(num_nodes == nproc){
		    freelist = cur_node;
		    nptr = NULL;
		}

		if(nptr)
		    nptr->next = cur_node;		
		nptr = cur_node;
		num_nodes ++;
	    }
	    else break;
	}
	if(fscanf(fp, "%s", buf) != 1)
	    break;
    }

    if(num_nodes == 0){
	printf("incorrect scheduler_nodelist file\n");
	exit(1);
    }

    if(num_nodes < nproc){
	printf("#Processors is less than nproc\n");
	exit(1);
    }
}

// When starting to execute a new non-migratable job, we have to
// create a nodelist file telling it which CPU's to use.
char *Scheduler::create_nodes_file(Job *j){
    char *pgm_nodesfile = new char[20];
    sprintf(pgm_nodesfile, "nodelist.%d", j->dbid);

    char nodesfilepath[256];
    if(j->working_directory)
	sprintf(nodesfilepath, "%s/%s", j->working_directory, pgm_nodesfile);
    else 
	sprintf(nodesfilepath, "%s", pgm_nodesfile);

    FILE * fp = fopen(nodesfilepath, "w+");
    if((j->type == NCHARM) || (j->type == MCHARM))
	fprintf(fp, "group main\n");
    
    int i;
    Node *nptr;
    for(i=0, nptr = activelist; i < nproc; i++, nptr=nptr->next){
	if(j->type == MCHARM)
	    fprintf(fp, "host %s\n", nptr->name);
	else if(j->bit_map[i])
	    if(j->type == MPI)
		fprintf(fp, "%s\n", nptr->name);
	    else if(j->type == NCHARM)
		fprintf(fp, "host %s\n", nptr->name);
    }

    fclose(fp);
    return pgm_nodesfile;
}

Node *Scheduler::get_live_node(){
    Node *nptr = freelist;
    Node *prev = NULL;
    
    while(nptr != NULL)
	if(is_alive(nptr->name)){
	    if(prev)
		prev->next = nptr->next;
	    else
		freelist = freelist->next;
	    return nptr;
	}
	else{
	    prev = nptr;
	    nptr = nptr->next;
	}
    return NULL;
}

void Scheduler::activelist_manager(){

    // LOCK
    pthread_mutex_lock(proc_mutex);
    // There exists a mcharm job running on all processors.
    Job *jptr = runq;
    while(jptr != NULL)
	if(jptr->type == MCHARM){
	    // UNLOCK
	    pthread_mutex_unlock(proc_mutex);
	    return;
	}
	else jptr = jptr ->next;
    
    int node_changed = 0;

    Node *nptr = activelist;
    Node *prev = NULL;
    int pos = 0;
    while(nptr != NULL){
	if(!is_alive(nptr->name) && !job_running(pos)){
	    Node *node = get_live_node();
	    
	    if(node != NULL){

		if(prev){
		    prev->next = node;
		    node->next = nptr->next;
		}
		else{
		    activelist = node;
		    node->next = nptr->next;
		}

		node = nptr;
		nptr = nptr->next;

		// add down node to free list 
		node->next = freelist;
		freelist = node;
		node_changed = 1;
	    }
	}
	else{
	    prev = nptr;
	    nptr = nptr->next;
	}
	pos ++;
    }

    if(node_changed)
	update_nodesfile();

    //    printf("finished active list manager\n");
    pthread_mutex_unlock(proc_mutex);
}

int Scheduler::is_alive(char *node){
    int sockfd;

    struct hostent *he;
    struct sockaddr_in their_addr; /* connector's address information */

    if ((he=gethostbyname(node)) == NULL) {  /* get the host info */
	char buf[256];
	sprintf(buf, "Node %s :gethostbyname", node);
        perror(buf);
        return 0;
    }

    if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) == -1) {
	char buf[256];
	sprintf(buf, "Node %s :socket", node);
        perror(buf);
        return 0;
    }
    
    their_addr.sin_family = AF_INET;         /* host byte order */
    /* short, network byte order */
    their_addr.sin_port = htons(MACHINE_PERMANENT_PORT);   
    
    their_addr.sin_addr = *((struct in_addr *)he->h_addr);
    bzero(&(their_addr.sin_zero), 8);        /* zero the rest of the struct */

    // SIGPIPE is received if the socket stream connection is broken.
    signal(SIGPIPE, SIG_IGN);

    if (connect(sockfd, (struct sockaddr *)&their_addr,
		sizeof(struct sockaddr)) == -1) {
	char buf[256];
	sprintf(buf, "Node %s :connect", node);
        perror(buf);
        return 0;
    }

    close(sockfd);
    return 1;
}

int Scheduler::job_running(int pos){
    Job *jptr = runq;

    while(jptr != NULL)
	if(jptr->bit_map[pos])
	    return 1;
	else
	    jptr = jptr->next;

    return 0;
}

// Move to the database later.
void Scheduler::update_nodesfile(){
    FILE * fp = fopen(nodesfile, "w+");
    
    Node *nptr = activelist;

    while(nptr != NULL){
	fprintf(fp, "host %s\n", nptr->name);
	nptr = nptr->next;
    }

    nptr = freelist;

    while(nptr != NULL){
	fprintf(fp, "host %s\n", nptr->name);
	nptr = nptr->next;
    }
    fclose(fp);
}

void Scheduler::dump_stats(FILE *fp){

    int njobs = 0;
    Job *jptr = runq;
    
    while(jptr != NULL){
	njobs ++;
	jptr = jptr->next;
    }
    
    fprintf(fp, "%d %d\n", nproc, njobs);

    jptr = runq;
    while(jptr != NULL){
        fprintf(fp, "%s %d\n", jptr->name, jptr->num_allocated_proc);
        jptr = jptr->next;
    }
}

void Scheduler::list_jobs(FILE *fp){
    int njobs = 0;
    Job *jptr = runq;
    
    while(jptr != NULL){
	njobs ++;
	jptr = jptr->next;
    }
    
    //    fprintf(fp, "%d %d\n", nproc, njobs);

    jptr = runq;
    fprintf(fp, "Id\tName\t#Processors\tMinPE\tMaxPE\tStatus\tArguments\n");
    while(jptr != NULL){
        fprintf(fp, "%d\t%s\t%d\t%d\t%d\t%d\t%s\n", jptr->dbid, jptr->name, jptr->num_allocated_proc, jptr->min_proc, jptr->max_proc, jptr->status, jptr->argbuf);
        jptr = jptr->next;
    }
}
