//=========================================================================
//  Copyright (c) Kitware, Inc.
//  All rights reserved.
//  See LICENSE.txt for details.
//
//  This software is distributed WITHOUT ANY WARRANTY; without even
//  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
//  PURPOSE.  See the above copyright notice for more information.
//=========================================================================
#include "smtk/simulation/ace3p/qt/qtNewtJobSubmitter.h"

#include "smtk/newt/qtNewtInterface.h"
#include "smtk/simulation/ace3p/JobsManifest.h"
#include "smtk/simulation/ace3p/Project.h"
#include "smtk/simulation/ace3p/qt/qtProjectRuntime.h"
#include "smtk/simulation/ace3p/utility/AttributeUtils.h"

// SMTK includes
#include "smtk/attribute/Analyses.h"
#include "smtk/attribute/Attribute.h"
#include "smtk/attribute/FileItem.h"
#include "smtk/attribute/GroupItem.h"
#include "smtk/attribute/IntItem.h"
#include "smtk/attribute/ReferenceItem.h"
#include "smtk/attribute/Resource.h"
#include "smtk/attribute/ResourceItem.h"
#include "smtk/attribute/SearchStyle.h"
#include "smtk/attribute/StringItem.h"
#include "smtk/io/Logger.h"
#include "smtk/operation/Operation.h"
#include "smtk/project/Project.h"

#include <QCoreApplication>
#include <QFile>
#include <QFileInfo>
#include <QJsonObject>
#include <QNetworkReply>
#include <QSet>
#include <QString>
#include <QStringList>
#include <QTextStream>
#include <QTimer>

#include "nlohmann/json.hpp"

#include <cassert>
#include <ctime>
#include <regex>
#include <string>
#include <vector>

#ifndef NDEBUG
#include <iostream>
#endif

namespace
{
const smtk::attribute::SearchStyle RECURSIVE = smtk::attribute::SearchStyle::RECURSIVE;

// For tracking instance state
enum State
{
  Idle = 0,
  Running,
  Error,
  Canceled
};

// Macro to call method on instance. Used for setting fields in JobRecordGenerator
#define AsLambdaMacro(instance, method) [&instance](const std::string& s) { instance.method(s); }

// Copy from smtk ValueItem to job record
void toJobRecord(
  const smtk::attribute::ConstAttributePtr att,
  const std::string& itemName,
  std::function<void(const std::string)> setField)
{
  const auto valueItem = att->findAs<smtk::attribute::ValueItem>(
    itemName, smtk::attribute::SearchStyle::RECURSIVE_ACTIVE);
  if (valueItem == nullptr)
  {
    QString msg = QString("Warning: Did not find attribute ValueItem \"%1\"").arg(itemName.c_str());
    qDebug() << msg;
  }
  else if (valueItem->isEnabled() && valueItem->isSet())
  {
    setField(valueItem->valueAsString());
  }
}

// Macro for printing messages to stdout for debug builds only
#ifndef NDEBUG
#define DebugMessageMacro(msg)                                                                     \
  do                                                                                               \
  {                                                                                                \
    std::cout << __FILE__ << ":" << __LINE__ << " " << msg << std::endl;                           \
  } while (0)
#else
#define DebugMessageMacro(msg)
#endif
} // namespace

namespace smtk
{
namespace simulation
{
namespace ace3p
{
//-----------------------------------------------------------------------------
class qtNewtJobSubmitter::Internal
{
public:
  ::newt::qtNewtInterface* m_newt;
  std::shared_ptr<smtk::simulation::ace3p::Project> m_ace3pProject;
  smtk::attribute::ConstAttributePtr m_exportParams;
  smtk::attribute::ConstAttributePtr m_exportResults;

  State m_state = State::Idle;
  QString m_machine;
  QString m_remoteSubfolder;    // relative path from Scratch directory to job directory
  QString m_remoteJobDirectory; // absolute path to remote job directory
  QString m_resultsFolder;      // results subfolder
  std::vector<std::string> m_exportFileList; // files written by export operation
  QStringList m_uploadFileList;              // files to be uploaded (superset of m_exportFileList)
  QString m_localSlurmScript;

  unsigned int m_uploadCount = 0;
  JobRecordGenerator m_jobRecordGenerator;

  Internal();
  ~Internal() = default;

  // Returns true if no errors
  bool init(
    smtk::project::ProjectPtr project,
    smtk::attribute::ConstAttributePtr exportParams,
    smtk::attribute::ConstAttributePtr exportResults);

  // Injects JOB_DIRECTORY into script generated by export operation
  bool generateSlurmScript(QString& errMessage);

  // Find existing jobs that use the same results directory as the input record.
  // The data from the existing job(s) will be overwritten
  QSet<QString> findCoincidentJobs(const nlohmann::json& jobRecord);

  void reset()
  {
    m_ace3pProject.reset();
    m_exportParams.reset();
    m_exportResults.reset();

    m_state = State::Idle;
    m_machine.clear();
    m_remoteSubfolder.clear();
    m_remoteJobDirectory.clear();
    m_resultsFolder.clear();
    m_exportFileList.clear();
    m_uploadFileList.clear();
    m_localSlurmScript.clear();

    m_uploadCount = 0; // for tracking file uploads
    m_jobRecordGenerator.reset();
  }
};

qtNewtJobSubmitter::Internal::Internal()
  : m_newt(newt::qtNewtInterface::instance())
{
}

bool qtNewtJobSubmitter::Internal::init(
  smtk::project::ProjectPtr project,
  smtk::attribute::ConstAttributePtr exportParams,
  smtk::attribute::ConstAttributePtr exportResults)
{
  // Make sure submit item is enabled
  auto nerscItem = exportParams->findAs<smtk::attribute::GroupItem>("NERSCSimulation");
  if ((nerscItem == nullptr) || !nerscItem->isEnabled())
  {
    return false;
  }

  this->reset();
  m_ace3pProject = std::dynamic_pointer_cast<Project>(project);
  m_exportParams = exportParams;
  m_exportResults = exportResults;

  // Get remote subfolder path from export params
  auto subfolderItem = nerscItem->findAs<smtk::attribute::StringItem>("SubFolder", RECURSIVE);
  assert(subfolderItem != nullptr);
  m_remoteSubfolder = QString::fromStdString(subfolderItem->value());

  // Get results directory from export params
  auto resultsFolderItem =
    nerscItem->findAs<smtk::attribute::StringItem>("ResultsDirectory", RECURSIVE);
  if (resultsFolderItem && resultsFolderItem->isEnabled())
  {
    m_resultsFolder = QString::fromStdString(resultsFolderItem->value());
  }

  // Get upload file list from export results
  m_uploadFileList.clear();

  int stageIndex = exportParams->findInt("stage-index")->value();
  std::shared_ptr<Stage> stage = m_ace3pProject->stage(stageIndex);

  auto modelResource = stage->modelResource();
  if (modelResource == nullptr)
  {
    qCritical() << "Unable to export: stage has no model resource";
    return false;
  }

  if (!modelResource->properties().contains<std::string>(Metadata::METADATA_PROPERTY_KEY))
  {
    qCritical() << "Internal Error: model resource missing"
                << Metadata::METADATA_PROPERTY_KEY.c_str() << "property";
    return false;
  }

  std::string ssMeshfile =
    modelResource->properties().at<std::string>(Metadata::METADATA_PROPERTY_KEY);
  QString meshfile = QString::fromStdString(ssMeshfile);

  QFileInfo meshfileInfo(meshfile);
  if (!meshfileInfo.exists())
  {
    qCritical() << "Error: mesh file missing:" << meshfile;
    return false;
  }

  m_uploadFileList.append(meshfile);

  // Begin populating job record
  JobRecordGenerator& jrg(m_jobRecordGenerator); // shorthand

  QString remoteMeshfileName = meshfileInfo.completeBaseName() + ".ncdf";
  jrg.runtimeMeshFileName(remoteMeshfileName.toStdString());

  ::toJobRecord(exportParams, "JobName", AsLambdaMacro(jrg, jobName));
  ::toJobRecord(exportParams, "Machine", AsLambdaMacro(jrg, machine));
  auto record = jrg.get();
  auto jiter = record.find("machine");
  if (jiter == record.end())
  {
    qWarning() << "Error getting machine from export parameters; using cori";
    m_machine = "cori";
  }
  else
  {
    m_machine = QString::fromStdString(jiter->get<std::string>());
  }

  smtk::attribute::ResourcePtr attResource = stage->attributeResource();
  jrg.analysisID(attResource->id().toString());

  smtk::simulation::ace3p::AttributeUtils attUtils;
  auto analysisAtt = attUtils.getAnalysisAtt(attResource);
  if (analysisAtt)
  {
    ::toJobRecord(analysisAtt, "Analysis", AsLambdaMacro(jrg, analysis));

    // Check for ACDTool task - easiest way is to get from categories
    std::set<std::string> categories;
    attResource->analyses().getAnalysisAttributeCategories(analysisAtt, categories);
    if (categories.find("Rf-Postprocess") != categories.end())
    {
      jrg.acdtoolTask("rf postprocess");
    }
    else
    {
      jrg.acdtoolTask(""); // clears default value ACDTool
    }
  }
  else
  {
    qCritical() << "Internal Error: Failed to find analysis item in attribute resource";
    return false;
  }

  ::toJobRecord(exportParams, "JobNotes", AsLambdaMacro(jrg, notes));
  ::toJobRecord(exportResults, "NerscInputFolder", AsLambdaMacro(jrg, runtimeInputFolder));
  ::toJobRecord(exportParams, "ResultsDirectory", AsLambdaMacro(jrg, resultsSubfolder));
  ::toJobRecord(exportParams, "NumberOfNodes", AsLambdaMacro(jrg, nodes));
  ::toJobRecord(exportParams, "NumberOfTasks", AsLambdaMacro(jrg, processes));

  // update this continuously via cumulus
  jrg.elapsedTime(0);
  jrg.submissionTime("0");

  std::vector<std::string> fileItemNames = { "OutputFile", "SlurmScript" };
  for (const auto& name : fileItemNames)
  {
    auto fileItem = exportResults->findFile(name);
    assert(fileItem != nullptr && fileItem->isEnabled());
    QString uploadFilename = QString::fromStdString(fileItem->value());
    if (name == "SlurmScript")
    {
      // Store path for configuring after we get the job directory
      m_localSlurmScript = uploadFilename;
    }
    else
    {
      m_exportFileList.push_back(uploadFilename.toStdString());
      m_uploadFileList.push_back(uploadFilename);
    }
  }

  return true;
}

bool qtNewtJobSubmitter::Internal::generateSlurmScript(QString& errMessage)
{
  if (!m_localSlurmScript.endsWith(".in"))
  {
    errMessage =
      QString("Expected slurm script input %1 to end with .in extension.").arg(m_localSlurmScript);
    return false;
  }

  // Read input file
  QFile inFile(m_localSlurmScript);
  if (!inFile.open(QFile::ReadOnly | QFile::Text))
  {
    errMessage = QString("Unable to open slurm script input file %1.").arg(m_localSlurmScript);
    return false;
  }
  QString inputScript = inFile.readAll();
  inFile.close();

  // Replace $JOB_DIRECTORY
  QString outputScript = inputScript.replace("$JOB_DIRECTORY", m_remoteJobDirectory);

  // Write script file
  m_localSlurmScript.chop(3); // strip the ".in" suffix
  QFile outFile(m_localSlurmScript);
  if (!outFile.open(QFile::Truncate | QFile::WriteOnly))
  {
    errMessage =
      QString("Unable to open slurm script file %1 for writing.").arg(m_localSlurmScript);
    return false;
  }
  QTextStream out(&outFile);
  out << outputScript;
  outFile.close();

  m_exportFileList.push_back(m_localSlurmScript.toStdString());
  m_uploadFileList.push_back(m_localSlurmScript);

  // Update export results item
  smtk::attribute::ConstFileItemPtr fileItem = m_exportResults->findFile("SlurmScript");
  auto* rawFileItem = const_cast<smtk::attribute::FileItem*>(fileItem.get());
  rawFileItem->setValue(m_localSlurmScript.toStdString());

  return true;
}

QSet<QString> qtNewtJobSubmitter::Internal::findCoincidentJobs(const nlohmann::json& jobRecord)
{
  QSet<QString> jobIds; // return value

  // Traverse jobs manifest checking for matching results subfolder and job directory
  std::string otherSubfolder = jobRecord["results_subfolder"].get<std::string>();
  std::string otherJobFolder = jobRecord["runtime_job_folder"].get<std::string>();

  std::shared_ptr<JobsManifest> manifest = m_ace3pProject->jobsManifest();
  for (int i = 0; i < static_cast<int>(manifest->size()); ++i)
  {
    bool found;

    std::string subfolder;
    found = manifest->getField(i, "results_subfolder", subfolder);
    if (!found || (subfolder != otherSubfolder))
    {
      continue;
    }

    // The job folder should generally be the same for all jobs, but the user has the
    // option to change it, so we check this field too.
    std::string jobFolder;
    found = manifest->getField(i, "runtime_job_folder", jobFolder);
    if (!found || (jobFolder != otherJobFolder))
    {
      continue;
    }

    // For ACDTool, the acdtool_task fields must also match
    if (subfolder == "acdtool_results") // shortcut
    {
      std::string otherTask = jobRecord["acdtool_task"].get<std::string>();
      std::string task;
      found = manifest->getField(i, "acdtool_task", task);
      if (!found || (task != otherTask))
      {
        continue;
      }
    }

    // Folder and subfolder both match, so get job id and insert into the return set.
    std::string jobId;
    found = manifest->getField(i, "job_id", jobId);
    jobIds.insert(QString::fromStdString(jobId));
  } // for

  return jobIds;
}

//-----------------------------------------------------------------------------
qtNewtJobSubmitter::qtNewtJobSubmitter(QObject* parent)
  : QObject(parent)
{
  m_internal = new Internal;
}

qtNewtJobSubmitter::~qtNewtJobSubmitter()
{
  delete m_internal;
}

void qtNewtJobSubmitter::submitAnalysisJob(
  smtk::project::ProjectPtr project,
  smtk::attribute::ConstAttributePtr exportParams,
  smtk::attribute::ConstAttributePtr exportResults)
{
  // Sanity check - can only submit one job at a time
  if (m_internal->m_state == State::Running)
  {
    emit this->errorMessage(
      "Cannot call qtNewtJobSubmitter::submitAnalysisJob() when already running.");
    return;
  }

  // Make sure user is logged in
  if (!m_internal->m_newt->isLoggedIn())
  {
    emit this->errorMessage("Cannot submit job because user not logged into NERSC.");
    return;
  }

  m_internal->init(project, exportParams, exportResults);
  m_internal->m_state = State::Running;

  auto* runtime = qtProjectRuntime::instance();
  QString scratchPath = runtime->scratchPath();
  if (scratchPath.isEmpty())
  {
    emit this->progressMessage("Getting Scratch folder location.");
    QNetworkReply* reply = m_internal->m_newt->requestScratchPath(m_internal->m_machine);
    QObject::connect(reply, &QNetworkReply::finished, [this, runtime, reply]() {
      QString path;
      QString errMessage;
      this->m_internal->m_newt->getCommandReply(reply, path, errMessage);
      if (!errMessage.isEmpty())
      {
        emit this->errorMessage(errMessage);
        reply->deleteLater();
        return;
      }

      QString scratchMessage = QString("Scratch location is %1.").arg(path);
      emit this->progressMessage(scratchMessage);
      runtime->setScratchPath(path);

      this->createJobDirectory();
      reply->deleteLater();
    });
  }
  else
  {
    QString message = QString("Using scratch location is %1.").arg(scratchPath);
    emit this->progressMessage(message);
    this->createJobDirectory();
  }
}

QString qtNewtJobSubmitter::getUniqueJobName(
  smtk::project::ProjectPtr project,
  const QString& prefix)
{
  // Find max number used with this prefix
  // Returns prefix.N (e.g, prefix.1 or prefix.2 or ...)
  auto ace3pProject = std::dynamic_pointer_cast<Project>(project);
  auto manifest = ace3pProject->jobsManifest();

  unsigned long maxUsed = 0;
  std::string reString = QString("^%1\\.(\\d+)$").arg(prefix).toStdString();
  std::regex reNum(reString);
  std::smatch matchList;
  std::string jobName;
  for (int i = 0; i < static_cast<int>(manifest->size()); ++i)
  {
    manifest->getField(i, "job_name", jobName);
    std::regex_match(jobName, matchList, reNum);
    if (matchList.size() == 2)
    {
      std::string matchString = matchList[1].str();
      unsigned long n = std::stoul(matchString);
      maxUsed = n > maxUsed ? n : maxUsed;
    }
  }

  unsigned long uniqueNum = maxUsed + 1;
  QString uniqueName = QString("%1.%2").arg(prefix).arg(uniqueNum);
  return uniqueName;
}

void qtNewtJobSubmitter::createJobDirectory()
{
  if (m_internal->m_state != State::Running)
  {
    emit this->progressMessage("Submit-Job sequence was terminated.");
    return;
  }

  QString scratchPath = qtProjectRuntime::instance()->scratchPath();
  m_internal->m_remoteJobDirectory = scratchPath + "/" + m_internal->m_remoteSubfolder;
  m_internal->m_jobRecordGenerator.runtimeJobFolder(m_internal->m_remoteJobDirectory.toStdString());

  // Need to generate the slurm script from the .in file
  QString errMsg;
  if (!m_internal->generateSlurmScript(errMsg))
  {
    emit this->errorMessage(errMsg);
    return;
  }

  QString msg = QString("Setting up job directory %1.").arg(m_internal->m_remoteJobDirectory);
  emit this->progressMessage(msg);

  QString command = QString("/bin/mkdir -p %1").arg(m_internal->m_remoteJobDirectory);
  QNetworkReply* reply = m_internal->m_newt->sendCommand(command, m_internal->m_machine);
  QObject::connect(
    reply, &QNetworkReply::finished, this, &qtNewtJobSubmitter::onCreateJobDirectoryFinished);
}

void qtNewtJobSubmitter::onCreateJobDirectoryFinished()
{
  auto* reply = dynamic_cast<QNetworkReply*>(this->sender());
  QString outputText;
  QString errMessage;
  m_internal->m_newt->getCommandReply(reply, outputText, errMessage);
  reply->deleteLater();
  if (!errMessage.isEmpty())
  {
    emit this->errorMessage(errMessage);
    return;
  }

  // Next step is to make sure results directory is cleared out
  this->clearResultsDirectory();
}

void qtNewtJobSubmitter::clearResultsDirectory()
{
  if (m_internal->m_state != State::Running)
  {
    emit this->progressMessage("Submit-Job sequence was terminated.");
    return;
  }

  // Special case for acdtool_results directory, which is just a placeholder
  if (m_internal->m_resultsFolder == "acdtool_results")
  {
    // Schedule call to next slot
    QTimer::singleShot(0, this, &qtNewtJobSubmitter::uploadJobFiles);
    return;
  }

  QString resultsPath = m_internal->m_remoteJobDirectory + "/" + m_internal->m_resultsFolder;
  QString msg = QString("Clearing results directory %1.").arg(resultsPath);
  emit this->progressMessage(msg);

  QString command = QString("/bin/rm -rf %1").arg(resultsPath);
  QNetworkReply* reply = m_internal->m_newt->sendCommand(command, m_internal->m_machine);
  QObject::connect(
    reply, &QNetworkReply::finished, this, &qtNewtJobSubmitter::onClearResultsDirectoryFinished);
}

void qtNewtJobSubmitter::onClearResultsDirectoryFinished()
{
  auto* reply = dynamic_cast<QNetworkReply*>(this->sender());
  QString outputText;
  QString errMessage;
  m_internal->m_newt->getCommandReply(reply, outputText, errMessage);
  reply->deleteLater();
  if (!errMessage.isEmpty())
  {
    emit this->errorMessage(errMessage);
    return;
  }

  // Next step is to upload job files
  this->uploadJobFiles();
}

void qtNewtJobSubmitter::uploadJobFiles()
{
  if (m_internal->m_state != State::Running)
  {
    emit this->progressMessage("Submit-Job sequence was terminated.");
    return;
  }

  m_internal->m_uploadCount = m_internal->m_uploadFileList.size();
  foreach (QString localPath, m_internal->m_uploadFileList)
  {
    QNetworkReply* reply = m_internal->m_newt->requestFileUpload(
      localPath, m_internal->m_remoteJobDirectory, m_internal->m_machine);
    if (reply == nullptr)
    {
      QString msg = QString("Error loading local file: %1").arg(localPath);
      emit this->errorMessage(msg);
      return;
    }
    QObject::connect(
      reply, &QNetworkReply::finished, this, &qtNewtJobSubmitter::onUploadJobFilesFinished);
    QString msg = QString("Uploading %1.").arg(localPath);
    emit this->progressMessage(msg);
  }
}

void qtNewtJobSubmitter::onUploadJobFilesFinished()
{
  // DebugMessageMacro("onFileUploadFinished(): upload count " << m_internal->m_uploadCount);
  --m_internal->m_uploadCount;

  // If job got canceled or hit an error, don't proceed
  if (m_internal->m_state != State::Running)
  {
    return;
  }

  auto* reply = dynamic_cast<QNetworkReply*>(this->sender());
  QJsonObject j;
  QString errMessage;
  if (!m_internal->m_newt->parseReply(reply, j, errMessage))
  {
    qDebug() << "parseReply error, returned:" << reply->readAll();
    int errCode = static_cast<int>(reply->error());
    QString msg = QString("Upload Error: %1, error code %2").arg(errMessage).arg(errCode);
    emit this->errorMessage(msg);

    m_internal->m_state = State::Error;
    reply->deleteLater();
    return;
  }
  reply->deleteLater();

  // DebugMessageMacro("m_uploadCount " << m_internal->m_uploadCount);
  if (m_internal->m_uploadCount == 0)
  {
    emit this->progressMessage("File uploads complete.");
    this->submitJob();
  }
}

void qtNewtJobSubmitter::submitJob()
{
  if (m_internal->m_state != State::Running)
  {
    emit this->progressMessage("Submit-Job sequence was terminated.");
    return;
  }

  emit this->progressMessage("Submitting job request.");

  QFileInfo fileInfo(m_internal->m_localSlurmScript);
  QString remoteSlurmScript =
    QString("%1/%2").arg(m_internal->m_remoteJobDirectory, fileInfo.fileName());
  QNetworkReply* reply =
    m_internal->m_newt->requestJobSubmit(remoteSlurmScript, m_internal->m_machine);
  QObject::connect(reply, &QNetworkReply::finished, this, &qtNewtJobSubmitter::onSubmitJobFinished);
}

void qtNewtJobSubmitter::onSubmitJobFinished()
{
  auto* reply = dynamic_cast<QNetworkReply*>(this->sender());
  QJsonObject j;
  QString errMessage;
  if (!m_internal->m_newt->parseReply(reply, j, errMessage))
  {
    emit this->errorMessage(errMessage);
    reply->deleteLater();
    return;
  }

  if (j.contains("jobid"))
  {
    JobRecordGenerator& jrg(m_internal->m_jobRecordGenerator); // shorthand

    jrg.status("created");

    QString jobId = j.value("jobid").toString();
    jrg.jobID(jobId.toStdString());

    std::time_t seconds = std::time(nullptr);
    std::stringstream ss;
    ss << seconds;
    jrg.submissionTime(ss.str());

    nlohmann::json jobRec = jrg.get();

    // Check for jobs being overwritten (i.e., existing jobs using the same results folder)
    QSet<QString> otherJobIds = m_internal->findCoincidentJobs(jobRec);
    foreach (QString jobId, otherJobIds)
    {
      emit this->jobOverwritten(jobId);
      QCoreApplication::processEvents(); // wait for events to complete
    }

    // Update current project
    smtk::io::Logger logger;
    m_internal->m_ace3pProject->onJobSubmit(jobRec, m_internal->m_exportFileList, logger);
    if (logger.hasErrors())
    {
      std::string errString = logger.convertToString();
      emit this->errorMessage(QString::fromStdString(errString));
    }

    // Check if another job was overwritten

    // Notify system
    emit this->jobSubmitted(jobId);

    QString msg = QString("Job %1 was submitted to %2").arg(jobId, m_internal->m_machine);
    emit this->progressMessage(msg);
  }
  else
  {
    emit this->errorMessage("Submit-Job call did not return job id.");
  }
  reply->deleteLater();

  // For testing only
  // std::string recString = m_internal->m_jobRecordGenerator.get().dump(4);
  // DebugMessageMacro(recString);

  m_internal->m_state = State::Idle;
}

} // namespace ace3p
} // namespace simulation
} // namespace smtk
