/*
 * (C) Copyright 2009-2016 ECMWF.
 *
 * This software is licensed under the terms of the Apache Licence Version 2.0
 * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
 * In applying this licence, ECMWF does not waive the privileges and immunities
 * granted to it by virtue of its status as an intergovernmental organisation
 * nor does it submit to any jurisdiction.
 */

#ifndef OOPS_UTIL_RESTART_H_
#define OOPS_UTIL_RESTART_H_

#include <map>
#include <string>

#include "eckit/config/Configuration.h"
#include "eckit/memory/NonCopyable.h"
#include "util/Printable.h"
#include "util/RestartHelper.h"

namespace util {

// -----------------------------------------------------------------------------

/*! @brief Restart class

 * Restart class allows to checkpoint/restart the Model; it interacts with the
 *  RestartHelper class which handles the restart logbook. The restart logbook
 *  is supposed to contain information required to restart the Model.
 *
 * The model registers the callback functions allowing to checkpoint/restart its
 *  components. Each component should have a unique key. The restart logbook
 *  shall contain configuration JSON block necessary to restart that component.
 */

class Restart : public util::Printable, private eckit::NonCopyable {
  using chkptCallbackFunction_ =
      std::function<void(eckit::LocalConfiguration& conf)>;
  using rstrtCallbackFunction_ =
      std::function<void(const eckit::Configuration& conf)>;

 public:
  static const std::string classname() { return "util::Restart"; }

  /// Constructor, destructor
  explicit Restart(const std::vector<std::string>& restarts);
  ~Restart() {}

  /// Connect checkpoint/restart callbacks
  void connectChkptCallback(const std::string& str,
                            const chkptCallbackFunction_& cb);

  void connectRstrtCallback(const std::string& str,
                            const rstrtCallbackFunction_& cb);

  /// Checkpoint/restart
  void checkpoint(const eckit::Configuration& conf,
                  const std::vector<std::string>& strs);
  void restart(const std::vector<std::string>& strs);

  /// Reset Logbook file
  void reset();

  /// Other
  const bool& isRestarted() const { return isRestarted_; }

 private:
  bool getRestartConf(const std::string& str, eckit::LocalConfiguration&);
  eckit::LocalConfiguration& newRestartConf(const std::string& str);
  void print(std::ostream& os) const;

  bool isRestarted_;
  std::map<std::string, eckit::LocalConfiguration> restartConfigs_;
  std::map<std::string, chkptCallbackFunction_> cbChkpt_;
  std::map<std::string, rstrtCallbackFunction_> cbRstrt_;
};

// =============================================================================

Restart::Restart(const std::vector<std::string>& restarts)
    : isRestarted_(false) {
  // Setup restart configurations if the restart file is available
  if (util::RestartHelper::isReadIn()) {
    // Restart file has been read in
    isRestarted_ = true;
    // Set restart configurations from the restart file
    for (auto const& restart : restarts) {
      if (util::RestartHelper::has(restart))
        util::RestartHelper::query<eckit::LocalConfiguration>(
            restart, newRestartConf(restart));
    }
    oops::Log::info() << *this << std::endl;
  }
}

// -----------------------------------------------------------------------------

void Restart::connectChkptCallback(const std::string& str,
                                   const chkptCallbackFunction_& cb) {
  auto search = cbChkpt_.find(str);
  ASSERT(search == cbChkpt_.end());
  cbChkpt_[str] = cb;
}

// -----------------------------------------------------------------------------

void Restart::connectRstrtCallback(const std::string& str,
                                   const rstrtCallbackFunction_& cb) {
  auto search = cbRstrt_.find(str);
  ASSERT(search == cbRstrt_.end());
  cbRstrt_[str] = cb;
}

// -----------------------------------------------------------------------------

void Restart::checkpoint(const eckit::Configuration& conf,
                         const std::vector<std::string>& strs) {
  if (conf.has("checkpoint")) {
    // Loop over checkpoint configuration blocks
    for (auto const& str : strs) {
      if (conf.has("checkpoint." + str)) {
        oops::Log::info() << "Restart::checkpoint " << str << ": " << std::endl;
        // Get checkpoint block configuration from checkpoint configuration
        eckit::LocalConfiguration lconf(conf, "checkpoint." + str);
        // Call the appropriate checkpoint callback function providing
        // the checkpoint block configuration
        cbChkpt_[str](lconf);
        util::RestartHelper::record<eckit::LocalConfiguration>(str, lconf);
      } else {
        oops::Log::info() << "Restart::checkpoint " << str
                          << ": not available in the config file" << std::endl;
      }
    }
    oops::Log::info()
        << "Restart::checkpoint write restart config to file starting"
        << std::endl;
    // Write out the Logbook
    util::RestartHelper::write();
    oops::Log::info() << "Restart::checkpoint write restart config to file done"
                      << std::endl;
  }
}

// -----------------------------------------------------------------------------

void Restart::reset() { util::RestartHelper::reset(); }

// -----------------------------------------------------------------------------

void Restart::restart(const std::vector<std::string>& strs) {
  if (isRestarted_) {
    // Loop over restart configuration blocks
    for (auto const& str : strs) {
      oops::Log::info() << "Restart::restart " << str << ": " << std::endl;
      // Get restart block configuration from restart configuration
      eckit::LocalConfiguration lconf;
      if (getRestartConf(str, lconf)) {
        // Call the appropriate restart callback function providing
        // the restart block configuration
        cbRstrt_[str](lconf);
      } else {
        oops::Log::info() << "Restart::restart configuration " << str
                          << " not available in the Logbook" << std::endl;
      }
    }
  }
}

// -----------------------------------------------------------------------------

bool Restart::getRestartConf(const std::string& str,
                             eckit::LocalConfiguration& conf) {
  auto search = restartConfigs_.find(str);
  if (search != restartConfigs_.end()) {
    conf = search->second;
    return true;
  } else {
    return false;
  }
}

// -----------------------------------------------------------------------------

eckit::LocalConfiguration& Restart::newRestartConf(const std::string& str) {
  auto search = restartConfigs_.find(str);
  ASSERT(search == restartConfigs_.end());
  eckit::LocalConfiguration lconf;
  restartConfigs_[str] = lconf;
  return restartConfigs_[str];
}

// -----------------------------------------------------------------------------

void Restart::print(std::ostream& os) const {
  if (isRestarted_) {
    os << std::endl;
    os << "--------------------------------------------------------------------"
          "-"
       << std::endl;
    os << "------------------------- Restart "
          "-----------------------------------"
       << std::endl;
    os << "--------------------------------------------------------------------"
          "-"
       << std::endl;
    os << "  Recorded configuration: " << std::endl;
    for (auto const& restart : restartConfigs_) {
      os << "    " << restart.first << ": " << restart.second << std::endl;
      os << "------------------------------------------------------------------"
            "---"
         << std::endl;
    }
    os << "----------- Restart configuration setup successful "
          "------------------"
       << std::endl;
    os << "--------------------------------------------------------------------"
          "-"
       << std::endl;
  }
}

// =============================================================================

}  // namespace util

#endif  // OOPS_UTIL_RESTART_H_