/*
 * Copyright 2019 Bloomberg Finance LP
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <buildboxcommon_exception.h>
#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_stringutils.h>
#include <buildboxcommon_systemutils.h>

#include <array>
#include <cerrno>
#include <chrono>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <fcntl.h>
#include <grp.h>
#include <memory>
#include <optional>
#include <pwd.h>
#include <sstream>
#include <stdlib.h>
#include <string>
#include <sys/types.h>
#include <sys/utsname.h>
#include <sys/wait.h>
#include <system_error>
#include <unistd.h>

extern "C" char **environ;

using namespace buildboxcommon;
using namespace system_utils_exitcode;

std::array<int, 2> SystemUtils::createPipe()
{
    std::array<int, 2> pipe_fds = {-1, -1};

#ifdef HAVE_PIPE2
    // Prefer atomic `pipe2()` to avoid race conditions
    if (pipe2(pipe_fds.data(), O_CLOEXEC) == 0) {
        return pipe_fds;
    }

    // Fall through to pipe() + fcntl() if pipe2() is not available
    if (errno != ENOSYS) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(std::system_error, errno,
                                              std::system_category,
                                              "Error calling `pipe2()`");
    }
#endif

    if (pipe(pipe_fds.data()) == -1) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(std::system_error, errno,
                                              std::system_category,
                                              "Error calling `pipe()`");
    }
    if (fcntl(pipe_fds[0], F_SETFD, FD_CLOEXEC) == -1 ||
        fcntl(pipe_fds[1], F_SETFD, FD_CLOEXEC) == -1) {
        close(pipe_fds[0]);
        close(pipe_fds[1]);
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category,
            "Failed to set `FD_CLOEXEC` on pipe file descriptors");
    }

    return pipe_fds;
}

int SystemUtils::executeCommand(
    const std::vector<std::string> &command, bool closeFds,
    const std::optional<std::unordered_map<std::string, std::string>>
        &environment)
{
    if (closeFds) {
        // POSIX.1
#ifdef HAVE_CLOSEFROM
        // close() can be very slow with higher fd limits.
        // closefrom() uses optimized system calls making it faster than
        // iterating up to _SC_OPEN_MAX
        closefrom(STDERR_FILENO + 1);
#else
        const int fdLimit = static_cast<int>(sysconf(_SC_OPEN_MAX));
        for (int i = STDERR_FILENO + 1; i < fdLimit; i++) {
            (void)close(i);
        }
#endif
    }

    const auto argc = command.size();

    // `execvp()` takes a `const char*[]`, making the conversion:
    const auto argv = std::make_unique<const char *[]>(argc + 1);
    for (unsigned int i = 0; i < argc; ++i) {
        argv[i] = command[i].c_str();
    }
    argv[argc] = nullptr;

    std::vector<std::string> envStrs;
    std::unique_ptr<const char *[]> envp = nullptr;
    if (environment.has_value()) {
        const auto envSize = environment->size();
        // convert to K=V format
        envStrs.reserve(envSize);
        for (const auto &[envKey, envValue] : *environment) {
            envStrs.emplace_back(envKey + "=" + envValue);
        }
        // convert to cstr's
        envp = std::make_unique<const char *[]>(envSize + 1);
        for (size_t i = 0; i < envSize; ++i) {
            envp[i] = envStrs[i].c_str();
        }
        envp[envSize] = nullptr;
    }

    execve(argv[0],
           const_cast< // NOLINT(cppcoreguidelines-pro-type-const-cast)
               char *const *>(argv.get()),
           envp == nullptr
               ? environ
               : const_cast< // NOLINT(cppcoreguidelines-pro-type-const-cast)
                     char *const *>(envp.get()));

    // Following the Bash convention for exit codes.
    // (https://gnu.org/software/bash/manual/html_node/Exit-Status.html)
    if (errno == ENOENT) {
        return COMMAND_NOT_FOUND;
    }
    else {
        return COMMAND_CANNOT_EXECUTE;
    }
}

int SystemUtils::executeCommandAndWait(
    const std::vector<std::string> &command, bool closeFds,
    const std::optional<std::unordered_map<std::string, std::string>>
        &environment)
{
    const pid_t pid = fork();
    if (pid == -1) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(std::system_error, errno,
                                              std::system_category,
                                              "Error calling fork()");
    }

    if (pid == 0) { // Child process
        const auto status = executeCommand(command, closeFds, environment);

        // Only reached if the underlying call to exec() failed.
        perror(nullptr);
        _Exit(status);
    }
    else { // Parent process
        return waitPid(pid);
    }
}

SystemUtils::ExecutionResult SystemUtils::executeCommandWithResult(
    const std::vector<std::string> &command, bool pipeStdOut, bool pipeStdErr,
    const std::optional<std::string> &workingDirectory,
    const std::optional<std::string> &stdinInput,
    const std::optional<std::chrono::milliseconds> &timeout)
{
    // Convert the command to a char*[]
    size_t argc = command.size();
    std::unique_ptr<const char *[]> argv(new const char *[argc + 1]);
    for (size_t i = 0; i < argc; ++i) {
        argv[i] = command[i].c_str();
    }
    argv[argc] = nullptr;

    // Create pipes as needed
    std::array<int, 2> stdinPipeFDs = {-1, -1};
    if (stdinInput.has_value()) {
        stdinPipeFDs = createPipe();
        // Set stdin write FD to non-blocking mode
        int flags = fcntl(stdinPipeFDs[1], F_GETFL, 0);
        if (flags < 0 ||
            fcntl(stdinPipeFDs[1], F_SETFL, flags | O_NONBLOCK) < 0) {
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                std::system_error, errno, std::system_category,
                "Error marking stdin pipe as non-blocking with fcntl()");
        }
    }
    std::array<int, 2> stdOutPipeFDs = {-1, -1};
    if (pipeStdOut) {
        stdOutPipeFDs = createPipe();
    }
    std::array<int, 2> stdErrPipeFDs = {-1, -1};
    if (pipeStdErr) {
        stdErrPipeFDs = createPipe();
    }

    const auto pid = fork();
    if (pid == -1) {
        BUILDBOX_LOG_ERROR("Error calling `fork()`: " << strerror(errno));
        throw std::system_error(errno, std::system_category());
    }

    if (pid == 0) {
        // Child process

        // Switch to working directory if requested
        const int chdir_status = workingDirectory.has_value()
                                     ? chdir(workingDirectory->c_str())
                                     : 0;

        if (chdir_status != 0) {
            // Following the Bash convention for exit codes.
            // Matching posix_spawn() returns exit code 127 if chdir() failure
            // https://pubs.opengroup.org/onlinepubs/9799919799/functions/posix_spawn.html
            _Exit(COMMAND_NOT_FOUND);
        }

        // Redirect stdin from pipe if input provided
        if (stdinInput.has_value()) {
            close(stdinPipeFDs[1]); // Close write end in child
            dup2(stdinPipeFDs[0], STDIN_FILENO);
            close(stdinPipeFDs[0]);
        }

        // Redirect stdout if requested
        if (pipeStdOut) {
            close(stdOutPipeFDs[0]);
            dup2(stdOutPipeFDs[1],
                 STDOUT_FILENO); // redirect stdout to input end of pipe
            close(stdOutPipeFDs[1]);
        }

        // Redirect stderr if requested
        if (pipeStdErr) {
            close(stdErrPipeFDs[0]);
            dup2(stdErrPipeFDs[1],
                 STDERR_FILENO); // redirect stderr to input end of pipe
            close(stdErrPipeFDs[1]);
        }

        const auto exec_status =
            execvp(argv[0],
                   const_cast< // NOLINT(cppcoreguidelines-pro-type-const-cast)
                       char *const *>(argv.get()));

        int exit_code = 1;
        if (exec_status != 0) {
            const auto exec_error = errno;
            // Following the Bash convention for exit codes.
            // (https://gnu.org/software/bash/manual/html_node/Exit-Status.html)
            if (exec_error == ENOENT) {
                exit_code = COMMAND_NOT_FOUND;
            }
            else {
                exit_code = COMMAND_CANNOT_EXECUTE;
            }
        }

        _Exit(exit_code);
    }

    // Parent process
    ExecutionResult result;

    // Close read end of stdin pipe, we'll write to it
    if (stdinInput.has_value()) {
        close(stdinPipeFDs[0]);
    }

    // Close write ends of stdout/stderr pipes
    if (pipeStdOut) {
        close(stdOutPipeFDs[1]);
    }
    if (pipeStdErr) {
        close(stdErrPipeFDs[1]);
    }

    // Track stdin write progress
    size_t stdinOffset = 0;
    bool stdinOpen = stdinInput.has_value();

    if (stdinOpen && stdinInput.has_value() && stdinInput->empty()) {
        close(stdinPipeFDs[1]);
        stdinOpen = false;
    }

    // Track overall timeout
    const auto startTime = std::chrono::steady_clock::now();
    std::optional<std::chrono::steady_clock::time_point> deadline;
    if (timeout.has_value()) {
        deadline = startTime + timeout.value();
    }
    bool timedOut = false;

    // Get the output from the child process and write to stdin
    fd_set readFDSet;
    fd_set writeFDSet;
    FD_ZERO(&readFDSet);
    FD_ZERO(&writeFDSet);

    if (pipeStdOut) {
        FD_SET(stdOutPipeFDs[0], &readFDSet);
    }
    if (pipeStdErr) {
        FD_SET(stdErrPipeFDs[0], &readFDSet);
    }

    constexpr int BUFFER_SIZE = 4096;
    std::array<char, BUFFER_SIZE> buffer{};

    while ((pipeStdOut && FD_ISSET(stdOutPipeFDs[0], &readFDSet)) ||
           (pipeStdErr && FD_ISSET(stdErrPipeFDs[0], &readFDSet)) ||
           stdinOpen) {
        fd_set currentReadFDs = readFDSet;
        fd_set currentWriteFDs;
        FD_ZERO(&currentWriteFDs);

        // Add stdin to write set if we still have data to write
        if (stdinOpen && stdinInput.has_value() &&
            stdinOffset < stdinInput->size()) {
            FD_SET(stdinPipeFDs[1], &currentWriteFDs);
        }

        // Check if timeout expired
        if (deadline.has_value() &&
            std::chrono::steady_clock::now() >= deadline.value()) {
            timedOut = true;
            break;
        }

        constexpr int TIMEOUT_SECONDS = 5;
        struct timeval timeout_tv{};
        timeout_tv.tv_sec = TIMEOUT_SECONDS;
        timeout_tv.tv_usec = 0;
        if (select(FD_SETSIZE, &currentReadFDs, &currentWriteFDs, nullptr,
                   &timeout_tv) < 0 &&
            errno != EINTR) {
            int selectErrno = errno;
            if (pipeStdOut && FD_ISSET(stdOutPipeFDs[0], &readFDSet)) {
                close(stdOutPipeFDs[0]);
            }
            if (pipeStdErr && FD_ISSET(stdErrPipeFDs[0], &readFDSet)) {
                close(stdErrPipeFDs[0]);
            }
            if (stdinOpen) {
                close(stdinPipeFDs[1]);
            }
            kill(pid, SIGKILL);
            waitPid(pid);
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                std::system_error, selectErrno, std::system_category,
                "Error in select()");
        }

        // Handle stdout
        if (pipeStdOut && FD_ISSET(stdOutPipeFDs[0], &currentReadFDs)) {
            ssize_t bytesRead =
                read(stdOutPipeFDs[0], buffer.data(), buffer.size());
            if (bytesRead > 0) {
                result.d_stdOut.append(buffer.data(),
                                       static_cast<size_t>(bytesRead));
            }
            else {
                close(stdOutPipeFDs[0]);
                FD_CLR(stdOutPipeFDs[0], &readFDSet);
            }
        }

        // Handle stderr
        if (pipeStdErr && FD_ISSET(stdErrPipeFDs[0], &currentReadFDs)) {
            ssize_t bytesRead =
                read(stdErrPipeFDs[0], buffer.data(), buffer.size());
            if (bytesRead > 0) {
                result.d_stdErr.append(buffer.data(),
                                       static_cast<size_t>(bytesRead));
            }
            else {
                close(stdErrPipeFDs[0]);
                FD_CLR(stdErrPipeFDs[0], &readFDSet);
            }
        }

        // Handle stdin writes
        if (stdinOpen && stdinInput.has_value() &&
            FD_ISSET(stdinPipeFDs[1], &currentWriteFDs)) {
            size_t remaining = stdinInput->size() - stdinOffset;
            ssize_t bytesWritten =
                write(stdinPipeFDs[1], &(*stdinInput)[stdinOffset], remaining);

            if (bytesWritten > 0) {
                stdinOffset += static_cast<size_t>(bytesWritten);

                // If all data written, close stdin pipe to signal EOF
                if (stdinOffset >= stdinInput->size()) {
                    close(stdinPipeFDs[1]);
                    stdinOpen = false;
                }
            }
            else if (bytesWritten < 0 && errno != EAGAIN &&
                     errno != EWOULDBLOCK) {
                // Error writing to stdin (not just "would block")
                if (errno != EINTR && errno != EPIPE) {
                    BUILDBOX_LOG_WARNING(
                        "Error writing to child stdin: " << strerror(errno));
                }
                close(stdinPipeFDs[1]);
                stdinOpen = false;
            }
        }
    }
    if (!timedOut) {
        // Calculate remaining timeout for waitpid
        std::optional<std::chrono::milliseconds> remainingTimeout;
        if (deadline.has_value()) {
            auto now = std::chrono::steady_clock::now();
            if (now >= deadline.value()) {
                timedOut = true;
            }
            else {
                remainingTimeout =
                    std::chrono::duration_cast<std::chrono::milliseconds>(
                        deadline.value() - now);
            }
        }

        // Only wait for the process if we haven't timed out
        if (!timedOut) {
            result.d_exitCode =
                waitPidOrSignalWithTimeout(pid, remainingTimeout, &timedOut);
        }
    }
    if (timedOut) {
        // Close any remaining pipes if we timed out
        if (pipeStdOut && FD_ISSET(stdOutPipeFDs[0], &readFDSet)) {
            close(stdOutPipeFDs[0]);
        }
        if (pipeStdErr && FD_ISSET(stdErrPipeFDs[0], &readFDSet)) {
            close(stdErrPipeFDs[0]);
        }
        if (stdinOpen) {
            close(stdinPipeFDs[1]);
        }
        // Kill the process
        kill(pid, SIGKILL);
        result.d_timedOut = true;
        result.d_exitCode = waitPid(pid);
    }

    return result;
}

std::string SystemUtils::getPathToCommand(const std::string &command)
{
    if (command.find('/') != std::string::npos) {
        return command; // `command` is a path, no need to search.
    }

    // Reading $PATH, parsing it, and looking for the binary:
    char *path_pointer = getenv("PATH");
    if (path_pointer == nullptr) {
        BUILDBOX_LOG_ERROR("Could not read $PATH");
        return "";
    }

    std::istringstream iss(path_pointer);

    char separator = ':';
    for (std::string token; std::getline(iss, token, separator);) {
        std::string path = token + "/" + command;

        if (FileUtils::isRegularFile(path.c_str()) &&
            FileUtils::isExecutable(path.c_str())) {
            return path;
        }
    }

    return "";
}

int SystemUtils::waitPid(const pid_t pid)
{
    while (true) {
        const int exit_code = waitPidOrSignal(pid);
        if (exit_code == -EINTR) {
            // The child can still run, keep waiting for it.
            continue;
        }
        else {
            return exit_code;
        }
    }
}

int SystemUtils::waitPidOrSignal(const pid_t pid)
{
    bool timedOut = false;
    return waitPidOrSignalWithTimeout(pid, {}, &timedOut);
}

int SystemUtils::waitPidOrSignalWithTimeout(
    const pid_t pid, const std::optional<std::chrono::milliseconds> &timeout,
    bool *timedOut)
{
    int status = 0;
    int options = 0;
    auto end_time = std::chrono::steady_clock::now();
    const timespec poll_interval = {
        .tv_sec = 0,
        .tv_nsec = 10000000 // 10 ms
    };

    if (timeout.has_value()) {
        options |= WNOHANG; // Don't block waiting for the child to exit.
        end_time += timeout.value();
    }
    while (true) {
        const pid_t child_pid = waitpid(pid, &status, options);
        if (child_pid == -1) { // `waitpid()` failed
            const auto waitpid_error = errno;

            if (waitpid_error == EINTR) {
                // Signal caught before child terminated.
                return -EINTR;
            }

            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(std::system_error, errno,
                                                  std::system_category,
                                                  "Error in waitpid()");
        }
        if (child_pid == pid) {
            break;
        }
        if (std::chrono::steady_clock::now() >= end_time) {
            // Timeout reached.
            *timedOut = true;
            return -ETIMEDOUT;
        }
        // Child process is still running, wait a bit before checking again.
        const int ret = nanosleep(&poll_interval, nullptr);
        if (ret == -1) {
            if (errno == EINTR) {
                // Interrupted by a signal
                return -EINTR;
            }
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(std::system_error, errno,
                                                  std::system_category,
                                                  "Error in nanosleep()");
        }
    }

    if (WIFEXITED(status)) {
        return WEXITSTATUS(status);
    }

    if (WIFSIGNALED(status)) {
        constexpr int SIGNAL_OFFSET = 128;
        return SIGNAL_OFFSET + WTERMSIG(status);
        // Exit code as returned by Bash.
        // (https://gnu.org/software/bash/manual/html_node/Exit-Status.html)
    }

    /* According to the documentation for `waitpid(2)` we should never get
     * here:
     *
     * "If the information pointed to by stat_loc was stored by a call to
     * waitpid() that did not specify the WUNTRACED  or
     * CONTINUED flags, or by a call to the wait() function,
     * exactly one of the macros WIFEXITED(*stat_loc) and
     * WIFSIGNALED(*stat_loc) shall evaluate to a non-zero value."
     *
     * (https://pubs.opengroup.org/onlinepubs/009695399/functions/wait.html)
     */
    BUILDBOXCOMMON_THROW_EXCEPTION(
        std::runtime_error,
        "`waitpid()` returned an unexpected status: " << status);
}

std::string SystemUtils::getCurrentWorkingDirectory()
{
    constexpr unsigned int BUFFER_SIZE = 1024;
    unsigned int bufferSize = BUFFER_SIZE;
    while (true) {
        std::unique_ptr<char[]> buffer(new char[bufferSize]);
        char *cwd = getcwd(buffer.get(), bufferSize);

        if (cwd != nullptr) {
            return {cwd};
        }
        else if (errno == ERANGE) {
            bufferSize *= 2;
        }
        else {
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error, "current working directory not found");
        }
    }
}

void SystemUtils::redirectStandardOutputToFile(const int standardOutputFd,
                                               const std::string &path)
{
    if (standardOutputFd != STDOUT_FILENO &&
        standardOutputFd != STDERR_FILENO) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::invalid_argument,
            "File descriptor it not `STDOUT_FILENO` or `STDERR_FILENO`.");
    }

    const FileDescriptor fileFd(
        open(path.c_str(), O_CREAT | O_TRUNC | O_APPEND | O_WRONLY | O_CLOEXEC,
             S_IRUSR | S_IWUSR));

    if (fileFd.get() == -1 || dup2(fileFd.get(), standardOutputFd) == -1) {
        const auto outputName =
            (standardOutputFd == STDOUT_FILENO) ? "stdout" : "stderr";

        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category,
            "Error redirecting " << outputName << " to \"" << path << "\"");
    }
}

ProcessCredentials SystemUtils::getProcessCredentials()
{
    ProcessCredentials creds{};
    creds.uid = geteuid();
    creds.gid = getegid();
    return creds;
}

std::unordered_set<gid_t> SystemUtils::getGids()
{
    std::unordered_set<gid_t> gids;

    // effective gid
    // It is unspecified whether the effective group ID of the calling process
    // of `getgroups` is included in the returned list.
    gids.emplace(getegid());

    // supplement gids
    // get total size
    int count = getgroups(0, nullptr);
    if (count < 0) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "Failed to get the size of group access list "
                                "for current process, error:"
                                    << std::strerror(errno));
    }
    std::vector<gid_t> gidBuffer(static_cast<size_t>(count));

    if (getgroups(count, gidBuffer.data()) == -1) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Failed to get the group access list for current process, error:"
                << std::strerror(errno));
    }

    gids.insert(gidBuffer.begin(), gidBuffer.end());
    return gids;
}

uid_t SystemUtils::getUidFromName(const std::string &name)
{
    struct passwd *pw = getpwnam(name.c_str());
    if (pw == nullptr) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "Cannot retrieve passwd struct from username "
                                    << name
                                    << ", error: " << std::strerror(errno));
    }
    return pw->pw_uid;
}

gid_t SystemUtils::getGidFromName(const std::string &name)
{
    struct group *grp = getgrnam(name.c_str());
    if (grp == nullptr) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "Cannot retrieve group struct from group name "
                                    << name
                                    << ", error: " << std::strerror(errno));
    }
    return grp->gr_gid;
}

const PlatformInfo &PlatformInfo::get()
{
    // thread-safe after C++11
    static const PlatformInfo instance;
    return instance;
}

PlatformInfo::PlatformInfo()
{
    struct utsname uts{};
    // Solaris returns 1 on success, Linux returns 0
    if (uname(&uts) < 0) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(std::system_error, errno,
                                              std::system_category,
                                              "Error calling `uname()`");
    }

    d_osName = uts.sysname;
    const auto releaseParts =
        StringUtils::split(std::string(uts.release), ".");
    if (releaseParts.size() >= 1) {
        d_majorVer = std::stoi(releaseParts[0]);
    }
    if (releaseParts.size() >= 2) {
        d_minorVer = std::stoi(releaseParts[1]);
    }

#if defined(__linux__)
    // Linux 4.5 introduced the copy_file_range syscall
    // There are older kernel releases that have buggy backports
    // so checking ENOSYS is not enough

    // NOLINTBEGIN (cppcoreguidelines-avoid-magic-numbers)
    d_enableCopyFileRange =
        d_majorVer > 4 || (d_majorVer == 4 && d_minorVer >= 5);
    // NOLINTEND (cppcoreguidelines-avoid-magic-numbers)

    // Available in Linux 5.6+ upstream but also available via backports.
    int fd = FileUtils::openat2(AT_FDCWD, ".", O_RDONLY, 0,
                                FileUtils::OPENAT2_RESOLVE_IN_ROOT);
    if (fd >= 0) {
        close(fd);
        d_enableOpenat2 = true;
    }
#endif
}

bool PlatformInfo::enableCopyFileRange() const
{
    return d_enableCopyFileRange;
}

bool PlatformInfo::enableOpenat2() const { return d_enableOpenat2; }

int PlatformInfo::majorVersion() const { return d_majorVer; }

int PlatformInfo::minorVersion() const { return d_minorVer; }
