#include <iomanip>
#include <iterator>
#include <mutex>
#include <stdexcept>
#ifdef DAGGY_ENABLE_SLURM
#include <slurm/slurm.h>
#include <string.h>
#include <sys/resource.h>
#include <sys/stat.h>
#include <sys/types.h>

#include <csignal>
#include <cstdlib>
#include <daggy/Utilities.hpp>
#include <daggy/executors/task/SlurmTaskExecutor.hpp>
#include <filesystem>
#include <fstream>
#include <random>

namespace fs = std::filesystem;

namespace daggy::executors::task {
  std::string getUniqueTag(size_t nChars = 6)
  {
    std::string result(nChars, '\0');
    static std::random_device dev;
    static std::mt19937 rng(dev());

    std::uniform_int_distribution<int> dist(0, 61);

    const char *v =
        "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";

    for (size_t i = 0; i < nChars; i++) {
      result[i] = v[dist(rng)];
    }
    return result;
  }

  void readAndClean(const fs::path &fn, std::string &dest)
  {
    if (!fs::exists(fn))
      return;

    std::ifstream ifh;
    ifh.open(fn);
    std::string contents(std::istreambuf_iterator<char>{ifh}, {});
    ifh.close();
    fs::remove_all(fn);

    dest.swap(contents);
  }

  SlurmTaskExecutor::SlurmTaskExecutor()
    : running_(true)
    , monitorWorker_(&SlurmTaskExecutor::monitor, this)
  {
    std::string priority =
        "SLURM_PRIO_PROCESS=" + std::to_string(getpriority(PRIO_PROCESS, 0));
    std::string submitDir = "SLURM_SUBMIT_DIR=" + fs::current_path().string();

    const size_t MAX_HOSTNAME_LENGTH = 50;
    std::string submitHost(MAX_HOSTNAME_LENGTH, '\0');
    gethostname(submitHost.data(), MAX_HOSTNAME_LENGTH);
    submitHost = "SLURM_SUBMIT_HOST=" + submitHost;
    submitHost.resize(submitHost.find('\0'));

    uint32_t mask = umask(0);
    umask(mask);  // Restore the old mask

    std::stringstream ss;
    ss << "SLURM_UMASK=0" << uint32_t{((mask >> 6) & 07)}
       << uint32_t{((mask >> 3) & 07)} << uint32_t{(mask & 07)};

    // Set some environment variables
    putenv(const_cast<char *>(priority.c_str()));
    putenv(const_cast<char *>(submitDir.c_str()));
    putenv(const_cast<char *>(submitHost.c_str()));
    putenv(const_cast<char *>(ss.str().c_str()));
  }

  SlurmTaskExecutor::~SlurmTaskExecutor()
  {
    running_ = false;
    monitorWorker_.join();

    // Resolve the remaining futures
    std::lock_guard<std::mutex> lock(promiseGuard_);
    for (auto &[jobID, job] : runningJobs_) {
      job.fut->set(AttemptRecord{.rc = -1, .executorLog = "executor killed"});
    }
    runningJobs_.clear();
  }

  std::string SlurmTaskExecutor::description() const
  {
    return "SlurmTaskExecutor";
  }

  // Validates the job to ensure that all required values are set and are of
  // the right type,
  bool SlurmTaskExecutor::validateTaskParameters(const ConfigValues &job)
  {
    const std::unordered_set<std::string> requiredFields{
        "minCPUs",          "minMemoryMB", "minTmpDiskMB", "priority",
        "timeLimitSeconds", "userID",      "workDir",      "tmpDir"};

    for (const auto &requiredField : requiredFields) {
      if (job.count(requiredField) == 0) {
        throw std::runtime_error("Missing field " + requiredField);
      }
    }

    // Require command or commandString
    if (job.count("command") + job.count("commandString") == 0)
      throw std::runtime_error(
          "Either command or commandString must be specified");

    if (job.count("environment")) {
      if (!std::holds_alternative<Command>(job.at("environment")))
        throw std::runtime_error(R"(environment must be an array of strings)");
    }

    return true;
  }

  std::vector<ConfigValues> SlurmTaskExecutor::expandTaskParameters(
      const ConfigValues &job, const ConfigValues &expansionValues)
  {
    std::vector<ConfigValues> newValues;

    auto command =
        (job.count("command") == 0 ? Command{}
                                   : std::get<Command>(job.at("command")));

    auto environment = (job.count("environment") == 0
                            ? Command{}
                            : std::get<Command>(job.at("environment")));

    Command both(command);
    std::copy(environment.begin(), environment.end(), std::back_inserter(both));

    for (const auto &parts : interpolateValues(both, expansionValues)) {
      ConfigValues newCommand{job};
      newCommand["command"] =
          Command(parts.begin(), parts.begin() + command.size());
      newCommand["environment"] =
          Command(parts.begin() + command.size(), parts.end());
      newValues.emplace_back(newCommand);
    }

    return newValues;
  }

  TaskFuture SlurmTaskExecutor::execute(DAGRunID runID,
                                        const std::string &taskName,
                                        const Task &task)
  {
    std::stringstream executorLog;

    const auto &job           = task.job;
    const auto uniqueTaskName = taskName + "_" + getUniqueTag(6);

    fs::path tmpDir        = std::get<std::string>(job.at("tmpDir"));
    std::string stdoutFile = (tmpDir / (uniqueTaskName + ".stdout")).string();
    std::string stderrFile = (tmpDir / (uniqueTaskName + ".stderr")).string();
    std::string workDir    = std::get<std::string>(job.at("workDir"));

    // Convert command to argc / argv
    std::vector<char *> argv{nullptr};
    // Populate the command
    Command command;
    if (task.job.count("commandString")) {
      std::stringstream ss;
      ss << std::get<std::string>(task.job.at("commandString"));
      std::string tok;
      while (ss >> std::quoted(tok)) {
        command.push_back(tok);
      }
    }
    else {
      const auto cmd = std::get<Command>(task.job.at("command"));
      std::copy(cmd.begin(), cmd.end(), std::back_inserter(command));
    }
    std::transform(
        command.begin(), command.end(), std::back_inserter(argv),
        [](const std::string &s) { return const_cast<char *>(s.c_str()); });
    argv.push_back(nullptr);

    std::vector<std::string> env{""};
    std::vector<char *> envp;
    auto it = task.job.find("environment");
    if (it != task.job.end()) {
      const auto environment = std::get<Command>(task.job.at("environment"));
      std::copy(environment.begin(), environment.end(),
                std::back_inserter(env));
    }
    std::transform(
        env.begin(), env.end(), std::back_inserter(envp),
        [](const std::string &s) { return const_cast<char *>(s.c_str()); });

    char script[]    = "#!/bin/bash\n$@\n";
    char stdinFile[] = "/dev/null";

    // taken from slurm
    int error_code;
    job_desc_msg_t jd;
    submit_response_msg_t *resp_msg;

    slurm_init_job_desc_msg(&jd);
    jd.contiguous = 1;
    jd.name       = const_cast<char *>(taskName.c_str());
    jd.min_cpus   = std::stoi(std::get<std::string>(job.at("minCPUs")));

    jd.pn_min_memory = std::stoi(std::get<std::string>(job.at("minMemoryMB")));
    jd.pn_min_tmp_disk =
        std::stoi(std::get<std::string>(job.at("minTmpDiskMB")));
    jd.priority = std::stoi(std::get<std::string>(job.at("priority")));
    jd.shared   = 0;
    jd.time_limit =
        std::stoi(std::get<std::string>(job.at("timeLimitSeconds")));
    jd.min_nodes = 1;
    jd.user_id   = std::stoi(std::get<std::string>(job.at("userID")));
    jd.argv      = argv.data();
    jd.argc      = argv.size();
    // TODO figure out the script to run
    jd.script   = script;
    jd.std_in   = stdinFile;
    jd.std_err  = const_cast<char *>(stderrFile.c_str());
    jd.std_out  = const_cast<char *>(stdoutFile.c_str());
    jd.work_dir = const_cast<char *>(workDir.c_str());

    // jd.env_size    = 1;
    // jd.environment = env;
    jd.env_size    = envp.size();
    jd.environment = envp.data();

    error_code = slurm_submit_batch_job(&jd, &resp_msg);
    if (error_code) {
      std::stringstream ss;
      ss << "Unable to submit slurm job: " << slurm_strerror(error_code);
      throw std::runtime_error(ss.str());
    }

    uint32_t jobID = resp_msg->job_id;
    executorLog << "Job " << resp_msg->job_submit_user_msg << '\n';
    slurm_free_submit_response_response_msg(resp_msg);

    std::lock_guard<std::mutex> lock(promiseGuard_);
    Job newJob{.fut        = std::make_shared<Future<AttemptRecord>>(),
               .stdoutFile = stdoutFile,
               .stderrFile = stderrFile,
               .runID      = runID,
               .taskName   = taskName};
    auto fut = newJob.fut;
    runningJobs_.emplace(jobID, std::move(newJob));

    return fut;
  }

  bool SlurmTaskExecutor::stop(DAGRunID runID, const std::string &taskName)
  {
    // Hopefully this isn't a common thing, so just scrap the current jobs and
    // kill them
    size_t jobID = 0;
    {
      std::lock_guard<std::mutex> lock(promiseGuard_);
      for (const auto &[k, v] : runningJobs_) {
        if (v.runID == runID and v.taskName == taskName) {
          jobID = k;
          break;
        }
      }
      if (jobID == 0)
        return true;
    }

    // Send the kill message to slurm
    slurm_kill_job(jobID, SIGKILL, KILL_HURRY);
    return true;
  }

  void SlurmTaskExecutor::monitor()
  {
    std::unordered_set<size_t> resolvedJobs;
    while (running_) {
      {
        std::lock_guard<std::mutex> lock(promiseGuard_);
        for (auto &[jobID, job] : runningJobs_) {
          job_info_msg_t *jobStatus;
          int error_code =
              slurm_load_job(&jobStatus, jobID, SHOW_ALL | SHOW_DETAIL);
          if (error_code != SLURM_SUCCESS)
            continue;

          uint32_t idx = jobStatus->record_count;
          if (idx == 0)
            continue;
          idx--;
          const slurm_job_info_t &jobInfo = jobStatus->job_array[idx];
          AttemptRecord record;
          switch (jobInfo.job_state) {
            case JOB_PENDING:
            case JOB_SUSPENDED:
            case JOB_RUNNING:
              continue;
              // Job has finished
            case JOB_COMPLETE: /* completed execution successfully */
              record.rc = jobInfo.exit_code;
              break;
            case JOB_FAILED: /* completed execution unsuccessfully */
              record.rc          = jobInfo.exit_code;
              record.executorLog = "Script errored.\n";
              break;
            case JOB_CANCELLED:        /* cancelled by user */
              record.rc          = 9;  // matches SIGKILL
              record.executorLog = "Job cancelled by user.\n";
              break;
            case JOB_TIMEOUT: /* terminated on reaching time limit */
              record.rc          = jobInfo.exit_code;
              record.executorLog = "Job exceeded time limit.\n";
              break;
            case JOB_NODE_FAIL: /* terminated on node failure */
              record.rc          = jobInfo.exit_code;
              record.executorLog = "Node failed during execution\n";
              break;
            case JOB_PREEMPTED: /* terminated due to preemption */
              record.rc          = jobInfo.exit_code;
              record.executorLog = "Job terminated due to pre-emption.\n";
              break;
            case JOB_BOOT_FAIL: /* terminated due to node boot failure */
              record.rc = jobInfo.exit_code;
              record.executorLog =
                  "Job failed to run due to failure of compute node to "
                  "boot.\n";
              break;
            case JOB_DEADLINE: /* terminated on deadline */
              record.rc          = jobInfo.exit_code;
              record.executorLog = "Job terminated due to deadline.\n";
              break;
            case JOB_OOM: /* experienced out of memory error */
              record.rc          = jobInfo.exit_code;
              record.executorLog = "Job terminated due to out-of-memory.\n";
              break;
          }
          slurm_free_job_info_msg(jobStatus);

          readAndClean(job.stdoutFile, record.outputLog);
          readAndClean(job.stderrFile, record.errorLog);

          job.fut->set(std::move(record));
          resolvedJobs.insert(jobID);
        }

        for (const auto &jobID : resolvedJobs) {
          runningJobs_.extract(jobID);
        }
      }

      std::this_thread::sleep_for(std::chrono::seconds(1));
    }
  }
}  // namespace daggy::executors::task
#endif