// Copyright 2018 Google Inc. All Rights Reserved.
//
// 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 "subprocess.h"

#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <launchpad/launchpad.h>
#include <poll.h>
#include <unistd.h>
#include <stdio.h>
#include <string.h>
#include <zircon/errors.h>
#include <zircon/process.h>
#include <zircon/syscalls.h>
#include <zircon/types.h>

extern char** environ;

#include "util.h"

Subprocess::Subprocess(bool use_console) : fd_(-1), handle_(ZX_HANDLE_INVALID),
                                           use_console_(use_console) {
}

Subprocess::~Subprocess() {
  if (fd_ >= 0)
    close(fd_);
  // Reap child if forgotten.
  if (handle_ != ZX_HANDLE_INVALID)
    Finish();
}

bool Subprocess::Start(SubprocessSet* set, const string& command) {
  printf("executing %s\n", command.c_str());

  int output_pipe[2];
  if (pipe(output_pipe) < 0)
    Fatal("pipe: %s", strerror(errno));
  fd_ = output_pipe[0];

  launchpad_t* lp = NULL;
  launchpad_create(ZX_HANDLE_INVALID, "sh", &lp);
  const char* args[] = { "/boot/bin/sh", "-c", command.c_str() };
  launchpad_load_from_file(lp, args[0]);
  launchpad_set_args(lp, 3, args);
  launchpad_clone(lp, LP_CLONE_ALL & (~LP_CLONE_FDIO_STDIO));

  if (!use_console_) {
    launchpad_clone_fd(lp, output_pipe[1], 1);
    launchpad_clone_fd(lp, output_pipe[1], 2);
  }

  const char *errmsg = NULL;
  if (launchpad_go(lp, &handle_, &errmsg) != ZX_OK)
    Fatal(errmsg);
  launchpad_destroy(lp);

  close(output_pipe[1]);
  return true;
}

void Subprocess::OnPipeReady() {
  char buf[4 << 10];
  ssize_t len = read(fd_, buf, sizeof(buf));
  if (len > 0) {
    buf_.append(buf, len);
  } else {
    if (len < 0)
      Fatal("read: %s", strerror(errno));
    close(fd_);
    fd_ = -1;
  }
}

ExitStatus Subprocess::Finish() {
  assert(handle_ != ZX_HANDLE_INVALID);
  zx_status_t status = zx_object_wait_one(handle_, ZX_PROCESS_TERMINATED, ZX_TIME_INFINITE, nullptr);
  if (status != ZX_OK)
    Fatal("zx_object_wait_one failed");

  zx_info_process_t info;
  status = zx_object_get_info(handle_, ZX_INFO_PROCESS, &info, sizeof(info), nullptr, nullptr);
  if (status != ZX_OK)
    Fatal("zx_object_get_info failed to get process status");

  zx_handle_close(handle_);
  handle_ = ZX_HANDLE_INVALID;

  if (info.exited) {
    if (info.return_code == 0)
      return ExitSuccess;
  }
  return ExitFailure;
}

bool Subprocess::Done() const {
  return fd_ == -1;
}

const string& Subprocess::GetOutput() const {
  return buf_;
}

SubprocessSet::SubprocessSet() {
}

SubprocessSet::~SubprocessSet() {
  Clear();
}

Subprocess *SubprocessSet::Add(const string& command, bool use_console) {
  Subprocess *subprocess = new Subprocess(use_console);
  if (!subprocess->Start(this, command)) {
    delete subprocess;
    return 0;
  }
  running_.push_back(subprocess);
  return subprocess;
}

bool SubprocessSet::DoWork() {
  vector<pollfd> fds;
  nfds_t nfds = 0;

  for (vector<Subprocess*>::iterator i = running_.begin();
       i != running_.end(); ++i) {
    int fd = (*i)->fd_;
    if (fd < 0)
      continue;
    pollfd pfd = { fd, POLLIN | POLLPRI, 0 };
    fds.push_back(pfd);
    ++nfds;
  }

  int ret = ppoll(&fds.front(), nfds, NULL, NULL);
  if (ret == -1) {
    perror("ninja: ppoll");
    return false;
  }

  nfds_t cur_nfd = 0;
  for (vector<Subprocess*>::iterator i = running_.begin();
       i != running_.end(); ) {
    int fd = (*i)->fd_;
    if (fd < 0)
      continue;
    assert(fd == fds[cur_nfd].fd);
    if (fds[cur_nfd++].revents) {
      (*i)->OnPipeReady();
      if ((*i)->Done()) {
        finished_.push(*i);
        i = running_.erase(i);
        continue;
      }
    }
    ++i;
  }

  return false;
}

Subprocess* SubprocessSet::NextFinished() {
  if (finished_.empty())
    return NULL;
  Subprocess* subproc = finished_.front();
  finished_.pop();
  return subproc;
}

void SubprocessSet::Clear() {
  for (vector<Subprocess*>::iterator i = running_.begin();
       i != running_.end(); ++i)
    // Since the foreground process is in our process group, it will receive
    // the interruption signal (i.e. SIGINT or SIGTERM) at the same time as us.
    if (!(*i)->use_console_)
      zx_task_kill((*i)->handle_);
  for (vector<Subprocess*>::iterator i = running_.begin();
       i != running_.end(); ++i)
    delete *i;
  running_.clear();
}