/* shell.c - the implementation of the shell.
 *
 * by Sam van Kampen, 11874716. */
#include "shell.h"
#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdio_ext.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <sys/utsname.h>
#include <pwd.h>
#include "parser/ast.h"
#include "builtins.h"
#include "front.h"
#include "parser/lex.yy.h"
#include "pidlist.h"

static void run_node_nofork(node_t*);
static void run_command_inner(node_t*);

/* Kills PIDs if they are foreground processes. */
void kill_if_fg(pid_node_t *p)
{
    if (p->foreground) kill(p->pid, SIGTERM);
}

/* Runs kill_if_fg using pl_foreach. */
static void kill_foreground_pids()
{
    pl_foreach(kill_if_fg);
}

/* Handles SIGINT and SIGCHLD. */
static void signal_handler(int signum)
{
    pid_t child;
    int status;

    switch (signum)
    {
        case SIGINT:
            // TERM all running foreground PIDs.
            kill_foreground_pids();
            break;
        case SIGCHLD:
            while ((child = waitpid(-1, &status, WNOHANG)) > 0)
            {
                pl_remove_pid(child);
            }
            break;
        default:
            fprintf(stderr, "Received unhandled signal %s\n",
                    strsignal(signum));
            break;
    }
}

/* Wrapper around sigaction that prints an error and exits when it fails. */
static void register_sigaction(int sig, struct sigaction *sa)
{
    if (sigaction(sig, sa, NULL))
    {
        perror("sigaction");
        exit(EXIT_FAILURE);
    }
}

/* Wrapper around setsigmask, just like register_sigaction. */
static void setsigmask(int flag, sigset_t* new, sigset_t* old)
{
    int s = sigprocmask(flag, new, old);
    if (s) perror("sigprocmask");
}

void initialize(void)
{
    /* Initialize the PID list. Destroy it on exit. */
    pl_init();
    atexit(pl_destroy);

    /* Register SIGINT and SIGCHLD handlers */
    sigset_t empty;
    sigemptyset(&empty);

    struct sigaction sa;
    sa.sa_handler = signal_handler;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = SA_RESTART;

    register_sigaction(SIGINT, &sa);
    register_sigaction(SIGCHLD, &sa);

    setsigmask(SIG_SETMASK, &empty, 0);
}

/* Fork, switch the dealloc_on_pop_all flag, and add the child process
 * to the PID list. */
static pid_t vork()
{
    pid_t child_pid    = fork();
    sigset_t s = block_sigchld();

    if (child_pid)
        pl_add_pid(child_pid)->foreground = true;

    reset_sigmask(s);
    return child_pid;
}

/* Wait on a child process, setting errno to its exit status. */
static int wait_on_child(pid_t child)
{
    int wstatus;
    waitpid(child, &wstatus, WUNTRACED);
    pl_remove_pid(child);
    errno = WEXITSTATUS(wstatus);
    return 0;
}

/* Run a command without forking. */
static void run_command_nofork(node_t *node)
{
    char *program = node->command.program;
    char **argv   = node->command.argv;

    /* The builtin call may be in code twice, but it only gets called once!
     * The fact that we need different behavior for a fork basically forces
     * this. */
    int (*builtin_fn)(node_t *);
    if ((builtin_fn = get_builtin(program)))
    {
        errno = 0;
        errno = builtin_fn(node);
    }
    else
    {
        execvp(program, argv);
        perror(program);
        exit(0);
    }
}

/* Run a command. */
static void run_node_command(node_t *node)
{
    char *program = node->command.program;

    int (*builtin_fn)(node_t *);
    if ((builtin_fn = get_builtin(program)))
    {
        errno = 0;
        errno = builtin_fn(node);
    }
    else
    {
        pid_t child_pid = vork();

        if (!child_pid)
            run_command_nofork(node);
        else
            wait_on_child(child_pid);
    }
}

/* Set up a REDIRECT_DUP. */
static void redirect_dup(node_t *node)
{
    int fd = node->redirect.fd;
    int fd2 = node->redirect.fd2;
    if (fd == -1)
    {
        dup2(fd2, STDOUT_FILENO);
        dup2(fd2, STDERR_FILENO);
    }
    else
        dup2(fd2, fd);
}

/* Set up a redirect that is NOT a REDIRECT_DUP. */
static int redirect_otherwise(node_t *node)
{
    int flags = 0;
    int target_fd = -1;
    int fd = node->redirect.fd;
    switch (node->redirect.mode)
    {
        case REDIRECT_INPUT:
            flags |= O_RDONLY; break;
        case REDIRECT_OUTPUT:
            flags |= O_WRONLY | O_TRUNC | O_CREAT; break;
        case REDIRECT_APPEND:
            flags |= O_WRONLY | O_APPEND | O_CREAT; break;
        default: break;
    };

    target_fd = open(node->redirect.target, flags, 0644);
    if (target_fd == -1)
    {
        perror("open");
        return -1;
    }

    if (fd == -1)
    {
        dup2(target_fd, STDOUT_FILENO);
        dup2(target_fd, STDERR_FILENO);
    }
    else
    {
        dup2(target_fd, fd);
    }
    return target_fd;
}

/* Run a redirect, without forking. */
static void run_redirect_nofork(node_t *node)
{
    int target_fd = -1;
    if (node->redirect.mode == REDIRECT_DUP)
        redirect_dup(node);
    else
        target_fd = redirect_otherwise(node);

    run_node_nofork(node->redirect.child);

    if (target_fd != -1)
        close(target_fd);
}

/* Run a node, without' forking.
 *
 *  ') may contain traces of forking. */
static void run_node_nofork(node_t *node)
{
    pid_t child;
    switch (node->type)
    {
        case NODE_COMMAND:
            run_command_nofork(node);
            break;
        case NODE_REDIRECT:
            run_redirect_nofork(node);
            break;
        case NODE_SUBSHELL:
            /* this still requires a fork, even in _nofork. */
            child = vork();
            if (!child)
                run_node_nofork(node->subshell.child);
            else
                wait_on_child(child);
            break;
        case NODE_SEQUENCE:
            /* We have to fork once, since calling run_node_nofork leads to
             * exit. */
            run_command_inner(node->sequence.first);
            run_node_nofork(node->sequence.second);
            break;
        default:
            break;
    }
}

/* Connect the pipes to each of the proper file descriptors. */
static void connect_pipes(int command_idx, int wpipe_idx, int rpipe_idx, int n_pipes,
                          int *pipes)
{
    /* If we're not on the last command, connect
     * the command's stdout to a write pipe. */
    if (command_idx < n_pipes)
    {
        if (dup2(pipes[wpipe_idx], 1) < 0)
        {
            perror("dup2");
            exit(EXIT_FAILURE);
        }
    }

    /* If this isn't the first command,
     * connect a read pipe to stdin. */
    if (rpipe_idx != 0)
    {
        if (dup2(pipes[rpipe_idx - 2], 0) < 0)
        {
            perror("dup2");
            exit(EXIT_FAILURE);
        }
    }

    for (int i = 0; i < 2 * n_pipes; ++i)
    {
        close(pipes[i]);
    }
}

/* Creates n_pipes pipes in an out-array. */
static void create_pipes(int n_pipes, int *out_pipes)
{
    for (int i = 0; i < n_pipes; ++i)
    {
        if (pipe(out_pipes + i * 2) < 0)
        {
            perror("Unable to create pipe");
            return;
        }
    }
}

/* Create the subprocesses for each of the pipe parts. */
static void create_pipe_subprocesses(node_t *node, int n_pipes, int *pipes)
{
    pid_t child_pid;
    int rpipe_idx = 0;
    int wpipe_idx = 1;

    for (int i = 0; i <= n_pipes; i++, rpipe_idx += 2, wpipe_idx += 2)
    {
        node_t *child_node = node->pipe.parts[i];

        child_pid = vork();

        if (!child_pid)
        {
            connect_pipes(i, wpipe_idx, rpipe_idx, n_pipes, pipes);
            run_node_nofork(child_node);
            exit(0);
        }
        else if (child_pid < 0)
        {
            perror("fork");
            exit(EXIT_FAILURE);
        }
    }
}

/* Run a node pipe. */
static void run_node_pipe(node_t *node)
{
    int status;
    int n_commands = node->pipe.n_parts;
    int n_pipes    = (n_commands - 1);

    /* pipes[0] refers to the read end,
     * pipes[1] refers to the write end. */
    int pipes[2 * n_pipes];
    create_pipes(n_pipes, pipes);

    create_pipe_subprocesses(node, n_pipes, pipes);

    for (int i = 0; i < 2 * n_pipes; ++i)
    {
        close(pipes[i]);
    }

    pid_t child;
    while ((child = waitpid(-1, &status, 0)) > 0)
    {
        if (WIFEXITED(status) ||
            (WIFSIGNALED(status) && WTERMSIG(status) == SIGPIPE))
        {
            pl_remove_pid(child);
        }
    }
}

/* Run a redirect node. */
static void run_node_redirect(node_t *node)
{
    pid_t child_pid = vork();

    if (!child_pid)
        run_redirect_nofork(node);
    else
        wait_on_child(child_pid);
}

/* Run a node in a subshell. */
static void run_node_subshell(node_t *node)
{
    pid_t child_pid = vork();

    if (!child_pid)
    {
        run_command_inner(node->subshell.child);
        exit(0);
    }
    else
        wait_on_child(child_pid);
}

/* Runs a command without registering it for cleanup;
 * used because sequences get cleaned up entirely by
 * free_tree, so we can't register different parts of them. */
static void run_command_inner(node_t *node)
{
    pid_t child;
    switch (node->type)
    {
        case NODE_COMMAND:
            run_node_command(node);
            break;
        case NODE_PIPE:
            run_node_pipe(node);
            break;
        case NODE_SEQUENCE:
            run_command_inner(node->sequence.first);
            run_command_inner(node->sequence.second);
            break;
        case NODE_REDIRECT:
            run_node_redirect(node);
            break;
        case NODE_SUBSHELL:
            run_node_subshell(node);
            break;
        case NODE_DETACH:
            child = vork();
            if (!child)
            {
                run_command_inner(node->detach.child);
                exit(0);
            }
        default:
            break;
    };
}

/* Runs a command. */
void run_command(node_t *node)
{
    run_command_inner(node);
}