/**
 * \file core/slurm_process/slurmprocess.cpp
 * \date 21/09/2023
 * \author Olivier Langella
 * \brief handles execution of processes 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 "slurmprocess.h"
#include <QDebug>
#include "../../utils/utils.h"

SlurmProcess::SlurmProcess()
{
  /*
  #!/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

#SBATCH --nodes=3
#SBATCH --ntasks=36
#SBATCH --cpus-per-task=2
#SBATCH --mem-per-cpu=2000

export OMP_NUM_THREADS=2

for i in {1..36}; do
    srun -N 1 -n 1 ./program input${i} >& out${i} &
done
  */

  QSettings settings;
  // m_condorSubmitCommand =
  // settings.value("condor/submit", "/usr/bin/condor_submit").toString();
  m_slurmSbatchCommand =
    settings.value("slurm/sbatch", "/usr/bin/sbatch").toString();
  m_slurmSrunCommand = settings.value("slurm/srun", "/usr/bin/srun").toString();
  m_slurmScancelCommand =
    settings.value("slurm/scancel", "/usr/bin/scancel").toString();
  m_slurmSinfoCommand =
    settings.value("slurm/sinfo", "/usr/bin/sinfo").toString();
  m_slurmSqueueCommand =
    settings.value("slurm/squeue", "/usr/bin/squeue").toString();

  m_slurmStatCommand =
    settings.value("slurm/sstat", "/usr/bin/sstat").toString();
}

SlurmProcess::~SlurmProcess()
{
}

std::size_t
SlurmProcess::startSlurmBatch(const QFileInfo &batch_file_name)
{

  if(!batch_file_name.absoluteDir().exists())
    {
      QString err = tr("The directory %1 does not exists (srun)")
                      .arg(batch_file_name.absolutePath());
      throw pappso::PappsoException(err);
    }
  QDir::setCurrent(batch_file_name.absolutePath());
  QStringList arguments;


  qDebug() << batch_file_name;

  // arguments.clear();
  arguments << "--parsable" << batch_file_name.absoluteFilePath();

  start(m_slurmSbatchCommand, arguments);

  if(!waitForStarted())
    {
      QString err =
        tr("Could not start slurm process '%1' with arguments '%2': %3")
          .arg(m_slurmSbatchCommand)
          .arg(arguments.join(QStringLiteral(" ")))
          .arg(errorString());
      throw pappso::PappsoException(err);
    }


  if(!waitForFinished(m_maxLaunchTime))
    {
      throw pappso::PappsoException(
        QObject::tr("Slurm batch process failed to finish"));
    }

  QString perr = readAllStandardError();
  if(perr.length())
    {

      qDebug() << "readAllStandardError perr:" << perr;
      throw pappso::PappsoException(
        QObject::tr("Slurm batch process failed :\n%1").arg(perr));
    }
  else
    {
      qDebug() << "readAllStandardError OK " << perr;
    }

  QString pjob = readAllStandardOutput();
  if(pjob.length())
    {
      qDebug() << "readAllStandardOutput OK pjob:" << pjob;
    }
  else
    {
      qDebug() << "readAllStandardOutput " << pjob;
      throw pappso::PappsoException(
        QObject::tr("Slurm batch process failed :\n%1").arg(pjob));
    }

  qDebug() << pjob;

  std::size_t job_number = pjob.toULongLong();
  qDebug() << job_number;
  // Submitting job(s).\n1 job(s) submitted to cluster 29.\n
  // parseCondorJobNumber(pjob);

  // squeue --json
  m_slurmJobId = job_number;
  return job_number;
}

void
SlurmProcess::srun(const QDir &current_dir, const QStringList &arguments)
{
  qDebug() << m_slurmSrunCommand << " " << arguments.join(" ");
  if(!current_dir.exists())
    {
      QString err = tr("The directory %1 does not exists (srun)")
                      .arg(current_dir.absolutePath());
      throw pappso::PappsoException(err);
    }
  QDir::setCurrent(current_dir.absolutePath());
  start(m_slurmSrunCommand, arguments);
}


bool
SlurmProcess::isRunning()
{
  // refreshJsonJobStatus();
  if(m_slurmJobState == SlurmJobState::not_listed)
    return false;
  if(m_slurmJobState == SlurmJobState::CANCELLED)
    return false;
  if(m_slurmJobState == SlurmJobState::COMPLETED)
    return false;
  return true;
}

bool
SlurmProcess::killJob()
{

  QStringList arguments;


  qDebug() << m_slurmJobId;

  // arguments.clear();
  arguments << QString("%1").arg(m_slurmJobId);

  start(m_slurmScancelCommand, arguments);

  if(!waitForStarted())
    {
      QString err = tr("Could not stop slurm process '%1'").arg(m_slurmJobId);
      throw pappso::PappsoException(err);
    }


  if(!waitForFinished(m_maxLaunchTime))
    {
      throw pappso::PappsoException(
        QObject::tr("Slurm scancel process failed to finish"));
    }

  QString perr = readAllStandardError();
  if(perr.length())
    {

      qDebug() << "readAllStandardError " << perr;
      throw pappso::PappsoException(
        QObject::tr("Slurm scancel process failed :\n%1").arg(perr));
    }
  else
    {
      qDebug() << "readAllStandardError OK " << perr;
    }

  return true;
}


QJsonDocument
SlurmProcess::getJsonDocument(const QString &slurm_command)
{
  QStringList arguments;


  qDebug() << m_slurmJobId;

  // arguments.clear();
  arguments << "--json";

  start(slurm_command, arguments);

  if(!waitForStarted())
    {
      QString err = tr("Could not start %1 command").arg(slurm_command);
      throw pappso::PappsoException(err);
    }


  if(!waitForFinished(m_maxLaunchTime))
    {
      throw pappso::PappsoException(
        QObject::tr("Slurm %1 process failed to finish").arg(slurm_command));
    }

  QString perr = readAllStandardError();
  if(perr.length())
    {

      qDebug() << "readAllStandardError " << perr;
      throw pappso::PappsoException(QObject::tr("Slurm %1 process failed :\n%2")
                                      .arg(slurm_command)
                                      .arg(perr));
    }
  else
    {
      qDebug() << "readAllStandardError OK " << perr;
    }

  QString sinfo_txt = readAllStandardOutput();
  if(sinfo_txt.length())
    {
      qDebug() << "readAllStandardOutput OK " << sinfo_txt;
    }
  else
    {
      qDebug() << "readAllStandardOutput " << sinfo_txt;
      throw pappso::PappsoException(QObject::tr("Slurm %1 process failed :\n%2")
                                      .arg(slurm_command)
                                      .arg(sinfo_txt));
    }
  QByteArray json_bytes;
  qDebug() << json_bytes.size();
  json_bytes = QByteArray::fromStdString(sinfo_txt.toStdString());

  qDebug() << json_bytes.size();

  // qDebug() << json_bytes.toStdString().c_str();
  QJsonDocument json_doc;
  QJsonParseError error;
  json_doc = QJsonDocument::fromJson(json_bytes, &error);
  qDebug() << json_doc.toJson();
  qDebug() << error.errorString();
  return json_doc;
}

void
SlurmProcess::refreshJsonJobStatus(const QJsonDocument &json_slurm_squeue)
{
  if(json_slurm_squeue.isObject())
    {
      qDebug() << "is object";
    }
  if(json_slurm_squeue.isArray())
    {
      qDebug() << "is array";
    }
  if(json_slurm_squeue.object().find("jobs") !=
     json_slurm_squeue.object().end())
    {
      qDebug() << "jobs not found";
    }
  QJsonArray array_jobs = json_slurm_squeue.object().value("jobs").toArray();

  m_slurmJobState = SlurmJobState::not_listed;
  qDebug();
  for(const QJsonValue &element : array_jobs)
    {
      if(element.isObject())
        {
          qDebug();
          m_jsonObjectJob         = element.toObject();
          const QJsonValue &jobid = m_jsonObjectJob.value("job_id");
          if(jobid != QJsonValue::Undefined)
            {
              qDebug() << jobid.toString() << " " << jobid.toInteger();
              if((std::size_t)jobid.toInteger() == m_slurmJobId)
                {
                  setSlurmJobState(m_jsonObjectJob.value("job_state"));
                }
            }
        }
    }

  qDebug();
}

void
SlurmProcess::setSlurmJobState(const QJsonValue &jobstate)
{
  m_slurmJobState = SlurmJobState::undefined;
  if(jobstate != QJsonValue::Undefined)
    {
      QString jobstate_str = jobstate.toString();
      qDebug() << jobstate_str;
      if(jobstate == "COMPLETED")
        {
          m_slurmJobState = SlurmJobState::COMPLETED;
        }
      else if(jobstate == "CANCELLED")
        {
          m_slurmJobState = SlurmJobState::CANCELLED;
        }
      else if(jobstate == "PENDING")
        {
          m_slurmJobState = SlurmJobState::PENDING;
        }

      else if(jobstate == "RUNNING")
        {
          m_slurmJobState = SlurmJobState::RUNNING;
        }
      else if(jobstate == "COMPLETING")
        {
          m_slurmJobState = SlurmJobState::COMPLETING;
        }
      else if(jobstate == "FAILED")
        {
          m_slurmJobState = SlurmJobState::FAILED;
        }

      else if(jobstate == "PREEMPTED")
        {
          m_slurmJobState = SlurmJobState::PREEMPTED;
        }
      else if(jobstate == "STOPPED")
        {
          m_slurmJobState = SlurmJobState::STOPPED;
        }
      else if(jobstate == "SUSPENDED")
        {
          m_slurmJobState = SlurmJobState::SUSPENDED;
        }
      else
        {
          m_slurmJobState = SlurmJobState::unknown;
          qWarning() << "not known " << jobstate_str;
        }
    }
}

SlurmJobState
SlurmProcess::getSlurmJobState() const
{
  return m_slurmJobState;
}

QString
SlurmProcess::getSlurmJobPendingReason() const
{

  QString reason;
  if(m_slurmJobState == SlurmJobState::PENDING)
    {
      if(!m_jsonObjectJob.isEmpty())
        {
          const QJsonValue &js_reason = m_jsonObjectJob.value("state_reason");
          if(js_reason != QJsonValue::Undefined)
            {
              reason = js_reason.toString();
            }
        }
    }
  qDebug() << "m_slurmJobId=" << m_slurmJobId << " reason=" << reason;
  return reason;
}


QString
SlurmProcess::whatAboutMyJob() const
{
  QString message;

  QString reason;
  if(!m_jsonObjectJob.isEmpty())
    {
      const QJsonValue &js_reason = m_jsonObjectJob.value("state_reason");
      if(js_reason != QJsonValue::Undefined)
        {
          reason = js_reason.toString();
        }
    }
  if(m_slurmJobState == SlurmJobState::PENDING)
    {
      message = QString("your slurm job %1 is pending :\n%2\n")
                  .arg(m_slurmJobId)
                  .arg(reason);
    }

  message.append(QString("job state : %1, reason : %2")
                   .arg((std::uint8_t)m_slurmJobState)
                   .arg(reason));
  return (message);
}

QJsonDocument
SlurmProcess::getJsonSlurmJobStatusDocument()
{
  return getJsonDocument(m_slurmSqueueCommand);
}

std::size_t
SlurmProcess::getAverageMemorySizeMegaBytes()
{
  // sstat --format=AveVMSize -j 11
  // sstat --format=AveVMSize 71 -n
  //    459340K

  QStringList arguments;


  qDebug() << m_slurmJobId;

  // arguments.clear();
  arguments << "--format=AveVMSize"
            << "-n"
            << "-j" << QString("%1").arg(m_slurmJobId);

  start(m_slurmStatCommand, arguments);

  if(!waitForStarted())
    {
      QString err = tr("Could not sstat process '%1'").arg(m_slurmJobId);
      throw pappso::PappsoException(err);
    }


  if(!waitForFinished(m_maxLaunchTime))
    {
      throw pappso::PappsoException(
        QObject::tr("Slurm sstat process failed to finish"));
    }

  QString perr = readAllStandardError();
  if(perr.length())
    {

      qDebug() << "readAllStandardError " << perr;
      return 0;
    }
  else
    {
      qDebug() << "readAllStandardError OK " << perr;
    }
  QString mem = readAllStandardOutput();
  qDebug() << mem;
  bool is_ok;
  return Utils::getMegaBytesFromString(mem, is_ok);
}
