//=========================================================================
//  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/examples/cxx/qtProjectSubmitTestWidget.h"

#include "smtk/newt/qtDownloadFolderRequester.h"
#include "smtk/newt/qtNewtInterface.h"
#include "smtk/simulation/ace3p/Metadata.h"
#include "smtk/simulation/ace3p/Project.h"
#include "smtk/simulation/ace3p/operations/Export.h"
#include "smtk/simulation/ace3p/operations/Read.h"
#include "smtk/simulation/ace3p/qt/qtNewtJobSubmitter.h"
#include "smtk/simulation/ace3p/qt/qtNewtJobTracker.h"
#include "smtk/simulation/ace3p/utility/AttributeUtils.h"

#include "smtk/simulation/ace3p/examples/cxx/ui_qtProjectSubmitTestWidget.h"

// SMTK includes
#include "smtk/attribute/Attribute.h"
#include "smtk/attribute/DirectoryItem.h"
#include "smtk/attribute/FileItem.h"
#include "smtk/attribute/GroupItem.h"
#include "smtk/attribute/IntItem.h"
#include "smtk/attribute/Item.h"
#include "smtk/attribute/Resource.h"
#include "smtk/attribute/ResourceItem.h"
#include "smtk/attribute/SearchStyle.h"
#include "smtk/attribute/StringItem.h"
#include "smtk/common/StringUtil.h"
#include "smtk/io/Logger.h"
#include "smtk/project/Manager.h"
#include "smtk/project/Project.h"

#include <QDebug>
#include <QDir>
#include <QFile>
#include <QFileInfo>
#include <QJsonArray>
#include <QJsonDocument>
#include <QJsonObject>
#include <QMessageBox>
#include <QNetworkReply>
#include <QString>
#include <QStringList>
#include <QTextStream>

#include <ctime>
#include <set>
#include <string>

namespace
{
const int OP_SUCCEEDED = static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED);
const smtk::attribute::SearchStyle RECURSIVE = smtk::attribute::SearchStyle::RECURSIVE;
const QString MACHINE("cori");

// 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::AttributePtr 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)
  {
    qWarning() << "Warning: Did not find attribute ValueItem \"" << itemName.c_str() << "\".";
  }
  else if (valueItem->isEnabled() && valueItem->isSet())
  {
    setField(valueItem->valueAsString());
  }
}

} // namespace

namespace smtk
{
namespace simulation
{
namespace ace3p
{

qtProjectSubmitTestWidget::qtProjectSubmitTestWidget(
  smtk::project::ManagerPtr projectManager,
  QWidget* parentWidget)
  : QWidget(parentWidget)
  , m_projectManager(projectManager)
#ifdef USE_NEWT_INTERFACE
  , m_downloader(new newt::qtDownloadFolderRequester(this))
  , m_jobSubmitter(new qtNewtJobSubmitter(this))
  , m_jobTracker(new qtNewtJobTracker(this))
#else
#endif
  , ui(new Ui::qtProjectSubmitTestWidget)
{
  this->ui->setupUi(this);

  // Check for PYTHONPATH
  // Is needed same way that ctest executables need it.
  if (!qEnvironmentVariableIsSet("PYTHONPATH"))
  {
    QString msg = "WARNING: The PYTHONPATH environment variable is not set.\n"
                  "If the Export operation fails, you probably need to set PYTHONPATH to point to\n"
                  "<smtk-build-directory>/lib/python3.9/site-packages.\n";
    this->ui->m_messageWidget->addItem(msg);
    this->ui->m_messageWidget->item(0)->setForeground(Qt::red);
  }

#ifdef USE_NEWT_INTERFACE
  this->setWindowTitle("NEWT Job Submit");
#else
#endif

  m_jobTracker->setPollingIntervalSeconds(5);

  QObject::connect(
    this->ui->m_loginButton,
    &QPushButton::clicked,
    this,
    &qtProjectSubmitTestWidget::onLoginTrigger);
  QObject::connect(
    this->ui->m_loadProjectButton,
    &QPushButton::clicked,
    this,
    &qtProjectSubmitTestWidget::onProjectTrigger);
  QObject::connect(
    this->ui->m_pollOnceButton,
    &QPushButton::clicked,
    this,
    &qtProjectSubmitTestWidget::onPollOnceTrigger);
  QObject::connect(
    this->ui->m_pollContinuousCheckBox,
    &QCheckBox::toggled,
    this,
    &qtProjectSubmitTestWidget::onPollContinuousToggled);
  QObject::connect(
    this->ui->m_submitJobButton,
    &QPushButton::clicked,
    this,
    &qtProjectSubmitTestWidget::onSubmitJobTrigger);
  QObject::connect(
    this->ui->m_downloadButton,
    &QPushButton::clicked,
    this,
    &qtProjectSubmitTestWidget::onDownloadTrigger);
  QObject::connect(
    this->ui->m_clearMessagesButton,
    &QPushButton::clicked,
    this,
    &qtProjectSubmitTestWidget::onClearMessagesTrigger);

  QObject::connect(
    m_jobSubmitter, &qtNewtJobSubmitter::progressMessage, [this](const QString& msg) {
      this->ui->m_messageWidget->addItem(msg);
    });
  QObject::connect(m_jobSubmitter, &qtNewtJobSubmitter::errorMessage, [this](const QString& msg) {
    int last = this->ui->m_messageWidget->count();
    this->ui->m_messageWidget->addItem(msg);
    this->ui->m_messageWidget->item(last)->setForeground(Qt::red);
  });
  QObject::connect(
    m_jobSubmitter, &qtNewtJobSubmitter::jobOverwritten, [this](const QString& jobId) {
      QString msg = QString("Job %1 is being overwritten").arg(jobId);
      this->ui->m_messageWidget->addItem(msg);
    });
  QObject::connect(m_jobSubmitter, &qtNewtJobSubmitter::jobSubmitted, [this](const QString& jobId) {
    this->ui->m_pollContinuousCheckBox->setChecked(true);
    this->m_jobTracker->addNewJob(jobId);
  });

  QObject::connect(
    m_jobTracker, &qtAbstractJobTracker::jobStatus, this, &qtProjectSubmitTestWidget::onJobStatus);
  QObject::connect(m_jobTracker, &qtAbstractJobTracker::pollingStateChanged, [this](bool polling) {
    QString text = polling ? "ON" : "Off";
    this->ui->m_pollingStateLabel->setText(text);
  });

  QObject::connect(
    m_downloader, &::newt::qtDownloadFolderRequester::progressMessage, [this](const QString& msg) {
      this->ui->m_messageWidget->addItem(msg);
    });
  QObject::connect(
    m_downloader, &::newt::qtDownloadFolderRequester::errorMessage, [this](const QString& msg) {
      int last = this->ui->m_messageWidget->count();
      this->ui->m_messageWidget->addItem(msg);
      this->ui->m_messageWidget->item(last)->setForeground(Qt::red);
    });
  QObject::connect(m_downloader, &::newt::qtDownloadFolderRequester::downloadComplete, [this]() {
    this->ui->m_messageWidget->addItem("Downloads complete");
  });
}

void qtProjectSubmitTestWidget::setUser(const QString& user)
{
  this->ui->m_userName->setText(user);
}

void qtProjectSubmitTestWidget::setNewtSessionId(const QString& id)
{
  this->ui->m_sessionId->setText(id);
}

void qtProjectSubmitTestWidget::setNerscRepository(const QString& repository)
{
  this->ui->m_repository->setText(repository);
}

void qtProjectSubmitTestWidget::setProjectDirectory(const QString& path)
{
  this->ui->m_projectDirectory->setText(path);
}

void qtProjectSubmitTestWidget::onClearMessagesTrigger()
{
  this->ui->m_messageWidget->clear();
}

void qtProjectSubmitTestWidget::onDownloadTrigger()
{
  QDir qDir;
  auto* newt = newt::qtNewtInterface::instance();

  // Get selected job
  QString jobId = this->ui->m_jobSelect->currentText();
  qDebug() << "Selected job id" << jobId;
  int jobIndex = m_project->jobsManifest()->findIndex(jobId.toStdString());
  // qDebug() << "Manifiest index" << jobIndex;

  std::string remoteJobDirectory;
  m_project->getJobRecordField(jobIndex, "runtime_job_folder", remoteJobDirectory);
  qDebug() << "Remote directory:" << remoteJobDirectory.c_str();

  // See if we need to download the ncdf file
  std::string ncdfFilename;
  m_project->getJobRecordField(jobIndex, "runtime_mesh_filename", ncdfFilename);
  QString localMeshPath =
    QString("%1/%2").arg(m_project->assetsDirectory().c_str()).arg(ncdfFilename.c_str());
  if (qDir.exists(localMeshPath))
  {
    QString msg = QString("Using model file at %1").arg(localMeshPath);
    this->ui->m_messageWidget->addItem(msg);
  }
  else
  {
    QFile* file = new QFile(localMeshPath);
    if (!file->open(QIODevice::WriteOnly)) // truncates any existing file
    {
      delete file;
      QString msg = QString("Unable to open local file for writing at %1").arg(localMeshPath);
      QMessageBox::critical(this, "Error Opening Local File", msg);
    }

    QString remotePath = QString("%1/%2").arg(remoteJobDirectory.c_str()).arg(ncdfFilename.c_str());
    QNetworkReply* reply = newt->requestFileDownload(MACHINE, remotePath);

    QObject::connect(
      reply, &QNetworkReply::readyRead, [this, reply, file]() { file->write(reply->readAll()); });

    QObject::connect(reply, &QNetworkReply::finished, [this, reply, localMeshPath]() {
      QString msg;
      if (reply->error() == 0)
      {
        msg = QString("Downloaded file %1").arg(localMeshPath);
        this->ui->m_messageWidget->addItem(msg);
      }
      else
      {
        msg = QString("ERROR Downloading to file %1: %2").arg(localMeshPath).arg(reply->error());
        int last = this->ui->m_messageWidget->count();
        this->ui->m_messageWidget->addItem(msg);
        this->ui->m_messageWidget->item(last)->setForeground(Qt::red);
      }
      reply->deleteLater();
    });
  }

  // Get paths to local and remote results directories
  std::string resultsFolder;
  m_project->getJobRecordField(jobIndex, "results_subfolder", resultsFolder);
  QString remoteResultsPath =
    QString("%1/%2").arg(remoteJobDirectory.c_str()).arg(resultsFolder.c_str());

  std::string localJobDirectory;
  m_project->getJobRecordField(jobIndex, "local_job_folder", localJobDirectory);
  std::string downloadDirectory = localJobDirectory + "/download/" + resultsFolder;
  QString downloadPath = QString::fromStdString(downloadDirectory);

  // Start the download
  bool emFilesOnly = true;
  if (emFilesOnly)
  {
    // Only download mode files
    m_downloader->startDownload(downloadPath, MACHINE, remoteResultsPath, false, "*.mod");
  }
  else
  {
    // Download everything
    m_downloader->startDownload(downloadPath, MACHINE, remoteResultsPath);
  }
}

void qtProjectSubmitTestWidget::onJobStatus(
  const QString& /* cumulusId */,
  const QString& status,
  const QString& jobId)
{
  QString message;
  QTextStream qs(&message);
  qs << "job status for " << jobId << ": " << status << ".";
  int last = this->ui->m_messageWidget->count();
  this->ui->m_messageWidget->addItem(message);
  this->ui->m_messageWidget->scrollToItem(this->ui->m_messageWidget->item(last));

  int index = m_project->jobsManifest()->findIndex(jobId.toStdString());
  m_project->setJobRecordField(index, "status", status.toStdString());
  m_project->writeJobsManifest();
}

void qtProjectSubmitTestWidget::onNewtLoginComplete(const QString& userName)
{
  QString message;
  QTextStream qs(&message);
  qs << "Logged in as user " << userName << ".";
  this->ui->m_messageWidget->addItem(message);

  this->ui->m_loginButton->setEnabled(false);
  this->updateControls();
}

void qtProjectSubmitTestWidget::onLoginTrigger()
{
  // Check input fields
  QString userName = this->ui->m_userName->text();
  if (userName.isEmpty())
  {
    this->ui->m_messageWidget->addItem("Cannot login: User name not specified");
    return;
  }

  QString newtSessionId = this->ui->m_sessionId->text();
  if (newtSessionId.isEmpty())
  {
    this->ui->m_messageWidget->addItem("Cannot login: NEWT session id not specified");
    return;
  }

  auto* newt = newt::qtNewtInterface::instance();
  QObject::connect(
    newt,
    &newt::qtNewtInterface::loginComplete,
    this,
    &qtProjectSubmitTestWidget::onNewtLoginComplete);

  newt->mockLogin(userName, newtSessionId);
}

void qtProjectSubmitTestWidget::onPollContinuousToggled(bool checked)
{
  bool started = m_jobTracker->enablePolling(checked);
  qDebug() << "Polling started?" << started;
}

void qtProjectSubmitTestWidget::onPollOnceTrigger()
{
  bool started = m_jobTracker->pollOnce();
  qDebug() << "Started poll? " << started;
  if (!started)
  {
    this->ui->m_messageWidget->addItem("Did NOT poll server; check console for more info.");
  }
}

void qtProjectSubmitTestWidget::onProjectTrigger()
{
  QString path = this->ui->m_projectDirectory->text();
  QDir dir(path);
  if (!dir.exists())
  {
    QString msg = QString("The project directory is not found: %1").arg(path);
    QMessageBox::critical(this, "Error", msg);
    return;
  }

  QString dirName = dir.dirName();
  QString projectFilePath;
  QTextStream qs(&projectFilePath);
  qs << dir.absolutePath() << "/" << dirName << ".project.smtk";
  if (!dir.exists(projectFilePath))
  {
    QString msg = QString("The project file is not found: %1").arg(projectFilePath);
    QMessageBox::critical(this, "Error", msg);
    return;
  }

  auto opManager = m_projectManager->operationManager();

  // Create Read operation
  auto readOp = opManager->create<smtk::simulation::ace3p::Read>();
  bool setFileOK =
    readOp->parameters()->findFile("filename")->setValue(projectFilePath.toStdString());
  auto result = readOp->operate();
  int outcome = result->findInt("outcome")->value();
  if (outcome != OP_SUCCEEDED)
  {
    QString msg =
      QString("Failed to load project, file: %1, Outcome %2").arg(projectFilePath).arg(outcome);
    QMessageBox::critical(this, "Error Loading Project", msg);
    return;
  }

  smtk::attribute::ResourceItemPtr projectItem = result->findResource("resource");
  smtk::resource::ResourcePtr resource = projectItem->value();
  m_project = std::dynamic_pointer_cast<smtk::simulation::ace3p::Project>(resource);
  this->ui->m_messageWidget->addItem("Project loaded");

  this->ui->m_loadProjectButton->setEnabled(false);
  this->updateControls();
}

void qtProjectSubmitTestWidget::onSubmitJobTrigger()
{
  this->ui->m_messageWidget->addItem("Exporting project...");
  m_exportFiles.clear();
  auto opManager = m_projectManager->operationManager();
  auto resManager = m_projectManager->resourceManager();

  // Create export operation to generate input file and slurm script (only)
  auto exportOp = opManager->create<smtk::simulation::ace3p::Export>();
  exportOp->parameters()->associate(m_project);

  // Setup inputs (much of this copied from pqACE3PExportBehavior.cxx)
  int stageIndex = static_cast<int>(m_project->currentStageIndex());
  exportOp->parameters()->findInt("stage-index")->setValue(stageIndex);

  std::string analysisName = "analysis"; // default
  std::string solverName;                // which ACE3P code (Omega3P, S3P, T3P, etc)

  // Check that Analysis is specified in the attribute resource
  auto attResource = m_project->stage(stageIndex)->attributeResource();
  smtk::simulation::ace3p::AttributeUtils attUtils;
  auto analysisAtt = attUtils.getAnalysisAtt(attResource);
  // InternalCheckMacro(analysisAtt != nullptr, "Internal Error: Analysis attribute not found.");

  auto analysisItem = analysisAtt->findString("Analysis");
  // InternalCheckMacro(analysisItem != nullptr, "Internal Error: Analysis item not found.");
  if (!analysisItem->isSet())
  {
    this->ui->m_messageWidget->addItem("No Analysis Set! Cannot export.");
    int last = this->ui->m_messageWidget->count() - 1;
    this->ui->m_messageWidget->item(last)->setForeground(Qt::red);
    return;
  }
  solverName = analysisItem->value();

  QFileInfo projectFileInfo(QString::fromStdString(m_project->location()));
  QDir projectDir = projectFileInfo.absoluteDir();
  QDir exportDir(projectDir);
  exportDir.mkdir("export");
  exportDir.cd("export");
  std::string outputFolder = exportDir.absolutePath().toStdString();
  exportOp->parameters()->findDirectory("OutputFolder")->setIsEnabled(true);
  exportOp->parameters()->findDirectory("OutputFolder")->setValue(outputFolder);

  const std::string outFilePrefix = "analysis1";
  auto prefixItem = exportOp->parameters()->findString("OutputFilePrefix");
  prefixItem->setIsEnabled(true);
  prefixItem->setValue(outFilePrefix);

  auto nerscItem = exportOp->parameters()->findGroup("NERSCSimulation");
  nerscItem->setIsEnabled(true);
  nerscItem->find("UseNewtInterface")->setIsEnabled(true);
  QString repo = this->ui->m_repository->text();
  if (repo.isEmpty())
  {
    QMessageBox::critical(this, "Missing Data", "ERROR: NERSC Repository not set.");
    return;
  }
  auto repoItem = nerscItem->findAs<smtk::attribute::StringItem>("NERSCRepository", RECURSIVE);
  repoItem->setValue(repo.toStdString());
  QString newtSessionId = this->ui->m_sessionId->text();
  auto idItem = nerscItem->findAs<smtk::attribute::StringItem>("NEWTSessionId", RECURSIVE);
  idItem->setValue(newtSessionId.toStdString());
  auto folderItem = nerscItem->findAs<smtk::attribute::StringItem>("SubFolder", RECURSIVE);
  folderItem->setValue("test/newt-job-submit");
  auto resultsDirItem =
    nerscItem->findAs<smtk::attribute::StringItem>("ResultsDirectory", RECURSIVE);
  resultsDirItem->setIsEnabled(true);
  resultsDirItem->setValue("omega3p_results");

  // Run exportOp to get input file and slurm script
  auto exportResult = exportOp->operate();
  int outcome = exportResult->findInt("outcome")->value();
  if (outcome != OP_SUCCEEDED)
  {
    qDebug() << QString::fromStdString(exportOp->log().convertToString());
    QString msg = QString("Export project failed, Outcome %1").arg(outcome);
    QMessageBox::critical(this, "Error Export Project", msg);
    return;
  }

  int last = this->ui->m_messageWidget->count() - 1;
  auto* item = this->ui->m_messageWidget->item(last);
  QString text = item->text() + "  Done";
  item->setText(text);

  // Set runtime input directory just for testing
  auto inputFolderItem = exportResult->findAs<smtk::attribute::StringItem>("NerscInputFolder");
  inputFolderItem->setIsEnabled(true);
  inputFolderItem->setValue("TestInputFolder");

  auto exportParams = exportOp->parameters();
  m_jobSubmitter->submitAnalysisJob(m_project, exportParams, exportResult);
}

void qtProjectSubmitTestWidget::updateControls()
{
  // Enable the newt controls if project is loaded and user is logged in.
  if (m_project == nullptr)
  {
    return;
  }

  auto* newt = newt::qtNewtInterface::instance();
  if (!newt->isLoggedIn())
  {
    return;
  }

  // (else)
  this->ui->m_pollOnceButton->setEnabled(true);
  this->ui->m_pollContinuousCheckBox->setEnabled(true);
  this->ui->m_submitJobButton->setEnabled(true);

  // Build combobox with project jobs
  QStringList jobList;
  nlohmann::json jobsData = m_project->jobsManifest()->getData()["Jobs"];
  for (auto rit = jobsData.rbegin(); rit != jobsData.rend(); ++rit)
  {
    const std::string jobId = (*rit)["slurm_id"].get<std::string>();
    if (!jobId.empty())
    {
      jobList.push_back(QString::fromStdString(jobId));
    }
  }
  this->ui->m_jobSelect->clear();
  this->ui->m_jobSelect->insertItems(0, jobList);
  this->ui->m_jobSelect->setEnabled(true);
  this->ui->m_downloadButton->setEnabled(true);
}

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