/**
 * \file core/tandem_run/tandemslurmprocess.h
 * \date 21/09/2023
 * \author Olivier Langella
 * \brief handles execution of a bunch of X!Tandem process through slurm
 * workload manager
 */

/*******************************************************************************
 * Copyright (c) 2023 Olivier Langella
 *<Olivier.Langella@universite-paris-saclay.fr>.
 *
 * This file is part of i2MassChroQ.
 *
 *     i2MassChroQ is free software: you can redistribute it and/or modify
 *     it under the terms of the GNU General Public License as published by
 *     the Free Software Foundation, either version 3 of the License, or
 *     (at your option) any later version.
 *
 *     i2MassChroQ is distributed in the hope that it will be useful,
 *     but WITHOUT ANY WARRANTY; without even the implied warranty of
 *     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *     GNU General Public License for more details.
 *
 *     You should have received a copy of the GNU General Public License
 *     along with i2MassChroQ.  If not, see <http://www.gnu.org/licenses/>.
 *
 ******************************************************************************/

#include "tandemslurmprocess.h"
#include <QDebug>
#include <QSettings>
#include <cutelee/engine.h>
#include <pappsomspp/core/pappsoexception.h>
#include <pappsomspp/core/exception/exceptioninterrupted.h>
#include "../../files/tandemparametersfile.h"
#include "../../utils/utils.h"

TandemSlurmProcess::TandemSlurmProcess(MainWindow *p_main_window,
                                       WorkMonitor *p_monitor,
                                       const TandemRunBatch &tandem_run_batch)
  : TandemBatchProcess(p_monitor, tandem_run_batch), mp_workMonitor(p_monitor)
{
  /*
  #!/bin/bash
#SBATCH --mail-user=olivier.langella@universite-paris-saclay.fr
#SBATCH --mail-type=end,fail
#SBATCH --partition=debug
#SBATCH --job-name="sbatch Example 01"
#SBATCH --cpus-per-task=10
# Your job commands
masschroq -t /tmp -c 10 balazadeh_dil.masschroqml

  */
  // mp_slurmProcess = new SlurmProcess(p_main_window, p_monitor, "X!Tandem");

  QSettings settings;

  m_condorRequestMemory =
    settings.value("tandem/condor_request_memory", "10000").toUInt();
  // m_condorRequestMemory =
  //  settings.value("tandem/condor_request_memory", "10000").toUInt();


  prepareTemporaryDirectory();
}

TandemSlurmProcess::~TandemSlurmProcess()
{

  for(const TandemSlurmProcess::ProcessUnit &process_unit : m_processUnitList)
    {
      delete process_unit.mp_slurmProcess;
    }
}

void
TandemSlurmProcess::prepareXmlDatabaseFile()
{

  QFile xml_database_file(QString("%1/database.xml").arg(mp_tmpDir->path()));

  if(xml_database_file.open(QIODevice::WriteOnly))
    {
      _xml_database_file =
        QFileInfo(xml_database_file.fileName()).absoluteFilePath();
      QXmlStreamWriter *p_out = new QXmlStreamWriter();
      p_out->setDevice(&xml_database_file);
      writeXmlDatabaseFile(p_out);
      xml_database_file.close();
      delete p_out;
    }
  else
    {
      throw pappso::PappsoException(
        QObject::tr("error : cannot open the XML database file : %1\n")
          .arg(xml_database_file.fileName()));
    }
}

QString
TandemSlurmProcess::getSlurmBatchContent(const TandemRunBatch &tandem_run_batch,
                                         const QString &input_tandem) const
{
  qDebug();
  Cutelee::Engine *engine_p = new Cutelee::Engine();
  auto loader = std::shared_ptr<Cutelee::FileSystemTemplateLoader>(
    new Cutelee::FileSystemTemplateLoader());

  qDebug();
  loader->setTemplateDirs({":/templates/resources/templates/slurm"});

  qDebug();
  engine_p->addTemplateLoader(loader);

  qDebug();
  Cutelee::Template slurm_tandem_template =
    engine_p->loadByName("slurm_tandemng_onejob.txt");

  qDebug();
  Cutelee::Context tandem_context;
  tandem_context.insert("job_name", "tandemng");
  tandem_context.insert("request_cpus", tandem_run_batch._number_of_threads);
  tandem_context.insert("request_mem", m_condorRequestMemory);
  tandem_context.insert("count_tandem_job", 1);
  tandem_context.insert("tandem_bin_path", tandem_run_batch._tandem_bin_path);

  tandem_context.insert("tmp_dir", mp_tmpDir->path());

  tandem_context.insert("tandem_input", input_tandem);


  qDebug();
  return slurm_tandem_template->render(&tandem_context);
}
void
TandemSlurmProcess::run()
{
  qDebug();
  qDebug() << _tandem_run_batch._preset_file;
  m_memoryMax    = 0;
  m_countJobDone = 0;

  TandemParametersFile orig_parameter_file(_tandem_run_batch._preset_file);
  _preset_file = QString("%1/%2")
                   .arg(mp_tmpDir->path())
                   .arg(orig_parameter_file.getFilename());

  TandemParametersFile new_parameter_file(_preset_file);
  TandemParameters new_param = orig_parameter_file.getTandemParameters();
  //  <note type="input" label="spectrum, threads">1</note>
  new_param.setParamLabelValue(
    "spectrum, threads",
    QString("%1").arg(_tandem_run_batch._number_of_threads));
  new_parameter_file.setTandemParameters(new_param);

  prepareXmlDatabaseFile();


  // std::vector<QTemporaryFile *> input_file_list;

  int i = 0;
  _p_monitor->setTotalSteps(_tandem_run_batch._mz_file_list.size());
  for(QString mz_file : _tandem_run_batch._mz_file_list)
    {


      QTemporaryFile *p_xml_input_file =
        new QTemporaryFile(QString("%1/tandem").arg(mp_tmpDir->path()));
      p_xml_input_file->setAutoRemove(false);
      if(p_xml_input_file->open())
        {
          QXmlStreamWriter *p_xml_out = new QXmlStreamWriter();
          p_xml_out->setDevice(p_xml_input_file);

          writeXmlInputFile(p_xml_out, mz_file);

          p_xml_input_file->close();
          delete p_xml_out;
          if(Utils::guessDataFileFormatFromFile(mz_file) !=
             pappso::Enums::MsDataFormat::unknown)
            {
              QSettings settings;
              QString tims_tmp_dir_path =
                settings.value("timstof/tmp_dir_path", QDir::tempPath())
                  .toString();
              if(!QDir(tims_tmp_dir_path).exists())
                {
                  QDir().mkdir(tims_tmp_dir_path);
                }
            }
          else
            {
              throw pappso::PappsoException(
                QObject::tr("The loaded file:\n %1 \n hasn't a known format")
                  .arg(mz_file));
            }

          startProcessUnit(mz_file,
                           QFileInfo(p_xml_input_file->fileName()).fileName());
        }
      else
        {
          throw pappso::PappsoException(
            QObject::tr(
              "error : cannot open the XML X!Tandem input file : %1\n")
              .arg(p_xml_input_file->fileName()));
        }


      i++;
    }


  while(batchIsRunning())
    {
      if(mp_workMonitor->shouldIstop())
        {
          killJobs();
          throw pappso::ExceptionInterrupted(
            "X!Tandem jobs stopped by the user");
        }
      _p_monitor->message(whatAboutMyJobs());

      QThread::msleep(m_slurmStatusTimerMillisecond);
    }


  _p_monitor->finished(QObject::tr("%1 slurm X!Tandem job(s) finished")
                         .arg(_tandem_run_batch._mz_file_list.size()));
  qDebug();
}


void
TandemSlurmProcess::prepareTemporaryDirectory()
{
  // /gorgone/pappso/tmp
  QSettings settings;
  QString slurm_tmp_dir =
    QString("%1/i2masschroq")
      .arg(settings.value("slurm/tmp_dir", "/tmp").toString());

  if(mp_tmpDir != nullptr)
    {
      delete mp_tmpDir;
      mp_tmpDir = nullptr;
    }
  mp_tmpDir = new QTemporaryDir(slurm_tmp_dir);
  mp_tmpDir->setAutoRemove(
    settings.value("slurm/tmp_dir_autoremove", true).toBool());

  if(!mp_tmpDir->isValid())
    {
      // dir.path() returns the unique directory path
      throw pappso::PappsoException(
        QObject::tr("problem creating slurm temporary directory in %1\n")
          .arg(slurm_tmp_dir));
    }
}

const QTemporaryDir *
TandemSlurmProcess::getTmpDir() const
{
  return mp_tmpDir;
}

void
TandemSlurmProcess::startProcessUnit(const QString &mz_input_file,
                                     const QString &input_tandem_file)
{

  qDebug() << "input_tandem_file=" << input_tandem_file;

  // slurm sbatch submit file :
  QFileInfo sbatch_file_info(
    QString("%1/sbatch_%2.txt").arg(mp_tmpDir->path()).arg(input_tandem_file));
  QFile submit_file(sbatch_file_info.absoluteFilePath());
  QTextStream *p_out = nullptr;

  if(submit_file.open(QIODevice::WriteOnly))
    {
      p_out = new QTextStream();
      p_out->setDevice(&submit_file);

      *p_out << getSlurmBatchContent(_tandem_run_batch, input_tandem_file);
      p_out->flush();
    }
  else
    {
      throw pappso::PappsoException(
        QObject::tr("error : cannot open slurm batch file : %1\n")
          .arg(submit_file.fileName()));
    }
  submit_file.close();


  // now run condor job on submit_file
  SlurmProcess *p_slurm_process = new SlurmProcess();
  m_processUnitList.push_back(
    {p_slurm_process, mz_input_file, input_tandem_file});

  p_slurm_process->startSlurmBatch(sbatch_file_info);
  // surveyCondorJob();
}

bool
TandemSlurmProcess::batchIsRunning()
{
  std::size_t old_job_count_done = m_countJobDone;
  m_countJobDone                 = 0;
  bool is_running                = false;
  if(m_processUnitList.size() > 0)
    {
      m_jsonSlurmJobStatusDocument =
        m_processUnitList.front()
          .mp_slurmProcess->getJsonSlurmJobStatusDocument();

      for(TandemSlurmProcess::ProcessUnit &process_unit : m_processUnitList)
        {
          process_unit.mp_slurmProcess->refreshJsonJobStatus(
            m_jsonSlurmJobStatusDocument);
          if(process_unit.mp_slurmProcess->isRunning())
            {
              is_running = true;
            }

          if(process_unit.mp_slurmProcess->getSlurmJobState() ==
             SlurmJobState::RUNNING)
            {
              m_memoryMax = std::max(
                m_memoryMax,
                process_unit.mp_slurmProcess->getAverageMemorySizeMegaBytes());
            }
          if(process_unit.mp_slurmProcess->getSlurmJobState() ==
             SlurmJobState::COMPLETED)
            {
              process_unit.done = true;
            }
          if(process_unit.done)
            m_countJobDone++;
        }
    }
  std::size_t count_diff = m_countJobDone - old_job_count_done;
  for(std::size_t i = 0; i < count_diff; i++)
    {
      _p_monitor->count();
    }
  return is_running;
}

void
TandemSlurmProcess::killJobs()
{
  for(const TandemSlurmProcess::ProcessUnit &process_unit : m_processUnitList)
    {
      process_unit.mp_slurmProcess->killJob();
      delete process_unit.mp_slurmProcess;
    }
  m_processUnitList.clear();
}

QString
TandemSlurmProcess::whatAboutMyJobs()
{

  QStringList reasons;
  for(const TandemSlurmProcess::ProcessUnit &process_unit : m_processUnitList)
    {
      QString reason = process_unit.mp_slurmProcess->getSlurmJobPendingReason();
      if(!reason.isEmpty())
        {
          reasons << reason;
        }
    }
  reasons.removeDuplicates();

  QStringList message_list;
  message_list << QString("%1 pending jobs (%2)")
                    .arg(countSlurmJobState(SlurmJobState::PENDING))
                    .arg(reasons.join(", "));
  message_list << QString("%1 running jobs using %2Mb of memory each (maximum)")
                    .arg(countSlurmJobState(SlurmJobState::RUNNING))
                    .arg(m_memoryMax);
  message_list << QString("%1 completed jobs")
                    .arg(countSlurmJobState(SlurmJobState::COMPLETED));
  message_list
    << QString("%1 failed jobs").arg(countSlurmJobState(SlurmJobState::FAILED));
  return message_list.join("\n");
}

std::size_t
TandemSlurmProcess::countSlurmJobState(SlurmJobState slurm_state) const
{
  std::size_t count = 0;
  for(const TandemSlurmProcess::ProcessUnit &process_unit : m_processUnitList)
    {
      if(process_unit.mp_slurmProcess->getSlurmJobState() == slurm_state)
        count++;
    }
  return count;
}
