#include "../include/Program.hpp"
#include <unistd.h>
#include <sys/wait.h>
#include <errno.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <unordered_map>
#include <mutex>
#include <signal.h>
#include <string>

#define READ_END 0
#define WRITE_END 1

int accumulate_string(char *data, int size, void *userdata) {
    std::string *str = (std::string*)userdata;
    if(str->size() + size > 1024 * 1024 * 100) // 100mb sane limit, TODO: make configurable
        return 1;
    str->append(data, size);
    return 0;
}

struct ReadWriteProgram {
    pid_t pid = -1;
    int read_fd = -1;
    int write_fd = -1;
};

struct ThreadProgram {
    ReadWriteProgram program;
    bool killed;
};

static std::unordered_map<std::thread::id, ThreadProgram> thread_current_program;
static std::mutex thread_current_program_mutex;

class CurrentThreadProgram {
public:
    CurrentThreadProgram() {
        std::lock_guard<std::mutex> lock(thread_current_program_mutex);
        ThreadProgram thread_program;
        thread_program.program.pid = -1;
        thread_program.program.read_fd = -1;
        thread_program.program.write_fd = -1;
        thread_program.killed = false;
        thread_current_program[std::this_thread::get_id()] = std::move(thread_program);
    }

    ~CurrentThreadProgram() {
        std::lock_guard<std::mutex> lock(thread_current_program_mutex);
        thread_current_program.erase(std::this_thread::get_id());
    }

    // TODO: Make sure the thread specific program has been stopped before this is called. exec_program_pipe needs to be modified for that
    void set(ReadProgram read_program) {
        std::lock_guard<std::mutex> lock(thread_current_program_mutex);
        auto it = thread_current_program.find(std::this_thread::get_id());
        if(it != thread_current_program.end()) {
            it->second.program.pid = read_program.pid;
            it->second.program.read_fd = read_program.read_fd;
            it->second.program.write_fd = -1;
        }
    }

    // TODO: Make sure the thread specific program has been stopped before this is called. exec_program_pipe needs to be modified for that
    void set(ReadWriteProgram program) {
        std::lock_guard<std::mutex> lock(thread_current_program_mutex);
        auto it = thread_current_program.find(std::this_thread::get_id());
        if(it != thread_current_program.end())
            it->second.program = std::move(program);
    }

    void clear() {
        std::lock_guard<std::mutex> lock(thread_current_program_mutex);
        auto it = thread_current_program.find(std::this_thread::get_id());
        if(it != thread_current_program.end()) {
            it->second.program.pid = -1;
            it->second.program.read_fd = -1;
            it->second.program.write_fd = -1;
        }
    }

    // TODO: This same mutex should be used in the exec_... functions when they do kill() etc to make sure we dont accidentally kill another program here if another process gets the killed process id!
    void kill_in_thread(const std::thread::id &thread_id) {
        std::lock_guard<std::mutex> lock(thread_current_program_mutex);
        auto it = thread_current_program.find(thread_id);
        if(it != thread_current_program.end()) {
            if(it->second.program.read_fd != -1) {
                close(it->second.program.read_fd);
                it->second.program.read_fd = -1;
            }
            if(it->second.program.write_fd != -1) {
                close(it->second.program.write_fd);
                it->second.program.write_fd = -1;
            }
            if(it->second.program.pid != -1) {
                kill(it->second.program.pid, SIGKILL);
                int status;
                waitpid(it->second.program.pid, &status, 0);
                it->second.program.pid = -1;
            }
            it->second.killed = true;
        }
    }

    bool is_killed() {
        std::lock_guard<std::mutex> lock(thread_current_program_mutex);
        auto it = thread_current_program.find(std::this_thread::get_id());
        if(it != thread_current_program.end())
            return it->second.killed;
        return false;
    }
};

thread_local CurrentThreadProgram current_thread_program;

int exec_program_pipe(const char **args, ReadProgram *read_program) {
    read_program->pid = -1;
    read_program->read_fd = -1;

    /* 1 arguments */
    if(args[0] == NULL)
        return -1;

    if(current_thread_program.is_killed())
        return -1;
    
    int fd[2];
    if(pipe(fd) == -1) {
        perror("Failed to open pipe");
        return -2;
    }

    pid_t pid = vfork();
    if(pid == -1) {
        perror("Failed to vfork");
        close(fd[READ_END]);
        close(fd[WRITE_END]);
        return -3;
    } else if(pid == 0) { /* child */
        dup2(fd[WRITE_END], STDOUT_FILENO);
        close(fd[READ_END]);
        close(fd[WRITE_END]);

        execvp(args[0], (char* const*)args);
        perror("execvp");
        _exit(127);
    } else { /* parent */
        close(fd[WRITE_END]);
        read_program->pid = pid;
        read_program->read_fd = fd[READ_END];
        current_thread_program.set(*read_program);
        return 0;
    }
}

static int exec_program_pipe2(const char **args, ReadWriteProgram *program) {
    program->pid = -1;
    program->read_fd = -1;
    program->write_fd = -1;

    /* 1 arguments */
    if(args[0] == NULL)
        return -1;

    if(current_thread_program.is_killed())
        return -1;
    
    int read_fd[2];
    if(pipe(read_fd) == -1) {
        perror("Failed to open pipe");
        return -2;
    }

    int write_fd[2];
    if(pipe(write_fd) == -1) {
        close(read_fd[0]);
        close(read_fd[1]);
        perror("Failed to open pipe");
        return -2;
    }

    pid_t pid = vfork();
    if(pid == -1) {
        perror("Failed to vfork");
        close(read_fd[READ_END]);
        close(read_fd[WRITE_END]);
        close(write_fd[READ_END]);
        close(write_fd[WRITE_END]);
        return -3;
    } else if(pid == 0) { /* child */
        dup2(read_fd[WRITE_END], STDOUT_FILENO);
        close(read_fd[READ_END]);
        close(read_fd[WRITE_END]);

        dup2(write_fd[READ_END], STDIN_FILENO);
        close(write_fd[READ_END]);
        close(write_fd[WRITE_END]);

        execvp(args[0], (char* const*)args);
        perror("execvp");
        _exit(127);
    } else { /* parent */
        close(read_fd[WRITE_END]);
        close(write_fd[READ_END]);
        program->pid = pid;
        program->read_fd = read_fd[READ_END];
        program->write_fd = write_fd[WRITE_END];
        current_thread_program.set(*program);
        return 0;
    }
}

int exec_program_write_stdin(const char **args, const char *str, size_t size, ProgramOutputCallback output_callback, void *userdata, int buffer_size) {
    ReadWriteProgram program;
    int res = exec_program_pipe2(args, &program);
    if(res != 0)
        return res;

    int result = 0;
    int status;
    int exit_status;

    assert(buffer_size >= 1 && buffer_size <= 65536);
    char *buffer = (char*)alloca(buffer_size + 1);

    const ssize_t write_buffer_size = 8192;
    size_t write_offset = 0;
    while(write_offset < size) {
        ssize_t write_size = (ssize_t)size - (ssize_t)write_offset;
        if(write_size > write_buffer_size)
            write_size = write_buffer_size;

        ssize_t bytes_written = write(program.write_fd, str + write_offset, write_size);
        if(bytes_written == -1) {
            int err = errno;
            fprintf(stderr, "Failed to write to pipe to program %s, error: %s\n", args[0], strerror(err));
            result = -err;
            break;
        }

        if(bytes_written < write_size)
            write_size = bytes_written;

        write_offset += write_size;
    }

    close(program.write_fd);

    if(result == 0) {
        for(;;) {
            ssize_t bytes_read = read(program.read_fd, buffer, buffer_size);
            if(bytes_read == 0) {
                break;
            } else if(bytes_read == -1) {
                int err = errno;
                fprintf(stderr, "Failed to read from pipe to program %s, error: %s\n", args[0], strerror(err));
                result = -err;
                break;
            }

            buffer[bytes_read] = '\0';
            if(output_callback) {
                result = output_callback(buffer, bytes_read, userdata);
                if(result != 0)
                    break;
            }
        }
    }

    // TODO: Set program.pid to -1 and with currenthreadprogram mutex. Same in other places
    if(result != 0)
        kill(program.pid, SIGKILL);

    if(waitpid(program.pid, &status, 0) == -1) {
        perror("waitpid failed");
        result = -5;
        goto cleanup;
    }

    if(!WIFEXITED(status)) {
        result = -4;
        goto cleanup;
    }

    exit_status = WEXITSTATUS(status);
    if(exit_status != 0) {
        fprintf(stderr, "Failed to execute program (");
        const char **arg = args;
        while(*arg) {
            if(arg != args)
                fputc(' ', stderr);
            fprintf(stderr, "'%s'", *arg);
            ++arg;
        }
        fprintf(stderr, "), exit status %d\n", exit_status);
        result = -exit_status;
    }

    cleanup:
    program_clear_current_thread();
    close(program.read_fd);
    return result;
}

int exec_program(const char **args, ProgramOutputCallback output_callback, void *userdata, int buffer_size) {
    int allowed_exit_status[1] = {0};
    return exec_program(args, output_callback, userdata, allowed_exit_status, 1, buffer_size);
}

int exec_program(const char **args, ProgramOutputCallback output_callback, void *userdata, int *allowed_exit_status, int num_allowed_exit_status, int buffer_size) {
    ReadProgram read_program;
    int res = exec_program_pipe(args, &read_program);
    if(res != 0)
        return res;

    int result = 0;
    int status;
    int exit_status;
    bool is_error = true;

    assert(buffer_size >= 1 && buffer_size <= 65536);
    char *buffer = (char*)alloca(buffer_size + 1);

    for(;;) {
        ssize_t bytes_read = read(read_program.read_fd, buffer, buffer_size);
        if(bytes_read == 0) {
            break;
        } else if(bytes_read == -1) {
            int err = errno;
            fprintf(stderr, "Failed to read from pipe to program %s, error: %s\n", args[0], strerror(err));
            result = -err;
            break;
        }

        buffer[bytes_read] = '\0';
        if(output_callback) {
            result = output_callback(buffer, bytes_read, userdata);
            if(result != 0)
                break;
        }
    }

    if(result != 0)
        kill(read_program.pid, SIGKILL);

    if(waitpid(read_program.pid, &status, 0) == -1) {
        perror("waitpid failed");
        result = -5;
        goto cleanup;
    }

    if(!WIFEXITED(status)) {
        result = -4;
        goto cleanup;
    }

    exit_status = WEXITSTATUS(status);
    for(int i = 0; i < num_allowed_exit_status; ++i) {
        if(exit_status == allowed_exit_status[i]) {
            is_error = false;
            break;
        }
    }

    if(is_error) {
        fprintf(stderr, "Failed to execute program (");
        const char **arg = args;
        while(*arg) {
            if(arg != args)
                fputc(' ', stderr);
            fprintf(stderr, "'%s'", *arg);
            ++arg;
        }
        fprintf(stderr, "), exit status %d\n", exit_status);
        result = -exit_status;
        if(result == 0)
            result = -1;
    }

    cleanup:
    program_clear_current_thread();
    close(read_program.read_fd);
    return result;
}

int wait_program(pid_t process_id) {
    int status;
    if(waitpid(process_id, &status, 0) == -1) {
        int err = -errno;
        perror("waitpid failed");
        return err;
    }

    if(!WIFEXITED(status))
        return -4;

    return WEXITSTATUS(status);
}

int wait_program_non_blocking(pid_t process_id, int *status) {
    int s;
    int wait_result = waitpid(process_id, &s, WNOHANG);
    if(wait_result == -1) {
        int err = -errno;
        perror("waitpid failed");
        *status = err;
        return 0;
    } else if(wait_result == 0) {
        /* the child process is still running */
        *status = 0;
        return 0;
    }

    if(!WIFEXITED(s)) {
        *status = -4;
        return 0;
    }

    *status = WEXITSTATUS(s);
    return 1;
}

// TODO: Verify if this can cause issues when |result_process_id| is null, because |args| may be deallocated
// by the time its used in the last execvp.
int exec_program_async(const char **args, pid_t *result_process_id) {
    /* 1 arguments */
    if(args[0] == NULL)
        return -1;

    pid_t pid = vfork();
    if(pid == -1) {
        int err = errno;
        perror("Failed to vfork");
        return -err;
    } else if(pid == 0) { /* child */
        if(result_process_id) {
            execvp(args[0], (char* const*)args);
            perror("execvp");
            _exit(127);
        } else {
            setsid();
            signal(SIGHUP, SIG_IGN);

            // Daemonize child to make the init process the parent which will reap the zombie child
            pid_t second_child = vfork();
            if(second_child == 0) { // child
                execvp(args[0], (char* const*)args);
                perror("execvp");
                _exit(127);
            } else if(second_child != -1) {
                _exit(0);
            }
        }
    } else { /* parent */
        if(result_process_id)
            *result_process_id = pid;
        else
            waitpid(pid, nullptr, 0);
    }
    return 0;
}

void program_clear_current_thread() {
    current_thread_program.clear();
}

void program_kill_in_thread(std::thread::id thread_id) {
    current_thread_program.kill_in_thread(thread_id);
}

bool program_is_dead_in_thread(std::thread::id thread_id) {
    std::lock_guard<std::mutex> lock(thread_current_program_mutex);
    auto it = thread_current_program.find(thread_id);
    if(it != thread_current_program.end())
        return it->second.killed;
    return false;
}

bool program_is_dead_in_current_thread() {
    return current_thread_program.is_killed();
}