#include <assert.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <signal.h>
#include <sys/wait.h>

#include "syntax.h"
#include "utils.h"
#include "stages.h"

void execute_command_pipeline(Pipeline pipeline, Channel read_cha, Channel write_cha, bool wait_here, int* pgid, ShellState* shstate);


DECLARE_MAYBE(StrPair, MaybeStrPair)

MaybeStrPair parse_assignment(char* token) {
    // TODO: ideally shouldn't mutate maybe
    if (token == NULL) return (MaybeStrPair){.some = false};
    char* eq = strchr(token, '=');
    if (eq == NULL || eq == token) return (MaybeStrPair){.some = false};
    *eq = '\0';
    if (!is_identifier(token)) exit(42);
    return (MaybeStrPair){.some = true, .value = (StrPair){.key = token, .val = eq+1}};
}


void source_script(char* fname, ShellState* shstate) {
    // TODO: add stdin/out (to other builtins too, actually)
    FILE* fp = fopen(fname, "r");        // TODO: checks
    if (fp == 0) exit(1);

    char* line = NULL;
    size_t len = 0;

    while (getline(&line, &len, fp) != -1) {
		line[strcspn(line, "\r\n")] = '\0';
        execute_string_line(line, shstate);
    }
    free(line);
    fclose(fp);

}


bool handle_builtin(ExecutableCommand cmd, ShellState* shstate) {
    MaybeStrPair m_assignment;
    assert(cmd.args[0] != NULL); 
    // TODO: handle id=value assignments
    
    m_assignment = parse_assignment(cmd.args[0]);
    if (m_assignment.some) {
        set_map_pair(&shstate->shell_vars, m_assignment.value);
        return true;
    }

    if (strcmp(cmd.args[0], "exit") == 0) {
        exit(0);
    } else if (strcmp(cmd.args[0], "cd") == 0) {
        assert(cmd.args[1] != NULL);
        chdir(cmd.args[1]);
        return true;
    } else if (strcmp(cmd.args[0], "alias") == 0) {
        // TODO: if no arg1 -> print alias list
        m_assignment = parse_assignment(cmd.args[1]);
        assert(m_assignment.some);
        set_map_pair(&shstate->aliases, m_assignment.value); 
        return true;
    } else if (strcmp(cmd.args[0], "source") == 0) {
        assert(cmd.args[1] != NULL);
        source_script(cmd.args[1], shstate);
        return true;
    } else if (strcmp(cmd.args[0], "export") == 0) {
        m_assignment = parse_assignment(cmd.args[1]);
        assert(m_assignment.some);
        setenv(m_assignment.value.key, m_assignment.value.val, 1);
        return true;
    } else if (strcmp(cmd.args[0], "unalias") == 0) { 
        assert(cmd.args[1] != NULL);
        unset_map_key(&shstate->aliases, cmd.args[1]);
        return true;
    }  else if (strcmp(cmd.args[0], "unset") == 0) {
        assert(cmd.args[1] != NULL);
        unset_map_key(&shstate->shell_vars, cmd.args[1]);
        return true;
    }
    return false;
}

void execute_cmd(ExecutableCommand cmd, int* pgid) {

    pid_t pid = fork();

    if (pid == 0) {

        if (*pgid == -1) {
            setpgid(0, 0);
        } else {
            setpgid(0, *pgid);
        }

        dup2(cmd.read_cha.fd, STDIN_FILENO);
        dup2(cmd.write_cha.fd, STDOUT_FILENO);

        if (cmd.read_cha.child_proc_closable) close(cmd.read_cha.fd);
        if (cmd.write_cha.child_proc_closable) close(cmd.write_cha.fd);

        for (int i = 0; i < cmd.redirects.count; i++) {
            FdRedirect fd_rd = *get_from_FdRedirectVec(&cmd.redirects, i);
            dup2(fd_rd.to, fd_rd.redirected);
        }

        for (int j = 0; j < cmd.other_open_channels.count; j++) {
            Channel cha = *get_from_ChannelVec(&cmd.other_open_channels, j);
            if (cha.child_proc_closable) close(cha.fd);
        }

        signal(SIGINT,  SIG_DFL);
        signal(SIGQUIT, SIG_DFL);
        signal(SIGTSTP, SIG_DFL);
        signal(SIGTTIN, SIG_DFL);
        signal(SIGTTOU, SIG_DFL);

        execvp(cmd.args[0], cmd.args);
        exit(67);
    }

    if (*pgid == -1) {
        setpgid(pid, pid);
        *pgid = pid;
    } else {
        setpgid(pid, *pgid);
    }
}


ExecutableCommand process_pl_command(PlCommand pl_cmd, int* pgid, ShellState* shstate) {

    StrVec arg_vec = make_StrVec(256);
    FdRedirectVec fd_redirs = make_FdRedirectVec(256);
    ChannelVec other_open_channels = make_ChannelVec(256);

    int ps_pipe_fds[2];

    for (int i = 0; i < pl_cmd.args.count; i++) {
        Arg arg = *get_from_ArgVec(&pl_cmd.args, i);
        Channel ps_cha;
        switch (arg.kind) {
            case ARG_LIT:
                push_to_StrVec(&arg_vec, arg.as.lit_str);
                break;
            case ARG_PS_IN:
                pipe(ps_pipe_fds);
                Channel write_cha = {.fd = ps_pipe_fds[1], .parent_proc_closable = true, .child_proc_closable = true};
                execute_command_pipeline(*(Pipeline*)arg.as.sub_pipeline, STDIN_CHA, write_cha, false, pgid, shstate);
                //close(ps_pipe_fds[1]);

                push_to_ChannelVec(&other_open_channels, (Channel){.fd = ps_pipe_fds[0], .parent_proc_closable = true, .child_proc_closable = false});
                char pname[64];
                snprintf(pname, sizeof(pname), "/dev/fd/%d", ps_pipe_fds[0]);
                push_to_StrVec(&arg_vec, strdup(pname));
                break;
            case ARG_PS_OUT:
                // TODO: do waiting correctly
                pipe(ps_pipe_fds);

                // TODO: don't do fcntl?
                fcntl(ps_pipe_fds[1], F_SETFD, FD_CLOEXEC);
                Channel read_cha = {.fd = ps_pipe_fds[0], .parent_proc_closable = true, .child_proc_closable = true};
                execute_command_pipeline(*(Pipeline*)arg.as.sub_pipeline, read_cha, STDOUT_CHA, false, pgid, shstate);
                fcntl(ps_pipe_fds[1], F_SETFD, 0);
                push_to_ChannelVec(&other_open_channels, (Channel){.fd = ps_pipe_fds[0], .parent_proc_closable = true, .child_proc_closable = false});
                //close(ps_pipe_fds[0]);
                char name[64];
                snprintf(name, sizeof(name), "/dev/fd/%d", ps_pipe_fds[1]);
                push_to_StrVec(&arg_vec, strdup(name));
                break;
            default:
                exit(67);
        }
    }

    for (int j = 0; j < pl_cmd.redirects.count; j++) {
        Redirect redir = *get_from_RedirectVec(&pl_cmd.redirects, j);
        FdRedirect fd_redir;
        int fd;
        Channel file_cha;
        switch (redir.kind) {
            case REDIR_READ:  
                fd = open(redir.as.filename, O_RDONLY, 0644);
                file_cha = (Channel){
                    .fd = fd, 
                    .parent_proc_closable = true, 
                    .child_proc_closable = true,
                };
                push_to_ChannelVec(&other_open_channels, file_cha);
                fd_redir = (FdRedirect){.to = fd, .redirected = STDIN_FILENO};
                break;
            case REDIR_WRITE_APPEND:
                fd = open(redir.as.filename, O_WRONLY | O_CREAT | O_APPEND, 0644);
                file_cha = (Channel){
                    .fd = fd,
                    .parent_proc_closable = true,
                    .child_proc_closable = true,
                };
                push_to_ChannelVec(&other_open_channels, file_cha);
                fd_redir = (FdRedirect){.to = fd, .redirected = STDOUT_FILENO};
                break;
            case REDIR_WRITE_TRUNC:
                fd = open(redir.as.filename, O_WRONLY | O_CREAT | O_TRUNC, 0644);
                file_cha = (Channel){
                    .fd = fd,
                    .parent_proc_closable = true,
                    .child_proc_closable = true,
                };
                push_to_ChannelVec(&other_open_channels, file_cha);
                fd_redir = (FdRedirect){.to = fd, .redirected = STDOUT_FILENO};
                break;
            default:
                exit(67);
        }
        push_to_FdRedirectVec(&fd_redirs, fd_redir);
}

    // TODO: I think dealiasing has certain preconditions, like arg0 not being expanded before(?).
    char** args = malloc(sizeof(char*) * (arg_vec.count + 1));
    for (int i = 0; i < arg_vec.count; i++) {
        if (i == 0) {
            char* arg = *get_from_StrVec(&arg_vec, 0);
            MaybeStr m_dealias = get_map_value(&shstate->aliases, arg);
            if (m_dealias.some) {
                args[0] = m_dealias.value;
            } else {
                args[0] = arg;
            }
        } else {
            args[i] = *get_from_StrVec(&arg_vec, i);
        }
    }
    args[arg_vec.count] = NULL;
    
    return (ExecutableCommand){
        .args = args, 
        .redirects = fd_redirs, 
        .other_open_channels = other_open_channels
    };
}


void execute_command_pipeline(Pipeline pipeline, Channel read_cha, Channel write_cha, bool wait_here, int* pgid, ShellState* shstate) {

    Channel trunk_read_cha = read_cha;

    for (int i = 0; i < pipeline.count; i++) {
        PlCommand cmd = *get_from_Pipeline(&pipeline, i);
        ExecutableCommand exec_cmd = process_pl_command(cmd, pgid, shstate);

        exec_cmd.read_cha = trunk_read_cha;

        exec_cmd.write_cha = write_cha;
        
        if (cmd.conn == CONN_PIPE) {
            int pipe_fds[2];
            pipe(pipe_fds);
            exec_cmd.write_cha = (Channel){.fd = pipe_fds[1], .parent_proc_closable = true, .child_proc_closable = true};
            trunk_read_cha = (Channel){.fd = pipe_fds[0], .parent_proc_closable = true, .child_proc_closable = true};
        }

        if (!handle_builtin(exec_cmd, shstate)) {
            execute_cmd(exec_cmd, pgid);
        }

        if (exec_cmd.read_cha.parent_proc_closable) close(exec_cmd.read_cha.fd);
        if (exec_cmd.write_cha.parent_proc_closable) close(exec_cmd.write_cha.fd);

        for (int j = 0; j < exec_cmd.other_open_channels.count; j++) {
            Channel cha = *get_from_ChannelVec(&exec_cmd.other_open_channels, j);
            if (cha.parent_proc_closable) close(cha.fd);
        }
    }

    if (wait_here) {
        tcsetpgrp(STDIN_FILENO, *pgid); 
        while (waitpid(-*pgid, NULL, 0) > 0); 
        tcsetpgrp(STDIN_FILENO, getpgrp());
    }
}

void execute_pipeline_tree(Pipeline root_pl, ShellState* shstate, int stdin_fd, int stdout_fd) {
    int pgid = -1;
    Channel read_cha = {.fd = stdin_fd, .parent_proc_closable = false, .child_proc_closable = false};
    Channel write_cha = {.fd = stdout_fd, .parent_proc_closable = false, .child_proc_closable = false};
    execute_command_pipeline(root_pl, read_cha, write_cha, true, &pgid, shstate); 
}