/*
 * cush - the customizable shell.
 *
 * Developed by Godmar Back for CS 3214 Summer 2020 
 * Virginia Tech.  Augmented to use posix_spawn in Fall 2021.
 */

#define __USE_GNU  /* need to set POSIX_SPAWN pgid */
#define _GNU_SOURCE    1
#include <stdio.h>
#include <readline/readline.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <termios.h>
#include <sys/wait.h>
#include <assert.h> 
#include <fcntl.h>  /* need for redirection */
#include <time.h>

/* Since the handed out code contains a number of unused functions. */
#pragma GCC diagnostic ignored "-Wunused-function"

#include "termstate_management.h"
#include "signal_support.h" 
#include "shell-ast.h"
#include "utils.h"
#include "../posix_spawn/spawn.h"  /* need to use POSIX_SPAWN */

static void sigchld_handler(int sig, siginfo_t *info, void *_ctxt);
/* Additional functions */
static void bin_cmd_custom_prompt(char *argv); 
static void build_changed_prompt(char *prompt_received, int size); 
static void bin_cmd_cd(char *arg); 

/* Menu to inform user */
static void
usage(char *progname)
{
    printf("Usage: %s -h\n"
        " -h            print this help\n"
	" jobs		\n"
	" fg		\n"
	" bg		\n"
	" kill		\n"
	" stop		\n"
	" cd    	\n"
    " custom_prompt	\n"
    " exit    	\n",
        progname);

    exit(EXIT_SUCCESS);
}

/* Build a prompt */
static char *
build_prompt(void)
{
    return strdup("cush> ");
}

/* Every single job can have four states */
enum job_status {
    FOREGROUND,     /* job is running in foreground.  Only one job can be
	                       in the foreground state. */
    BACKGROUND,     /* job is running in background */
    STOPPED,        /* job is stopped via SIGSTOP */
    NEEDSTERMINAL,  /* job is stopped because it was a background job
                       and requires exclusive terminal access */
};

/* Data structure to manage jobs */
struct job {
    struct list_elem    elem;   /* Link element for jobs list. */
    struct ast_pipeline *pipe;  /* The pipeline of commands this job represents */
    int  jid;                   /* Job id. */
    enum job_status status;     /* Job status. */ 
    int  num_processes_alive;   /* The number of processes that we know to be alive */
    struct termios saved_tty_state;  /* The state of the terminal when this job was 
                                        stopped after having been in foreground */

    /* Add additional fields here if needed. */
    pid_t pgid; /* process group id */
};

/* Utility functions for job list management.
 * We use 2 data structures: 
 * (a) an array jid2job to quickly find a job based on its id
 * (b) a linked list to support iteration
 */
#define MAXJOBS (1<<16)
#define MAXPIDS 1000

static struct list job_list;
static struct job * jid2job[MAXJOBS];
static struct job * pid2job[MAXPIDS];
static struct job *jobs2refresh[1000];
struct termios *terminal;
int job_jid_fg; 
int jobs_to_refresh = 0;
int num_of_command = 0;

/* Return job corresponding to jid */
static struct job * 
get_job_from_jid(int jid)
{
    if (jid > 0 && jid < MAXJOBS && jid2job[jid] != NULL)
        return jid2job[jid];

    return NULL;
}

/* Add a new job to the job list */
static struct job *
add_job(struct ast_pipeline *pipe)
{
    struct job * job = malloc(sizeof *job);
    job->pipe = pipe;
    job->num_processes_alive = 0;
    list_push_back(&job_list, &job->elem);
    for (int i = 1; i < MAXJOBS; i++) {
        if (jid2job[i] == NULL) {
            jid2job[i] = job;
            job->jid = i;
            return job;
        }
    }
    fprintf(stderr, "Maximum number of jobs exceeded\n");
    abort();
    return NULL;
}

/* Delete a job.
 * This should be called only when all processes that were
 * forked for this job are known to have terminated.
 */
static void
delete_job(struct job *job)
{
    int jid = job->jid;
    assert(jid != -1);
    jid2job[jid]->jid = -1;
    jid2job[jid] = NULL;
    ast_pipeline_free(job->pipe);
    free(job);
}

/* Get the status of a job */
static const char *
get_status(enum job_status status)
{
    switch (status) {
    case FOREGROUND:
        return "Foreground";
    case BACKGROUND:
        return "Running";
    case STOPPED:
        return "Stopped";
    case NEEDSTERMINAL:
        return "Stopped (tty)";
    default:
        return "Unknown";
    }
}

/* Print the command line that belongs to one job. */
static void
print_cmdline(struct ast_pipeline *pipeline)
{
    struct list_elem * e = list_begin (&pipeline->commands); 
    for (; e != list_end (&pipeline->commands); e = list_next(e)) {
        struct ast_command *cmd = list_entry(e, struct ast_command, elem);
        if (e != list_begin(&pipeline->commands))
            printf("| ");
        char **p = cmd->argv;
        printf("%s", *p++);
        while (*p)
            printf(" %s", *p++);
    }
}

/* Print a job */
static void
print_job(struct job *job)
{

	printf("[%d]\t%s\t\t(", job->jid, get_status(job->status));
    print_cmdline(job->pipe);
    printf(")\n");
}

/*  Refresh a finished job */
static void
refresh_job(struct job *job)
{
	struct list_elem *elem = &(job->elem);

	list_remove(elem);
	delete_job(job);
}

/*  Refresh every single finshed jobs */
static void
refresh_jobs(void) 
{
	for (int i = 0; i < jobs_to_refresh; i++) {
		struct job *job = jobs2refresh[i];
		jobs2refresh[i] = NULL;
		refresh_job(job);		
	}	
	jobs_to_refresh = 0;
}

/* Utility functions for child process & signal handling */

/* Child status handler */
static void
handle_child_status(pid_t pid, int status)
{
    // assert(signal_is_blocked(SIGCHLD));

    /* To be implemented.
     * Step 1. Given the pid, determine which job this pid is a part of
     *         (how to do this is not part of the provided code.)
     * Step 2. Determine what status change occurred using the
     *         WIF*() macros.
     * Step 3. Update the job status accordingly, and adjust
     *         num_processes_alive if appropriate.
     *         If a process was stopped, save the terminal state.
     */
    struct job *job = pid2job[(int)pid % MAXPIDS];

	if (WIFSTOPPED(status)) { 
			job->status = STOPPED;
			print_job(job);
	}

	if (WTERMSIG(status) == SIGABRT) {
			printf("aborted\n");
	}

	if (WTERMSIG(status) == SIGSEGV) {
			printf("segmentation fault\n");
	}

	if (WTERMSIG(status) == SIGKILL) {
			printf("killed\n");
	}

	if (WTERMSIG(status) == SIGTERM) {
			printf("terminated\n");
	}

	if (WTERMSIG(status) == SIGFPE) {
			printf("floating point exception\n");
	}
	
	if ((status == 0) || (WIFSIGNALED(status)) || WIFEXITED(status)){ 
			job->num_processes_alive--;

			if (job->num_processes_alive == 0) {
					if ((job->status == BACKGROUND) && (status == 0)) {
					}
					jobs2refresh[jobs_to_refresh] = job;
					jobs_to_refresh++;
					pid2job[(int)pid % MAXPIDS] = NULL;
			}
	}

}


/* Wait for all processes in this job to complete, or for
 * the job no longer to be in the foreground.
 * You should call this function from a) where you wait for
 * jobs started without the &; and b) where you implement the
 * 'fg' command.
 * 
 * Implement handle_child_status such that it records the 
 * information obtained from waitpid() for pid 'child.'
 *
 * If a process exited, it must find the job to which it
 * belongs and decrement num_processes_alive.
 *
 * However, not that it is not safe to call delete_job
 * in handle_child_status because wait_for_job assumes that
 * even jobs with no more num_processes_alive haven't been
 * deallocated.  You should postpone deleting completed
 * jobs from the job list until when your code will no
 * longer touch them.
 *
 * The code below relies on `job->status` having been set to FOREGROUND
 * and `job->num_processes_alive` having been set to the number of
 * processes successfully forked for this job.
 */
static void
wait_for_job(struct job *job)
{
    // assert(signal_is_blocked(SIGCHLD));

    while (job->status == FOREGROUND && job->num_processes_alive > 0) {
		
        int status;

        pid_t child = waitpid(-1, &status, WCONTINUED|WUNTRACED);

        // When called here, any error returned by waitpid indicates a logic bug in the shell.
        // In particular, ECHILD "No child process" means that there has
        // already been a successful waitpid() call that reaped the child, so
        // there's likely a bug in handle_child_status where it failed to update
        // the "job" status and/or num_processes_alive fields in the required
        // fashion.
        // Since SIGCHLD is blocked, there cannot be races where a child's exit
        // was handled via the SIGCHLD signal handler.
        if (child != -1)
            handle_child_status(child, status);
        else
            utils_fatal_error("waitpid failed, see code for explanation");
    }
}

/*
 * Suggested SIGCHLD handler.
 *
 * Call waitpid() to learn about any child processes that
 * have exited or changed status (been stopped, needed the
 * terminal, etc.)
 * Just record the information by updating the job list
 * data structures.  Since the call may be spurious (e.g.
 * an already pending SIGCHLD is delivered even though
 * a foreground process was already reaped), ignore when
 * waitpid returns -1.
 * Use a loop with WNOHANG since only a single SIGCHLD 
 * signal may be delivered for multiple children that have 
 * exited. All of them need to be reaped.
 */
static void
sigchld_handler(int sig, siginfo_t *info, void *_ctxt)
{
    pid_t child;
    int status;

    assert(sig == SIGCHLD);

	signal_block(SIGCHLD);
    while ((child = waitpid(-1, &status, WNOHANG|WUNTRACED)) > 0) {
		if (child != -1){
            handle_child_status(child, status);
		}
        else {
            utils_fatal_error("waitpid failed, see code for explanation");
    	}
    }

	signal_unblock(SIGCHLD);
}

/* SIGINT handler */
static void
sigint_handler(int signal)
{
	assert(signal == SIGINT);

	termstate_give_terminal_back_to_shell();

	printf("\ncush> ");
}

/* SIGTSTP handler */
static void
sigtstp_handler(int signal)
{
	assert(signal == SIGTSTP);

	struct job *job;

	if (job_jid_fg != -1) {
		job = get_job_from_jid(job_jid_fg);
		kill(job->pgid, 20);
		job_jid_fg = -1;
	}

}


/* Utility functions for built-in commands execution */

static char *bin_cmds[] = {"jobs", "fg", "bg", "kill", "stop", "cd", "custom_prompt", "exit"};
bool change_prompt = false;
char *myPrompt;

/* Built-in command "jobs" */
static void
bin_cmd_jobs(void)
{
	for(struct list_elem *elem = list_begin(&job_list);
			elem != list_prev(list_end(&job_list));
			elem = list_next(elem)) {

		struct job *job = list_entry(elem, struct job, elem);

		print_job(job);
	}
}

/* Built-in command "fg (foreground)" */
static void
bin_cmd_fg(char *cmd)
{
	struct job *job;
	int job_id = atoi(cmd);

	job = get_job_from_jid(job_id);
	job->status = FOREGROUND;
	job_jid_fg = job->jid;
	kill(job->pgid, SIGCONT); 
	tcsetpgrp(1, job->pgid);
	print_cmdline(job->pipe);
    printf("\n");
	termstate_give_terminal_to(terminal, job->pgid);
	wait_for_job(job);
	//termstate_restore(terminal);
	termstate_give_terminal_back_to_shell();
	job_jid_fg = -1;
}

/* Built-in command "bg (background)" */
static void
bin_cmd_bg(char *cmd)
{
	struct job *job;
	int job_id = atoi(cmd);
	int pgroup_id;

	job = get_job_from_jid(job_id);
	pgroup_id = (job->pgid) * -1;
	kill(pgroup_id, SIGCONT); 
	job->status = BACKGROUND;
	print_job(job);
}

/* Built-in command "kill" */
static void
bin_cmd_kill(char *cmd)
{
	struct job *job;
	int job_id = atoi(cmd);
	int pgroup_id;

	job = get_job_from_jid(job_id);
	pgroup_id = (job->pgid) * -1;
	kill(pgroup_id, SIGINT);
}

/* Built-in command "stop" */
static void
bin_cmd_stop(char *cmd)
{
	struct job *job;
	int job_id = atoi(cmd);
	int pgroup_id;

	job = get_job_from_jid(job_id);
	pgroup_id = (job->pgid) * -1;
	kill(pgroup_id, SIGTSTP);
}

/* Built-in command "exit" */
static void
bin_cmd_exit(char *cmd)
{
	if (cmd != NULL) 
		exit(atoi(cmd));
	else 
		exit(0);
}

/* Built-in commands execution main function */
static void
bin_cmd_main(char **argv)
{
	char *bin_cmd = argv[0];

	if (strcmp(bin_cmd, "jobs") == 0) {
		bin_cmd_jobs();
	}
	if (strcmp(bin_cmd, "fg") == 0) {
		bin_cmd_fg(argv[1]);
	}
	if (strcmp(bin_cmd, "bg") == 0) {
		bin_cmd_bg(argv[1]);
	}
	if (strcmp(bin_cmd, "kill") == 0) {
		bin_cmd_kill(argv[1]);
	}
	if (strcmp(bin_cmd, "stop") == 0) {
		bin_cmd_stop(argv[1]);
	}
	if (strcmp(bin_cmd, "exit") == 0) {
		bin_cmd_exit(argv[1]);
	}
    if (strcmp(bin_cmd, "cd") == 0) {
		bin_cmd_cd(argv[1]);
	}
	if (strcmp(bin_cmd, "custom_prompt") == 0){
		bin_cmd_custom_prompt(argv[1]);
	}
}

/* main function of cush, the customizable shell.
 * executes in 4 steps
 * 1. Initialize
 * 2. Read/Eval command line input
 * 3. Process job (pipe)
 * 4. Process commands in a job
 *   - built-in commands
 *   - non-built-in commands
 */
int
main(int ac, char *av[])
{

    /* Process command-line arguments. See getopt(3) */
    int opt;	
    while ((opt = getopt(ac, av, "h")) > 0) {
        switch (opt) {
        case 'h':
            usage(av[0]);
            break;
        }
    }

	/* 1. Initialize */
    list_init(&job_list); 
    termstate_init();
    signal_set_handler(SIGCHLD, sigchld_handler); 
	signal(SIGINT, sigint_handler); 
	signal(SIGTSTP, sigtstp_handler);

    /* 2. Read/Eval command line input */
    for (;;) {

		int commands_order, commands_size, posix; 	
		bool bin_cmd_path;	
		job_jid_fg = -1;
        num_of_command++;

		/* posix related */
		posix_spawnattr_t attribute;
		posix_spawnattr_t *attribute_point;
		posix_spawn_file_actions_t file_actions;
		posix_spawn_file_actions_t *file_actionsp;
		const int PIPE_READ = 0;  
		const int PIPE_WRITE = 1; 
		extern char **environ;
		pid_t child_pid;
		int pgroup_id = 0;
		char *cmd_path; 
		char **argvp;    			

        /* process prompt */
		char * prompt;
		if (!change_prompt) {
			prompt = isatty(0) ? build_prompt() : NULL;
		}
		else {
			if (isatty(0)) {
				prompt = malloc(300);
				build_changed_prompt(prompt, 300);
			}
			else {
				prompt = NULL;
			}
		}

		/* input command line */
        char * cmdline = readline(prompt);
        if (cmdline == NULL)  /* User typed EOF */
            break; 
		free (prompt);

		/* parse command line into pipes */
		struct ast_command_line * cline = ast_parse_command_line(cmdline);
		if (cline == NULL) { 	/* Error in command line */
            continue; 
		}
		free (cmdline);

	    /* User hit enter */
        if (list_empty(&cline->pipes)) {   
	    	sigchld_handler(SIGCHLD, NULL, NULL);
	    	refresh_jobs(); 
            ast_command_line_free(cline);
            continue;
        }

		/* 3. Process job (pipe) */
		for (struct list_elem *e = list_begin(&cline->pipes); e != list_end(&cline->pipes); e = list_next(e)) {

			/* input a pipe and add to job*/
			struct ast_pipeline *pipe = list_entry(e, struct ast_pipeline, elem);
			struct job *job = add_job(pipe);

			/* check BACKGROUND or FOREGROUND */
			job->pgid = (pid_t) - 1;
			if(pipe->bg_job) {
				job->status = BACKGROUND;
			}
			else {
				job->status = FOREGROUND;
				job_jid_fg = job->jid; 
			}

			commands_size = list_size(&pipe->commands);
			int fd[commands_size][2]; /* pipe file descriptor */

			/* 4. Process commands in a job */
			commands_order = 1;
			
			for (struct list_elem *e = list_begin(&pipe->commands);
					e != list_end(&pipe->commands);
					e = list_next(e)) 
			{
				struct ast_command *cmd = list_entry(e, struct ast_command, elem);
				argvp = cmd->argv;
				cmd_path = argvp[0];
				bin_cmd_path = false;

				/* built-in command processing */
				for (int i = 0; i < 8; i++) { /* no of built-in commands = 8 */
					char *bin_cmd = bin_cmds[i];
					if (strcmp(bin_cmd, cmd_path) == 0) {
						bin_cmd_path = true;
						break;
					}
				}
				
				if (bin_cmd_path) {
					signal_unblock(SIGCHLD);
					bin_cmd_main(argvp);				
					continue;
				}

				/* non built-in command processing */

				/* POSIX file action setting */
				int dup_sts = cmd->dup_stderr_to_stdout;
				file_actionsp = NULL;
				posix = posix_spawn_file_actions_init(&file_actions);				

				/* redirection */
				if (pipe->iored_input != NULL) {
					posix= posix_spawn_file_actions_addopen(&file_actions, STDIN_FILENO, pipe->iored_input, O_RDONLY, 0 /* mode */);
				}

				if ((pipe->iored_output != NULL) && (commands_order == commands_size)) {

					if (pipe->append_to_output==1) {
						posix= posix_spawn_file_actions_addopen(&file_actions, STDOUT_FILENO, pipe->iored_output, O_RDWR | O_CREAT | O_APPEND, 0777);
					}
					else if (dup_sts==1) {
						posix= posix_spawn_file_actions_addopen(&file_actions, STDOUT_FILENO, pipe->iored_output, O_RDWR | O_CREAT | O_APPEND, 0777);	
						posix= posix_spawn_file_actions_addopen(&file_actions, STDERR_FILENO, pipe->iored_output, O_RDWR | O_CREAT | O_APPEND, 0777);
					}
					else {
						posix= posix_spawn_file_actions_addopen(&file_actions, STDOUT_FILENO, pipe->iored_output, O_RDWR | O_CREAT | O_TRUNC, 0777);
					}
				}

				/* pipe connection with fds */
				if (commands_size > 1) {
					if (commands_order == 1){
						if (pipe2(fd[commands_order], O_CLOEXEC) == -1)
							perror("pipe"), exit(1);
						posix= posix_spawn_file_actions_adddup2(&file_actions,fd[commands_order][PIPE_WRITE], STDOUT_FILENO);
						if (dup_sts==1) {							
							posix= posix_spawn_file_actions_adddup2(&file_actions,fd[commands_order][PIPE_WRITE], STDERR_FILENO);
							dup_sts = 0;
						}
					}

					if ((commands_order > 1) && (commands_order != commands_size)) {
						if (pipe2(fd[commands_order], O_CLOEXEC) == -1)
							perror("pipe"), exit(1);
						posix= posix_spawn_file_actions_adddup2(&file_actions,fd[commands_order-1][PIPE_READ], STDIN_FILENO);
						posix= posix_spawn_file_actions_adddup2(&file_actions,fd[commands_order][PIPE_WRITE], STDOUT_FILENO);
						if (dup_sts==1) {
							posix= posix_spawn_file_actions_adddup2(&file_actions,fd[commands_order][PIPE_WRITE], STDERR_FILENO);
						}
					}

					if (commands_order == commands_size) {
						posix= posix_spawn_file_actions_adddup2(&file_actions,fd[commands_order-1][PIPE_READ], STDIN_FILENO);
					}
				}

				file_actionsp = &file_actions;

				/* POSIX attribute setting */
				posix_spawnattr_init (&attribute);
				posix_spawnattr_setflags (&attribute, POSIX_SPAWN_SETPGROUP);

				if (commands_order != 1) {
					posix_spawnattr_setpgroup (&attribute, pgroup_id);
				}
				else {
					posix_spawnattr_setpgroup (&attribute, 0);
				}
					
				attribute_point = &attribute;

				/* POSIX execution */
				posix= posix_spawnp(&child_pid, cmd_path, file_actionsp, attribute_point, argvp, environ);
			
				/* POSIX error case */
           		if (posix != 0){ 
					printf("%s\n", strerror(posix));
					continue;
				}

				/* POSIX is OK */
				pid2job[(int)child_pid % MAXPIDS] = job;
				job->num_processes_alive++;
				
				if (commands_order==1) {
					job->pgid = child_pid;
					pgroup_id = child_pid;
				}
				
				commands_order++;			 	
			} /* 4. Process commands in a job is finished */

			for (int i=1; i<commands_size; i++) {
				close(fd[i][0]);
				close(fd[i][1]);
			}

			if (pipe->bg_job) {
				printf("[%i] %i\n", job->jid, job->pgid);
			}

			if (job->num_processes_alive == 0) { 
				refresh_job(job);
			}
			else if (!(pipe->bg_job)) { /* execute foreground job */
				tcsetpgrp(1, pgroup_id);
				wait_for_job(job); 
				termstate_give_terminal_back_to_shell();
			}
		} /* 3. Process job (pipe) is finished */
		
		refresh_jobs();
		free(cline);
	} /* 2. Read/Eval command line input is finished */

    return 0;
} /* main() of cush is finished */


/* Additional Functions */

/* Built-in command "cd (change directory)" */
static void 
bin_cmd_cd(char *arg)
{
	if (arg == NULL) {
		printf("cd error\n");
	}
	else {
		int error = chdir(arg);
		if (error != 0) {
			printf("No such file or directory\n");
		}
	}
}

/* Built-in command "custom_prompt (user custom prompt)" */
static void 
bin_cmd_custom_prompt(char *arg)
{
	if (change_prompt) {
		free(myPrompt);
	}
	
	if (arg == NULL) {
		change_prompt = false;
	}
	else {
		myPrompt = strdup(arg);
		change_prompt = true;
	}
}

/* Create a custom prompt */
static void
build_changed_prompt(char *prompt_received, int size)
{
	char *cpp = myPrompt;
	
	int sizeof_prompt = sizeof(char) * size;
	
	char *current_prompt = prompt_received;
	char *end = prompt_received + sizeof_prompt;
	
	char *addTo = NULL;
	int addToInt = 0;
	
	for (int i = 0; i < strlen(cpp); i++) {
		
		if (cpp[i] != '\\') {
			if (current_prompt + 1 < end) {
				current_prompt[0] = cpp[i];
				current_prompt += sizeof(char);
				continue;
			}
			else {
				printf("my prompt error1\n");
				assert(false);
			}
		}
		
		else {
			i++;
			char data[100];
			if (cpp[i] == 'd') {
				time_t t = time(NULL);
				struct tm *time = localtime(&t);
				int length_date = strftime(data, sizeof(data), "%x", time);
				assert(length_date != 0);
				
				addTo = data;
			}
			else if (cpp[i] == 't') {
				time_t t = time(NULL);
				struct tm *time = localtime(&t);
				int length_time = strftime(data, sizeof(data), "%X", time);
				assert(length_time != 0);
				
				addTo = data;
			}
			else if (cpp[i] == 'h') {
				assert(gethostname(data, sizeof(data)) == 0);
				
				addTo = data;
			}
			else if (cpp[i] == 'j') {
				int job_num = 0;
				for (struct list_elem * e = list_begin (&job_list); 
					e != list_end (&job_list); 
					e = list_next (e)) {
					job_num++;
				}
				
				addTo = NULL;
				addToInt = job_num;
			}
			else if (cpp[i] == '#') {
				addTo = NULL;
				addToInt = num_of_command;
			}
			else if (cpp[i] == 'u') {
				char *user = getenv("USER");
				assert(user != NULL);
				
				addTo = user;
			}
			else if (cpp[i] == 'w') {
				assert(getcwd(data, sizeof(data)) != NULL);
				
				addTo = data;
			}
			
			if (current_prompt < end) {
				int inc_prompt;
				if (addTo != NULL) {
					inc_prompt = snprintf(current_prompt, end - current_prompt, "%s", addTo);
				}
				else {
					inc_prompt = snprintf(current_prompt, end - current_prompt, "%i", addToInt);
				}
				if (inc_prompt > 0) {
					current_prompt += inc_prompt;
				}
				else {
					printf("my prompt error2\n");
					assert(false);
				}
			}
		}
	}
}