// 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 <magenta/syscalls.h>
#include <magenta/processargs.h>
#include <mxio/util.h>
#include <magenta/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";
    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 mx_status_t launch(const char* filename,
                          int argc,
                          char** argv,
                          int input,
                          mx_handle_t* process) {
  launchpad_t* lp = NULL;
  mx_handle_t job_to_child = MX_HANDLE_INVALID;
  mx_handle_t job = mxio_get_startup_handle(MX_HND_INFO(MX_HND_TYPE_JOB, 0));
  if (job != MX_HANDLE_INVALID)
    mx_handle_duplicate(job, MX_RIGHT_SAME_RIGHTS, &job_to_child);
  mx_status_t status = launchpad_create(job_to_child, filename, &lp);
  if (status != NO_ERROR)
    return status;
  ScopedGuard guard([&]() { launchpad_destroy(lp); });
  status = launchpad_elf_load(lp, launchpad_vmo_from_file(filename));
  if (status != NO_ERROR)
    return status;
  status = launchpad_load_vdso(lp, MX_HANDLE_INVALID);
  if (status != NO_ERROR)
    return status;
  status = launchpad_add_vdso_vmo(lp);
  if (status != NO_ERROR)
    return status;
  status = launchpad_arguments(lp, argc, argv);
  if (status != NO_ERROR)
    return status;
  status = launchpad_environ(lp, (const char**)environ);
  if (status != NO_ERROR)
    return status;
  status = launchpad_clone_mxio_root(lp);
  if (status != NO_ERROR)
    return status;
  status = launchpad_clone_mxio_cwd(lp);
  if (status != NO_ERROR)
    return status;
  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 (status == NO_ERROR && input != -1) {
    status = launchpad_transfer_fd(lp, input, STDIN_FILENO);
  }
  if (status == NO_ERROR) {
    mx_handle_t result = launchpad_start(lp);
    if (result > 0)
      *process = result;
    else
      status = result;
  }
  return status;
}

int SpecFixture::LaunchProcess(char** args,
                               int args_length,
                               const char* input_filename) {
  mx_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);
    }
  });
  mx_status_t status = launch(args[0], args_length, args, input, &handle);
  if (status != NO_ERROR) {
    LOG(ERROR) << "Failed to launch " << args[0] << ":" << status;
    return handle;
  }
  status =
      mx_handle_wait_one(handle, MX_PROCESS_SIGNALED, MX_TIME_INFINITE, NULL);

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

  mx_info_process_t proc_info;
  ssize_t info_status = mx_object_get_info(handle, MX_INFO_PROCESS, &proc_info,
                                           sizeof(proc_info), NULL, NULL);
  mx_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