/*
 *  Naam      : R. Wacanno
 *  UvAnetID  : 11741163
 *  Studie    : BSc Informatica
 *
 *  shell.c
 *  -Contains functions to process a given node.
 */

#include "arena.h"
#include "ast.h"
#include "front.h"
#include <stdbool.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <signal.h>

/* Parses a given tree and calls type specific helper functions. */
static int parse_tree(node_t *node, bool c_fork, bool c_wait);

void my_free_tree(void *pt)
{
    free_tree((node_t *)pt);
}

/* Makes sure SIGTERM is ingnored in the parent process. */
void initialize(void)
{
    signal(SIGINT, SIG_IGN);
}

/* Forks and executes a given function which takes a given node as argument. */
static int fork_func(int (*f)(node_t *), node_t *arg_node, bool c_wait)
{
    int status;
    pid_t c_pid = fork();

    if (c_pid < 0)
    {
        perror("Error whilst forking");
        return 1;
    }
    else if (c_pid == 0)
    {
        (*f)(arg_node);
        exit(0);
    }
    else if (c_wait)
    {
        waitpid(c_pid, &status, 0);
    }

    return 0;
}

/* Forks and executes a given function which takes a given node as argument. */
static int fork_extr(int (*f)(node_t *, bool, bool),
                     node_t *arg_node,
                     bool c_fork,
                     bool c_wait)
{
    int status;
    pid_t c_pid = fork();

    if (c_pid < 0)
    {
        perror("Error whilst forking");
        return 1;
    }
    else if (c_pid == 0)
    {
        (*f)(arg_node, c_fork, true);
        exit(0);
    }
    else if (c_wait)
    {
        waitpid(c_pid, &status, 0);
    }

    return 0;
}

/* Exits with a given exit code. Exits with 0 if no exit code is given. */
static int cmd_exit(node_t *node)
{
    if (node->command.argc == 1)
    {
        exit(EXIT_SUCCESS);
    }
    else if (node->command.argc == 2)
    {
        int code = atoi(node->command.argv[1]);
        exit(code);
    }
    else
    {
        perror("Too many arguments");
        return 1;
    }
}

/* Executes a given path with given arguments. */
static int cmd_exec(node_t *node)
{
    char *program = node->command.program;
    char **argv = node->command.argv;
    signal(SIGINT, SIG_DFL);
    execvp(program, argv);
    perror("Error");
    return 1;
}

/* Assigns a given value to a given enviroment variable. */
static int cmd_set(node_t *node)
{
    char *arg = node->command.argv[1];
    char *name = strtok(arg, "=");
    char *value = strtok(NULL, "=");
    return setenv(name, value, 1);
}

/* Parses a given command node. */
static int parse_command(node_t *node, bool c_fork, bool c_wait)
{
    if (!strcmp(node->command.program, "cd"))
    {
        return chdir(node->command.argv[1]);
    }
    if (!strcmp(node->command.program, "exit"))
    {
        return cmd_exit(node);
    }
    if (!strcmp(node->command.program, "set"))
    {
        return cmd_set(node);
    }
    if (!strcmp(node->command.program, "unset"))
    {
        return unsetenv(node->command.argv[1]);
    }
    else if (c_fork)
    {
        return fork_func(&cmd_exec, node, c_wait);
    }
    else
    {
        cmd_exec(node);
    }
    return 1;
}

/* Parses a given sequence node. */
static int parse_sequence(node_t *node, bool c_fork, bool c_wait)
{
    if (parse_tree(node->sequence.first, c_fork, c_wait))
    {
        return 1;
    }
    if (parse_tree(node->sequence.second, c_fork, c_wait))
    {
        return 1;
    }
    return 0;
}

/* Duplicates the correct pipe filedestriptors in the child proces
   and parses the correct pipe part. */
static int child_pipe(node_t *node, int p_pos, int *fd, int fd_in)
{
    int p_count = node->pipe.n_parts;
    node_t **p_parts = node->pipe.parts;
    node_t *part = p_parts[p_pos];

    dup2(fd_in, 0);
    if (p_pos < (p_count - 1))
    {
        dup2(fd[1], 1);
        close(fd[1]);
    }
    close(fd[0]);
    parse_tree(part, false, true);
    exit(0);
    return 1;
}

/* Parses a given pipe node. */
static int parse_pipe(node_t *node, bool c_wait)
{
    int p_count = node->pipe.n_parts;
    int fd[2];
    int fd_in = 0;
    pid_t c_pid;
    int status;

    for (int i = 0; i < p_count; i++)
    {
        pipe(fd);
        c_pid = fork();
        if (c_pid == 0)
        {
            if (child_pipe(node, i, fd, fd_in))
            {
                perror("Error whilst piping");
                return 1;
            }
        }
        else if (c_pid < 0)
        {
            perror("Error whilst executing pipe (failed fork)");
            return 1;
        }
        close(fd[1]);
        fd_in = fd[0];
    }

    close(fd[0]);
    close(fd[1]);

    if (c_wait)
    {
        while (wait(&status) > 0);
    }
    return 0;
}

/* Duplicates the desired file destriptors and parses the child node */
static int child_red(node_t *node, int fd_o, int fd_n)
{
    if (fd_o == -1)
    {
        dup2(fd_n, 1);
        dup2(fd_n, 2);
    }
    else
    {
        dup2(fd_n, fd_o);
        close(fd_n);
    }
    parse_tree(node->redirect.child, false, true);
    exit(0);
    return 1;
}

/* Returns the correct flags based on the redirect type*/
static int get_flags(node_t *node)
{
    switch(node->redirect.mode)
    {
        case REDIRECT_OUTPUT: return O_WRONLY | O_CREAT | O_TRUNC;
        case REDIRECT_APPEND: return O_WRONLY | O_APPEND | O_CREAT;
        case REDIRECT_DUP: return O_WRONLY | O_CREAT | O_TRUNC;
        case REDIRECT_INPUT: return O_RDONLY;
    }
}

/* Returns the original file destriptor based on the redirect type*/
static int get_o_fd(node_t *node)
{
    switch(node->redirect.mode)
    {
        case REDIRECT_OUTPUT:
        case REDIRECT_APPEND: return 1;
        case REDIRECT_DUP: return -1;
        case REDIRECT_INPUT: return 0;
    }
}

/* Parses a given redirect node */
static int parse_red(node_t *node)
{
    int flags = get_flags(node);
    int fd_org = get_o_fd(node);
    int fd_new;
    int status;

    pid_t c_pid = fork();

    if (c_pid < 0)
    {
        perror("Error whilst parsing redirection (fork failed)");
        return 1;
    }
    else if (c_pid == 0)
    {
        if (node->redirect.mode == REDIRECT_DUP)
        {
            fd_new = node->redirect.fd2;
        }
        else
        {
            fd_new = open(node->redirect.target, flags, 0666);
        }
        child_red(node, fd_org, fd_new);
    }
    else
    {
        waitpid(c_pid, &status, 0);
    }
    return 0;
}

/* Parses a given detach node */
static int parse_det(node_t *node)
{
    return fork_extr(parse_tree, node->detach.child, true, false);
}

/* Parses a given sub shell node */
static int parse_sub(node_t *node)
{
    return fork_extr(parse_tree, node->subshell.child, true, true);
}

/* Parses a given tree and calls type specific helper functions. */
static int parse_tree(node_t *node, bool c_fork, bool c_wait)
{
    switch (node->type)
    {
        case NODE_COMMAND: return parse_command(node, c_fork, c_wait);
        case NODE_SEQUENCE: return parse_sequence(node, c_fork, c_wait);
        case NODE_PIPE: return parse_pipe(node, c_wait);
        case NODE_REDIRECT: return parse_red(node);
        case NODE_DETACH: return parse_det(node);
        case NODE_SUBSHELL: return parse_sub(node);
    }
    return 1;
}

void run_command(node_t *node)
{
    // Create a new memory arena to handle this function in. Please see the
    // hints and `arena.h` on how to use them, as you can very easily create
    // memory leaks if you are not careful.
    arena_push();
    arena_register_mem(node, &my_free_tree);

    parse_tree(node, true, true);

    // Pop the current memory arena, this frees all memory in the current arena.
    arena_pop();
}
