// Copyright 2021 The Tint Authors.
//
// 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 "src/utils/io/command.h"

#include <sys/poll.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <unistd.h>
#include <sstream>
#include <vector>

namespace tint {
namespace utils {

namespace {

/// File is a simple wrapper around a POSIX file descriptor
class File {
  constexpr static const int kClosed = -1;

 public:
  /// Constructor
  File() : handle_(kClosed) {}

  /// Constructor
  explicit File(int handle) : handle_(handle) {}

  /// Destructor
  ~File() { Close(); }

  /// Move assignment operator
  File& operator=(File&& rhs) {
    Close();
    handle_ = rhs.handle_;
    rhs.handle_ = kClosed;
    return *this;
  }

  /// Closes the file (if it wasn't already closed)
  void Close() {
    if (handle_ != kClosed) {
      close(handle_);
    }
    handle_ = kClosed;
  }

  /// @returns the file handle
  operator int() { return handle_; }

  /// @returns true if the file is not closed
  operator bool() { return handle_ != kClosed; }

 private:
  File(const File&) = delete;
  File& operator=(const File&) = delete;

  int handle_ = kClosed;
};

/// Pipe is a simple wrapper around a POSIX pipe() function
class Pipe {
 public:
  /// Constructs the pipe
  Pipe() {
    int pipes[2] = {};
    if (pipe(pipes) == 0) {
      read = File(pipes[0]);
      write = File(pipes[1]);
    }
  }

  /// Closes both the read and write files (if they're not already closed)
  void Close() {
    read.Close();
    write.Close();
  }

  /// @returns true if the pipe has an open read or write file
  operator bool() { return read || write; }

  /// The reader end of the pipe
  File read;

  /// The writer end of the pipe
  File write;
};

bool ExecutableExists(const std::string& path) {
  struct stat s {};
  if (stat(path.c_str(), &s) != 0) {
    return false;
  }
  return s.st_mode & S_IXUSR;
}

std::string FindExecutable(const std::string& name) {
  if (ExecutableExists(name)) {
    return name;
  }
  if (name.find("/") == std::string::npos) {
    auto* path_env = getenv("PATH");
    if (!path_env) {
      return "";
    }
    std::istringstream path{path_env};
    std::string dir;
    while (getline(path, dir, ':')) {
      auto test = dir + "/" + name;
      if (ExecutableExists(test)) {
        return test;
      }
    }
  }
  return "";
}

}  // namespace

Command::Command(const std::string& path) : path_(path) {}

Command Command::LookPath(const std::string& executable) {
  return Command(FindExecutable(executable));
}

bool Command::Found() const {
  return ExecutableExists(path_);
}

Command::Output Command::Exec(
    std::initializer_list<std::string> arguments) const {
  if (!Found()) {
    Output out;
    out.err = "Executable not found";
    return out;
  }

  // Pipes used for piping std[in,out,err] to / from the target process.
  Pipe stdin_pipe;
  Pipe stdout_pipe;
  Pipe stderr_pipe;

  if (!stdin_pipe || !stdout_pipe || !stderr_pipe) {
    Output output;
    output.err = "Command::Exec(): Failed to create pipes";
    return output;
  }

  // execv() and friends replace the current process image with the target
  // process image. To keep process that called this function going, we need to
  // fork() this process into a child and parent process.
  //
  // The child process is responsible for hooking up the pipes to
  // std[in,out,err]_pipes to STD[IN,OUT,ERR]_FILENO and then calling execv() to
  // run the target command.
  //
  // The parent process is responsible for feeding any input to the stdin_pipe
  // and collectting output from the std[out,err]_pipes.

  int child_id = fork();
  if (child_id < 0) {
    Output output;
    output.err = "Command::Exec(): fork() failed";
    return output;
  }

  if (child_id > 0) {
    // fork() - parent

    // Close the stdout and stderr writer pipes.
    // This is required for getting poll() POLLHUP events.
    stdout_pipe.write.Close();
    stderr_pipe.write.Close();

    // Write the input to the child process
    if (!input_.empty()) {
      ssize_t n = write(stdin_pipe.write, input_.data(), input_.size());
      if (n != static_cast<ssize_t>(input_.size())) {
        Output output;
        output.err = "Command::Exec(): write() for stdin failed";
        return output;
      }
    }
    stdin_pipe.write.Close();

    // Accumulate the stdout and stderr output from the child process
    pollfd poll_fds[2];
    poll_fds[0].fd = stdout_pipe.read;
    poll_fds[0].events = POLLIN;
    poll_fds[1].fd = stderr_pipe.read;
    poll_fds[1].events = POLLIN;

    Output output;
    bool stdout_open = true;
    bool stderr_open = true;
    while (stdout_open || stderr_open) {
      if (poll(poll_fds, 2, -1) < 0) {
        break;
      }
      char buf[256];
      if (poll_fds[0].revents & POLLIN) {
        auto n = read(stdout_pipe.read, buf, sizeof(buf));
        if (n > 0) {
          output.out += std::string(buf, buf + n);
        }
      }
      if (poll_fds[0].revents & POLLHUP) {
        stdout_open = false;
      }
      if (poll_fds[1].revents & POLLIN) {
        auto n = read(stderr_pipe.read, buf, sizeof(buf));
        if (n > 0) {
          output.err += std::string(buf, buf + n);
        }
      }
      if (poll_fds[1].revents & POLLHUP) {
        stderr_open = false;
      }
    }

    // Get the resulting error code
    waitpid(child_id, &output.error_code, 0);

    return output;
  } else {
    // fork() - child

    // Redirect the stdin, stdout, stderr pipes for the execv process
    if ((dup2(stdin_pipe.read, STDIN_FILENO) == -1) ||
        (dup2(stdout_pipe.write, STDOUT_FILENO) == -1) ||
        (dup2(stderr_pipe.write, STDERR_FILENO) == -1)) {
      fprintf(stderr, "Command::Exec(): Failed to redirect pipes");
      exit(errno);
    }

    // Close the pipes, once redirected above, we're now done with them.
    stdin_pipe.Close();
    stdout_pipe.Close();
    stderr_pipe.Close();

    // Run target executable
    std::vector<const char*> args;
    args.emplace_back(path_.c_str());
    for (auto& arg : arguments) {
      args.emplace_back(arg.c_str());
    }
    args.emplace_back(nullptr);
    auto res = execv(path_.c_str(), const_cast<char* const*>(args.data()));
    exit(res);
  }
}

}  // namespace utils
}  // namespace tint