// Copyright 2016 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "spec_fixture.h"
#include "spec_utils.h"

#include <benchmark/benchmark.h>
#include <fcntl.h>
#include <fstream>
#include <glog/logging.h>
#include <sys/stat.h>
#include <unistd.h>

#ifdef __unix__
#include <spawn.h>
#include <sys/wait.h>
#endif

#ifdef __Fuchsia__
#include <launchpad/launchpad.h>
#include <launchpad/vmo.h>
#include <zircon/syscalls.h>
#include <zircon/process.h>
#include <zircon/processargs.h>
#include <zircon/syscalls/object.h>
#endif

extern std::string executableDir;
extern std::string tmpDir;
extern std::ofstream result;

namespace {
class ScopedGuard {
 public:
  explicit ScopedGuard(std::function<void()> f) : f_(f) {}
  ~ScopedGuard() noexcept { f_(); }
  ScopedGuard(const ScopedGuard&) = delete;
  ScopedGuard(ScopedGuard&&) = delete;
  ScopedGuard& operator=(const ScopedGuard&) = delete;
  ScopedGuard& operator=(ScopedGuard&&) = delete;

 private:
  std::function<void()> f_;
};
}  // namespace

void SpecFixture::SetUp(benchmark::State& st) {
  DLOG(INFO) << "Preparing benchmark: " << benchmark_name;
  run_dir_name = tmpDir + "/" + benchmark_name;

  DLOG(INFO) << "Changing dir to: " << executableDir.c_str();
  CHECK_EQ(chdir(executableDir.c_str()), 0)
      << "Not able to change dir to " << executableDir;
  CHECK_EQ(mkdir(run_dir_name.c_str(), 0700), 0)
      << "Not able to create dir: " << run_dir_name
      << ", Error: " << strerror(errno);

  CHECK_EQ(CopyDir(run_dir_name, "data/" + benchmark_name), 0)
      << "Error copying dir: " << run_dir_name
      << ", Error: " << strerror(errno);

  std::string binary = run_dir_name + "/" + benchmark_name;
  CHECK_EQ(CopyFile(binary, benchmark_name), 0)
      << "Error copying binary, Error: " << strerror(errno);

  CHECK_EQ(chmod(binary.c_str(), 0700), 0)
      << "Error changing file permission for binary: " << strerror(errno);
}

void SpecFixture::TearDown(benchmark::State& st) {
  CHECK_EQ(DeleteDir(run_dir_name), 0)
      << "Error Deleting directory " << run_dir_name
      << ", Error: " << strerror(errno);
}

int SpecFixture::RunSpec(const char* args[],
                         int args_length,
                         const char* input_filename) {
  int page_size = getpagesize();
  char** argv = new char*[args_length + 1];
  argv[0] = (char*)(benchmark_name.c_str());
  if (args_length > 0) {
    memcpy(argv + 1, args, args_length * sizeof(char*));
  }

  if (chdir(run_dir_name.c_str()) != 0) {
    LOG(ERROR) << "Not able to change dir to " << run_dir_name
               << ", Error: " << strerror(errno);
    return errno;
  }

  DLOG(INFO) << "Running benchmark";

  // Capture output
  char buffer[page_size];
  int out_p[2];
  int err_p[2];
  int old_stdout = dup(STDOUT_FILENO);
  int old_stderr = dup(STDERR_FILENO);
  if (pipe2(out_p, O_NONBLOCK) != 0) {
    LOG(ERROR) << "Not able create pipe, Error: " << strerror(errno);
    return errno;
  }
  ScopedGuard out_guard([&]() { close(out_p[0]); });

  if (pipe2(err_p, O_NONBLOCK) != 0) {
    close(out_p[1]);
    LOG(ERROR) << "Not able create pipe, Error: " << strerror(errno);
    return errno;
  }
  ScopedGuard err_guard([&]() { close(err_p[0]); });
  dup2(err_p[1], STDERR_FILENO);
  close(err_p[1]);
  dup2(out_p[1], STDOUT_FILENO);
  close(out_p[1]);

  int status = LaunchProcess(argv, args_length + 1, input_filename);

  fflush(stdout);
  fflush(stderr);
  dup2(old_stdout, STDOUT_FILENO);
  dup2(old_stderr, STDERR_FILENO);
  if (chdir(executableDir.c_str()) != 0) {
    LOG(ERROR) << "Not able to change dir to " << executableDir
               << ", Error: " << strerror(errno);
    return errno;
  }
  delete[] argv;
  if (status != 0) {
    // Dump output
    LOG(ERROR) << "Error while running benchmark: " << benchmark_name;
    ssize_t len;
    do {
      len = read(out_p[0], buffer, page_size);
      if (len == -1) {
        LOG(ERROR) << "Error reading from pipe, Error: " << strerror(errno);
        return errno;
      }
      fwrite(buffer, sizeof(char), len, stderr);
    } while (len == page_size);
    fprintf(stderr, "\n");
    do {
      len = read(err_p[0], buffer, page_size);
      if (len == -1) {
        LOG(ERROR) << "Error reading from pipe: " << strerror(errno);
        return errno;
      }
      fwrite(buffer, sizeof(char), len, stderr);
    } while (len == page_size);
    return 1;
  }
  return 0;
}

#ifdef __Fuchsia__
static zx_status_t launch(const char* filename,
                          int argc,
                          char** argv,
                          int input,
                          zx_handle_t* process) {
  launchpad_t* lp = NULL;
  zx_handle_t job_to_child = ZX_HANDLE_INVALID;
  zx_handle_t job = zx_get_startup_handle(PA_HND(PA_JOB_DEFAULT, 0));
  if (job != ZX_HANDLE_INVALID) {
    zx_handle_duplicate(job, ZX_RIGHT_SAME_RIGHTS, &job_to_child);
  }
  launchpad_create(job_to_child, filename, &lp);
  launchpad_load_from_file(lp, filename);
  launchpad_set_args(lp, argc, argv);
  launchpad_set_environ(lp, (const char**)environ);
  launchpad_clone(lp, LP_CLONE_FDIO_NAMESPACE | LP_CLONE_FDIO_CWD);
  launchpad_clone_fd(lp, STDIN_FILENO, STDIN_FILENO);
  launchpad_clone_fd(lp, STDOUT_FILENO, STDOUT_FILENO);
  launchpad_clone_fd(lp, STDERR_FILENO, STDERR_FILENO);
  if(input != -1) {
    launchpad_transfer_fd(lp, input, STDIN_FILENO);
  }
  const char* errmsg;
  zx_status_t status = launchpad_go(lp, process, &errmsg);
  if (status == ZX_OK) {
    LOG(ERROR) << "Failed to launch: " << filename << ", error: " << errmsg
               << ": " << status;
  }
  return status;
}

int SpecFixture::LaunchProcess(char** args,
                               int args_length,
                               const char* input_filename) {
  zx_handle_t handle;
  int input = -1;
  if (input_filename != NULL) {
    if ((input = open(input_filename, O_RDONLY)) < 0) {
      LOG(ERROR) << "cannot open file for reading: " << input_filename
                 << ", error: " << strerror(errno);
      return errno;
    }
  }
  ScopedGuard guard([&]() {
    if (input != -1) {
      close(input);
    }
  });
  zx_status_t status = launch(args[0], args_length, args, input, &handle);
  if (status != ZX_OK) {
    return status;
  }
  status =
      zx_object_wait_one(handle, ZX_PROCESS_TERMINATED, ZX_TIME_INFINITE, NULL);

  if (status != ZX_OK) {
    LOG(ERROR) << "Failed to wait for process exiting " << args[0] << ":"
               << status;
    return status;
  }

  zx_info_process_t proc_info;
  ssize_t info_status = zx_object_get_info(handle, ZX_INFO_PROCESS, &proc_info,
                                           sizeof(proc_info), NULL, NULL);
  zx_handle_close(handle);

  if (info_status < 0) {
    LOG(ERROR) << "Failed to get process return code " << args[0] << ":"
               << info_status;
    return 1;
  }
  return proc_info.return_code;
}
#endif

#ifdef __unix__
int SpecFixture::LaunchProcess(char** args,
                               int args_length,
                               const char* input_filename) {
  run_dir_name = tmpDir + "/" + benchmark_name + "/";
  std::string command = run_dir_name + args[0];
  posix_spawn_file_actions_t* fd_actions = NULL;
  ScopedGuard spawn_action_guard([&]() {
    if (fd_actions != NULL) {
      delete fd_actions;
    }
  });
  int input = -1;
  ScopedGuard fd_guard([&]() {
    if (input != -1) {
      close(input);
    }
  });
  if (input_filename != NULL) {
    if ((input = open(input_filename, O_RDONLY)) < 0) {
      LOG(ERROR) << "cannot open file for reading: " << input_filename
                 << ", error: " << strerror(errno);
      return errno;
    }
    fd_actions = new posix_spawn_file_actions_t;
    if (int ret = posix_spawn_file_actions_init(fd_actions) != 0) {
      LOG(ERROR) << "Error initializing file actions: " << strerror(ret);
      return ret;
    }
    if (int ret = posix_spawn_file_actions_adddup2(fd_actions, input,
                                                   STDIN_FILENO) != 0) {
      LOG(ERROR) << "Error redirecting stdin: " << strerror(ret);
      return ret;
    }
  }

  char** argv = new char*[args_length + 1];
  ScopedGuard guard([&]() { delete[] argv; });
  for (int i = 0; i < args_length; i++) {
    argv[i] = args[i];
  }
  argv[args_length] = NULL;
  pid_t pid;
  int status =
      posix_spawn(&pid, command.c_str(), fd_actions, NULL, argv, environ);
  if (status == 0) {
    if (waitpid(pid, &status, 0) != -1) {
      return status;
    } else {
      LOG(ERROR) << "waitpid error: " << strerror(errno);
    }
  } else {
    LOG(ERROR) << "posix_spawn error: " << strerror(status);
  }
  return -1;
}
#endif