#include "../include/Exec.hpp"
#include "../include/env.hpp"

#if OS_FAMILY == OS_FAMILY_POSIX
#include <sys/wait.h>
#include <unistd.h>
#endif

using namespace std;

const int BUFSIZE = 4096;

// TODO: Redirect stderr to 
namespace sibs
{
#if OS_FAMILY == OS_FAMILY_POSIX
    Result<ExecResult> exec(const std::vector<FileString> &args, bool print_instead_of_pipe)
    {
        char buffer[BUFSIZE];
        std::string execStdout;

        if(args.empty())
            return Result<ExecResult>::Err("exec requires at least one argument (the program name)");

        std::vector<const char*> exec_args;
        for(const FileString &arg : args) {
            exec_args.push_back(arg.c_str());
        }
        exec_args.push_back(nullptr);

        int fd[2];
        if(!print_instead_of_pipe && pipe(fd) == -1)
            return Result<ExecResult>::Err(strerror(errno));

        pid_t pid = fork();
        if(pid == -1) {
            if(!print_instead_of_pipe) {
                close(fd[0]);
                close(fd[1]);
            }
            return Result<ExecResult>::Err("Failed to exec " + args[0] + " (failed to fork)");
        } else if(pid == 0) { // child
            if(!print_instead_of_pipe) {
                dup2(fd[1], STDOUT_FILENO);
                close(fd[0]);
                close(fd[1]);
            }
            execvp(exec_args[0], (char* const*)exec_args.data());
            perror("execvp");
            _exit(127);
        } else { // parent
            if(!print_instead_of_pipe)
                close(fd[1]);
        }

        if(!print_instead_of_pipe) {
            for(;;) {
                ssize_t bytes_read = read(fd[0], buffer, sizeof(buffer));
                if(bytes_read == 0) {
                    break;
                } else if(bytes_read == -1) {
                    std::string err_msg = "Failed to read from pipe to program " + args[0] + ", error: " + strerror(errno);
                    kill(pid, SIGTERM);
                    close(fd[0]);
                    return Result<ExecResult>::Err(err_msg);
                }

                execStdout.append(buffer, bytes_read);
            }
        }

        int status = 0;
        if(waitpid(pid, &status, 0) == -1) {
            std::string err_msg = std::string("waitpid failed, error: ") + strerror(errno);
            if(!print_instead_of_pipe)
                close(fd[0]);
            return Result<ExecResult>::Err(err_msg);
        }
        if(!print_instead_of_pipe)
            close(fd[0]);

        if(WIFEXITED(status))
        {
            int returned = WEXITSTATUS(status);
            ExecResult execResult;
            execResult.execStdout = move(execStdout);
            execResult.exitCode = returned;
            return Result<ExecResult>::Ok(execResult);
        }
        else if(WIFSIGNALED(status))
        {
            int signum = WSTOPSIG(status);
            string errMsg = "Exited due to receiving signal ";
            errMsg += to_string(signum);
            return Result<ExecResult>::Err(errMsg);
        }
        else if(WIFSTOPPED(status))
        {
            int signum = WSTOPSIG(status);
            string errMsg = "Stopped due to receiving signal ";
            errMsg += to_string(signum);
            return Result<ExecResult>::Err(errMsg);
        }
        else
        {
            string errMsg = "exec unexpected error on waitpid: ";
            errMsg += to_string(status);
            return Result<ExecResult>::Err(errMsg);
        }
    }

#else
    static FileString escape_arg(const FileString &arg) {
        FileString escaped = TINYDIR_STRING("\"");
        for(_tinydir_char_t c : arg) {
            if(c == '"') {
                escaped += TINYDIR_STRING("\"\"");
            } else {
                escaped += c;
            }
        }
        escaped += TINYDIR_STRING("\"");
        return escaped;
    }
    
    static FileString command_list_to_command_string(const std::vector<FileString> &args) {
        FileString cmd;
        for(size_t i = 0; i < args.size(); ++i) {
            if(i > 0)
                cmd += TINYDIR_STRING(" ");
            cmd += escape_arg(args[i]);
        }
        return cmd;
    }

    // Currently stdout is read in text mode so \n is replaced with \r\n, should we read in binary mode instead?
    Result<ExecResult> exec(const std::vector<FileString> &args, bool print_instead_of_pipe)
    {
        FileString cmdNonConst = command_list_to_command_string(args);
        std::string execStdout;

        SECURITY_ATTRIBUTES saAttr;
        saAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
        saAttr.bInheritHandle = TRUE;
        saAttr.lpSecurityDescriptor = nullptr;

        HANDLE childReadHandle = nullptr;
        HANDLE childStdoutHandle = nullptr;

        if(!print_instead_of_pipe) {
            if (!CreatePipe(&childReadHandle, &childStdoutHandle, &saAttr, 0))
            {
                string errMsg = "exec unexpected error: ";
                errMsg += toUtf8(getLastErrorAsString());
                return Result<ExecResult>::Err(errMsg);
            }

            if (!SetHandleInformation(childReadHandle, HANDLE_FLAG_INHERIT, 0))
                goto cleanupAndExit;
        }

        PROCESS_INFORMATION piProcInfo;
        ZeroMemory(&piProcInfo, sizeof(PROCESS_INFORMATION));

        STARTUPINFO siStartInfo;
        ZeroMemory(&siStartInfo, sizeof(STARTUPINFO));
        siStartInfo.cb = sizeof(STARTUPINFO);
        siStartInfo.hStdError = nullptr;
        siStartInfo.hStdOutput = print_instead_of_pipe ? nullptr : childStdoutHandle;
        siStartInfo.hStdInput = nullptr;
        siStartInfo.dwFlags |= STARTF_USESTDHANDLES;

        DWORD exitCode;

        if (!CreateProcessW(nullptr, (LPWSTR)cmdNonConst.data(), nullptr, nullptr, TRUE, 0, nullptr, nullptr, &siStartInfo, &piProcInfo))
            goto cleanupAndExit;

        if(!print_instead_of_pipe) {
            CloseHandle(childStdoutHandle);
            childStdoutHandle = nullptr;

            DWORD bytesRead;
            CHAR buffer[BUFSIZE];
            while (true)
            {
                BOOL bSuccess = ReadFile(childReadHandle, buffer, BUFSIZE, &bytesRead, nullptr);
                if (!bSuccess || bytesRead == 0)
                    break;

                execStdout.append(buffer, bytesRead);
                if (print)
                    printf("%.*s", bytesRead, buffer);
            }
        }

        WaitForSingleObject(piProcInfo.hProcess, INFINITE);
        GetExitCodeProcess(piProcInfo.hProcess, &exitCode);
        CloseHandle(piProcInfo.hProcess);
        CloseHandle(piProcInfo.hThread);

        {
            ExecResult execResult;
            execResult.execStdout = move(execStdout);
            execResult.exitCode = exitCode;
            CloseHandle(childReadHandle);
            return Result<ExecResult>::Ok(execResult);
        }

    cleanupAndExit:
        string errMsg = "exec unexpected error: ";
        errMsg += toUtf8(getLastErrorAsString());
        if(childReadHandle)
            CloseHandle(childReadHandle);
        if(childStdoutHandle)
            CloseHandle(childStdoutHandle);
        return Result<ExecResult>::Err(errMsg);
    }
#endif
}