#include <stdio.h>
#include <errno.h>
#include <assert.h>
#include <string.h>
#include <stdlib.h>
#include <stdbool.h>

#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <time.h>

// For reference
// ~cpen212/Public/lab5/crash-ref

#define MAXJOBS 32
#define MAXLINE 1024

char **environ;

struct job {
    pid_t pid;
    char cmd[MAXLINE];
    bool terminated;
    bool foreground;
    bool suspended;
};

struct job job_list[MAXJOBS];
int num_jobs = 0;

void int_to_str(int num, char* str);
char* string_combine(char* jid_str, char* pid_str, char* status, char* cmd);

void handle_sigchld(int sig) {
    int status, pid;
    while ((pid = waitpid(-1, &status, WNOHANG | WUNTRACED | WCONTINUED)) > 0) {
        for (int i = 0; i < num_jobs; i++) {
            if (job_list[i].pid == pid) {
                // Convert the job ID and PID to strings using a helper function
                char jid_str[16], pid_str[16];
                int_to_str(i+1, jid_str);
                int_to_str(job_list[i].pid, pid_str);

                if (WIFEXITED(status)) {
                    char* msg = string_combine(jid_str, pid_str, "finished", job_list[i].cmd);
                    int len = strlen(msg);
                    write(STDOUT_FILENO, msg, len);
                    job_list[i].terminated = true;
                } else if (WIFSIGNALED(status)) {
                    if (WCOREDUMP(status)) {
                        char* msg = string_combine(jid_str, pid_str, "killed (core dumped)", job_list[i].cmd);
                        int len = strlen(msg);
                        write(STDOUT_FILENO, msg, len);
                    }
                    else {
                        char* msg = string_combine(jid_str, pid_str, "killed", job_list[i].cmd);
                        int len = strlen(msg);
                        write(STDOUT_FILENO, msg, len);
                    }
                    job_list[i].terminated = true;
                } else if (WIFSTOPPED(status)) {
                    if (job_list[i].terminated == true) {
                        return;
                    }
                    char* msg = string_combine(jid_str, pid_str, "suspended", job_list[i].cmd);
                    int len = strlen(msg);
                    write(STDOUT_FILENO, msg, len);
                    job_list[i].suspended = true;
                } else if (WIFCONTINUED(status)) {
                    char* msg = string_combine(jid_str, pid_str, "running", job_list[i].cmd);
                    int len = strlen(msg);
                    write(STDOUT_FILENO, msg, len);
                    job_list[i].suspended = false;
                }
            }
        }
    }
}

void handle_sigtstp(int sig) {
    pid_t fg_pgid = tcgetpgrp(STDIN_FILENO);

    if (sig == SIGTSTP) {
        for (int i = 0; i < num_jobs; i++) {
            if (job_list[i].suspended == false && job_list[i].terminated == false) {
                char jid_str[16], pid_str[16];
                int_to_str(i+1, jid_str);
                int_to_str(job_list[i].pid, pid_str);

                char* msg = string_combine(jid_str, pid_str, "suspended", job_list[i].cmd);
                int len = strlen(msg);
                write(STDOUT_FILENO, msg, len);
                job_list[i].foreground = false;
                job_list[i].suspended = true;
            }
        }
    }
}

void handle_sigint(int sig) {
    pid_t fg_pgid = tcgetpgrp(STDIN_FILENO);

    if (sig == SIGINT) {
        for (int i = 0; i < num_jobs; i++) {
            if (job_list[i].foreground == true && job_list[i].terminated == false) {
                char jid_str[16], pid_str[16];
                int_to_str(i+1, jid_str);
                int_to_str(job_list[i].pid, pid_str);

                char* msg = string_combine(jid_str, pid_str, "killed", job_list[i].cmd);
                int len = strlen(msg);
                write(STDOUT_FILENO, msg, len);
                job_list[i].foreground = false;
            }
        }
    }
}

void handle_sigquit(int sig) {
    pid_t fg_pgid = tcgetpgrp(STDIN_FILENO);

    if (sig == SIGQUIT) {
        for (int i = 0; i < num_jobs; i++) {
            if (job_list[i].foreground == true && job_list[i].terminated == false) {
                char jid_str[16], pid_str[16];
                int_to_str(i+1, jid_str);
                int_to_str(job_list[i].pid, pid_str);

                char* msg = string_combine(jid_str, pid_str, "killed (core dumped)", job_list[i].cmd);
                int len = strlen(msg);
                write(STDOUT_FILENO, msg, len);
                job_list[i].foreground = false;
            }
        }
    }
}

void install_signal_handlers() {
    struct sigaction sa_chld, sa_tstp, sa_int, sa_quit;

    sa_chld.sa_handler = handle_sigchld;
    sa_tstp.sa_handler = handle_sigtstp;
    sa_int.sa_handler = handle_sigint;
    sa_quit.sa_handler = handle_sigquit;

    sigemptyset(&sa_chld.sa_mask);
    sigemptyset(&sa_tstp.sa_mask);
    sigemptyset(&sa_int.sa_mask);
    sigemptyset(&sa_quit.sa_mask);

    sa_chld.sa_flags = SA_RESTART;
    sa_tstp.sa_flags = SA_RESTART;
    sa_int.sa_flags = SA_RESTART;
    sa_quit.sa_flags = SA_RESTART;

    sigaction(SIGCHLD, &sa_chld, NULL);
    sigaction(SIGTSTP, &sa_tstp, NULL);
    sigaction(SIGINT, &sa_int, NULL);
    sigaction(SIGQUIT, &sa_quit, NULL);
}

void spawn(const char **toks, bool bg) { // bg is true iff command ended with &

    pid_t pid = fork(); // create a new process
    int status;
    
    if (pid < 0) {
        fprintf(stderr, "ERROR: failed to fork\n");
        exit(1);
    }
    else if (pid == 0) { // child process
        execvp(toks[0], (char * const *)toks); // execute the command
        fprintf(stderr, "ERROR: cannot run %s\n", toks[0]); // executes only if execvp fails
        exit(1);
    }
    else { // parent process
        if (num_jobs == MAXJOBS) { // 
            fprintf(stderr, "ERROR: too many jobs\n");
        }
        else {
            job_list[num_jobs].pid = pid;
            strncpy(job_list[num_jobs].cmd, *toks, MAXLINE);
            num_jobs++;
            if (bg) {
                printf("[%d] (%d)  running  %s\n", num_jobs, job_list[num_jobs - 1].pid, job_list[num_jobs - 1].cmd);
            }
        }
        if (!bg) { // waits if running in foreground
            job_list[num_jobs - 1].foreground = true;
            waitpid(pid, &status, 0);
            job_list[num_jobs - 1].terminated = true;
        }
    }
}

void cmd_jobs(const char **toks) {
    int term_counter = 0;
    if (toks[1] != NULL) {
        fprintf(stderr, "ERROR: jobs takes no arguments\n");
    }
    else {
        for (int i = 0; i < num_jobs; i++) {
            if (job_list[i].terminated) {
                continue;
            }
            if (job_list[i].suspended) {
                printf("[%d] (%d)  suspended  %s\n", i + 1, job_list[i].pid, job_list[i].cmd);
            }
            else {
                printf("[%d] (%d)  running  %s\n", i + 1, job_list[i].pid, job_list[i].cmd);
            }
        }
    }
}

void cmd_fg(const char **toks) {
    // Check if there is currently no foreground job running
    pid_t fg_pid = tcgetpgrp(STDIN_FILENO);
    if (fg_pid != getpgrp()) {
        fprintf(stderr, "ERROR: There is already a foreground job running\n");
        return;
    }

    if (toks[1] == NULL || toks[2] != NULL) {
        fprintf(stderr, "ERROR: fg takes exactly one argument\n");
    }
    else if (toks[1][0] == '%') { // by job
        char *endptr;
        int job_id = strtol(&toks[1][1], &endptr, 10);
        if (*endptr != '\0' || job_id < 0) {
            fprintf(stderr, "ERROR: bad argument for fg: %s\n", toks[1]);
        }
        else if (job_id > 0 && job_id < MAXJOBS && job_list[job_id - 1].pid > 0) {
            tcsetpgrp(STDIN_FILENO, job_list[job_id - 1].pid);

            // Wait for the job to finish using waitpid
            int status;
            waitpid(job_list[job_id - 1].pid, &status, WUNTRACED);

            // Restore the shell as the foreground process group using tcsetpgrp
            tcsetpgrp(STDIN_FILENO, getpgrp());

            // Update the job status
            if (WIFSTOPPED(status)) {
                printf("[%d] (%d)  suspended  %s\n", job_id, job_list[job_id - 1].pid, job_list[job_id - 1].cmd);
                job_list[job_id - 1].terminated;
            } else {
                job_list[job_id - 1].terminated;
            }
        }
        else {
            fprintf(stderr, "ERROR: no job %s\n", toks[1]);
        }
    }
    else if (atoi(toks[1]) > 0) { // by PID
        int temp_pid = atoi(toks[1]);
        int no_fg = 1;
        for (int i = 0; i < num_jobs; i++) {
            if (job_list[i].pid == temp_pid) {
                tcsetpgrp(STDIN_FILENO, job_list[i].pid);

                // Wait for the job to finish using waitpid
                int status;
                waitpid(job_list[i].pid, &status, WUNTRACED);

                // Restore the shell as the foreground process group using tcsetpgrp
                tcsetpgrp(STDIN_FILENO, getpgrp());

                // Update the job status
                if (WIFSTOPPED(status)) {
                    printf("[%d] (%d)  suspended  %s\n", i + 1, job_list[i].pid, job_list[i].cmd);
                    job_list[i].terminated;
                } else {
                    job_list[i].terminated;
                }
                no_fg = 0;
            }
        }
        if (no_fg) {
            fprintf(stderr, "ERROR: no PID %s\n", toks[1]);
        }
    }
    else {
        fprintf(stderr, "ERROR: bad argument for fg: %s\n", toks[1]);
    }
}

void cmd_bg(const char **toks) {
    if (toks[1] == NULL || toks[2] != NULL) {
        fprintf(stderr, "ERROR: bg takes exactly one argument\n");
    }
    else if (toks[1][0] == '%') { // by job
        char *endptr;
        int job_id = strtol(&toks[1][1], &endptr, 10);
        if (*endptr != '\0' || job_id < 0) {
            fprintf(stderr, "ERROR: bad argument for bg: %s\n", toks[1]);
        }
        else if (job_list[job_id - 1].suspended == false) {
            if (job_list[job_id - 1].terminated == true) {
                fprintf(stderr, "ERROR: no job %s\n", toks[1]);
            }
            else {
                return;
            }
        }
        else if (job_id > 0 && job_id < MAXJOBS && job_list[job_id - 1].pid > 0) {
            kill(job_list[job_id - 1].pid, SIGCONT);
        }
        else {
            fprintf(stderr, "ERROR: no job %s\n", toks[1]);
        }
    }
    else if (atoi(toks[1]) > 0) { // by PID
        int temp_pid = atoi(toks[1]);
        int no_bg = 1;
        for (int i = 0; i < num_jobs; i++) {
            if (job_list[i].terminated == true) {
                fprintf(stderr, "ERROR: no PID %s\n", toks[1]);
            }
            else if (job_list[i].suspended == true) {
                kill(job_list[i].pid, SIGCONT);
            }
            no_bg = 0;
        }
        if (no_bg) {
            fprintf(stderr, "ERROR: no PID %s\n", toks[1]);
        }
    }
    else {
        fprintf(stderr, "ERROR: bad argument for bg: %s\n", toks[1]);
    }
}

void cmd_nuke(const char **toks) {
    if (toks[2] != NULL) {
        fprintf(stderr, "ERROR: nuke takes exactly one argument\n");
    }
    else if (toks[1] == NULL) { // kill all
        for (int i = 0; i < num_jobs; i++) {
            kill(job_list[i].pid, SIGKILL);
        }
    }
    else if (toks[1][0] == '%') { // kill by job
        char *endptr;
        int job_id = strtol(&toks[1][1], &endptr, 10);
        if (*endptr != '\0' || job_id < 0) {
            fprintf(stderr, "ERROR: bad argument for nuke: %s\n", toks[1]);
        }
        else if (job_id > 0 && job_id < MAXJOBS && job_list[job_id - 1].pid > 0) {
            kill(job_list[job_id - 1].pid, SIGKILL);
        }
        else {
            fprintf(stderr, "ERROR: no job %s\n", toks[1]);
        }
    }
    else if (atoi(toks[1]) > 0) { // kill by PID
        int temp_pid = atoi(toks[1]);
        int no_kill = 1;
        for (int i = 0; i < num_jobs; i++) {
            if (job_list[i].pid == temp_pid) {
                kill(job_list[i].pid, SIGKILL);
                no_kill = 0;
            }
        }
        if (no_kill) {
            fprintf(stderr, "ERROR: no PID %s\n", toks[1]);
        }
    }
    else {
        fprintf(stderr, "ERROR: bad argument for nuke: %s\n", toks[1]);
    }
}

void cmd_quit(const char **toks) {
    if (toks[1] != NULL) {
        fprintf(stderr, "ERROR: quit takes no arguments\n");
    }
    else {
        exit(0);
    }
}

void eval(const char **toks, bool bg) { // bg is true iff command ended with &
    assert(toks);
    if (*toks == NULL) return;
    if (strcmp(toks[0], "quit") == 0) {
        cmd_quit(toks);
    }
    else if (strcmp(toks[0], "jobs") == 0) {
        cmd_jobs(toks);
    }
    else if (strcmp(toks[0], "nuke") == 0) {
        cmd_nuke(toks);
    }
    else if (strcmp(toks[0], "fg") == 0) {
        cmd_fg(toks);
    }
    else if (strcmp(toks[0], "bg") == 0) {
        cmd_bg(toks);
    }
    else {
        spawn(toks, bg);
    }
}

// you don't need to touch this unless you want to add debugging
void parse_and_eval(char *s) {
    assert(s);
    const char *toks[MAXLINE + 1];
    
    while (*s != '\0') {
        bool end = false;
        bool bg = false;
        int t = 0;

        while (*s != '\0' && !end) {
            while (*s == '\n' || *s == '\t' || *s == ' ') ++s;
            if (*s != ';' && *s != '&' && *s != '\0') toks[t++] = s;
            while (strchr("&;\n\t ", *s) == NULL) ++s;
            switch (*s) {
            case '&':
                bg = true;
                end = true;
                break;
            case ';':
                end = true;
                break;
            }
            if (*s) *s++ = '\0';
        }
        toks[t] = NULL;
        eval(toks, bg);
    }
}

// you don't need to touch this unless you want to add debugging
void prompt() {
    printf("crash> ");
    fflush(stdout);
}

// you don't need to touch this unless you want to add debugging
int repl() {
    char *buf = NULL;
    size_t len = 0;
    while (prompt(), getline(&buf, &len, stdin) != -1) {
        parse_and_eval(buf);
    }

    if (buf != NULL) free(buf);
    if (ferror(stdin)) {
        perror("ERROR");
        return 1;
    }
    return 0;
}

// you don't need to touch this unless you want to add debugging options
int main(int argc, char **argv) {
    install_signal_handlers();
    return repl();
}

void int_to_str(int num, char* str) {
    int i = 0;

    // Build string in reverse order
    while (num > 0) {
        str[i++] = (num % 10) + '0';
        num /= 10;
    }
    str[i] = '\0';

    // Reverse string
    int len = strlen(str);
    for (int j = 0; j < len / 2; j++) {
        char temp = str[j];
        str[j] = str[len - j - 1];
        str[len - j - 1] = temp;
    }
}

char* string_combine(char* jid_str, char* pid_str, char* status, char* cmd) {
    char* strings[] = {"[", jid_str, "] (", pid_str, ")  ", status, "  ", cmd, "\n"};
    int total_len = 0;
    for(int i = 0; i < 9; i++) {
        total_len += strlen(strings[i]);
    }
    char* result = (char*) malloc(total_len + 1); // +1 for null terminator
    result[0] = '\0';
    for(int i = 0; i < 9; i++) {
        strcat(result, strings[i]);
    }
    return result;
}