#include <utils/PugsUtils.hpp>

#include <dev/ParallelChecker.hpp>
#include <utils/BacktraceManager.hpp>
#include <utils/BuildInfo.hpp>
#include <utils/CommunicatorManager.hpp>
#include <utils/ConsoleManager.hpp>
#include <utils/ExecutionStatManager.hpp>
#include <utils/FPEManager.hpp>
#include <utils/Messenger.hpp>
#include <utils/PETScWrapper.hpp>
#include <utils/ReproducibleSumManager.hpp>
#include <utils/RevisionInfo.hpp>
#include <utils/SLEPcWrapper.hpp>
#include <utils/SignalManager.hpp>
#include <utils/checkpointing/ResumingManager.hpp>
#include <utils/pugs_build_info.hpp>

#include <rang.hpp>

#include <Kokkos_Core.hpp>

#include <CLI/CLI.hpp>

#include <iostream>
#include <thread>

std::string
pugsVersion()
{
  std::stringstream os;

  os << "pugs version: " << rang::style::bold << RevisionInfo::version() << rang::style::reset << '\n';

  os << "-------------------- " << rang::fg::green << "git info" << rang::fg::reset << " -------------------------"
     << '\n';
  os << "tag:  " << rang::style::bold << RevisionInfo::gitTag() << rang::style::reset << '\n';
  os << "HEAD: " << rang::style::bold << RevisionInfo::gitHead() << rang::style::reset << '\n';
  os << "hash: " << rang::style::bold << RevisionInfo::gitHash() << rang::style::reset << "  (";

  // LCOV_EXCL_START Cannot cover both situations at same time
  if (RevisionInfo::gitIsClean()) {
    os << rang::fgB::green << "clean" << rang::fg::reset;
  } else {
    os << rang::fgB::red << "dirty" << rang::fg::reset;
  }
  // LCOV_EXCL_STOP
  os << ")\n";
  os << "-------------------------------------------------------";

  return os.str();
}

std::string
pugsBuildInfo()
{
  std::ostringstream os;

  os << "-------------------- " << rang::fg::green << "build info" << rang::fg::reset << " -----------------------"
     << '\n';
  os << "type:     " << rang::style::bold << BuildInfo::type() << rang::style::reset << '\n';
  os << "compiler: " << rang::style::bold << BuildInfo::compiler() << rang::style::reset << '\n';
  os << "kokkos:   " << rang::style::bold << BuildInfo::kokkosDevices() << rang::style::reset << '\n';
  os << "MPI:      " << rang::style::bold << BuildInfo::mpiLibrary() << rang::style::reset << '\n';
  os << "PETSc:    " << rang::style::bold << BuildInfo::petscLibrary() << rang::style::reset << '\n';
  os << "SLEPc:    " << rang::style::bold << BuildInfo::slepcLibrary() << rang::style::reset << '\n';
  os << "HDF5:     " << rang::style::bold << BuildInfo::hdf5Library() << rang::style::reset << '\n';
  os << "SLURM:    " << rang::style::bold << BuildInfo::slurmLibrary() << rang::style::reset << '\n';
  os << "-------------------------------------------------------";

  return os.str();
}

void
setDefaultOMPEnvironment()
{
  if constexpr (std::string_view{PUGS_BUILD_KOKKOS_DEVICES}.find("OpenMP") != std::string_view::npos) {
    setenv("OMP_PROC_BIND", "spread,close", 0);
    setenv("OMP_PLACES", "threads", 0);
  }
}

// LCOV_EXCL_START

// This function cannot be unit-tested: run once when pugs starts

std::string
initialize(int& argc, char* argv[])
{
  bool enable_fpe      = true;
  bool enable_signals  = true;
  int nb_threads       = -1;
  bool parallel_output = false;

  ParallelChecker::Mode pc_mode = ParallelChecker::Mode::automatic;
  std::string pc_filename       = ParallelChecker::instance().filename();

  std::string filename;
  {
    CLI::App app{"pugs help"};

    app.add_option("filename", filename, "pugs script file")->check(CLI::ExistingFile)->required();

    bool is_resuming = false;
    app.add_flag("--resume", is_resuming, "Resume at checkpoint");

    app.set_version_flag("-v,--version", []() {
      ConsoleManager::init(true);
      std::stringstream os;
      os << pugsVersion() << '\n' << pugsBuildInfo();
      return os.str();
    });

    app.add_option("--threads", nb_threads, "Number of Kokkos threads")
      ->check(CLI::Range(1, std::numeric_limits<decltype(nb_threads)>::max()));

    bool enable_color = true;
    app.add_flag("--color,!--no-color", enable_color, "Colorize console output [default: true]");

    app.add_flag("--fpe,!--no-fpe", enable_fpe, "Trap floating point exceptions [default: true]");

    bool show_preamble = true;
    app.add_flag("--preamble,!--no-preamble", show_preamble, "Show execution info preamble [default: true]");

    bool print_exec_stat = true;
    app.add_flag("--exec-stat,!--no-exec-stat", print_exec_stat,
                 "Display memory and CPU usage after execution [default: true]");

    bool show_backtrace = false;
    app.add_flag("-b,--backtrace,!--no-backtrace", show_backtrace, "Show backtrace on failure [default: false]");

    app.add_flag("--signal,!--no-signal", enable_signals, "Catches signals [default: true]");

    bool pause_on_error = false;
    app.add_flag("-p,--pause-on-error", pause_on_error, "Pause for debugging on unexpected error [default: false]");

    bool reproducible_sums = true;
    app.add_flag("--reproducible-sums,!--no-reproducible-sums", reproducible_sums,
                 "Special treatment of array sums to ensure reproducibility [default: true]");

    app.add_flag("--parallel-output", parallel_output, "All MPI processes output to console [default: false]");

    std::map<std::string, ParallelChecker::Mode> pc_mode_map{{"auto", ParallelChecker::Mode::automatic},
                                                             {"write", ParallelChecker::Mode::write},
                                                             {"read", ParallelChecker::Mode::read}};
    app
      .add_option("--parallel-checker-mode", pc_mode,
                  "Parallel checker mode (auto: sequential write/parallel read) [default: auto]")
      ->transform(CLI::CheckedTransformer(pc_mode_map));

    app.add_option("--parallel-checker-file", pc_filename,
                   "Parallel checker filename   [default: " + pc_filename + "]");

    int mpi_split_color = -1;
    app.add_option("--mpi-split-color", mpi_split_color, "Sets the MPI split color value (for MPMD applications)")
      ->check(CLI::Range(0, std::numeric_limits<decltype(mpi_split_color)>::max()));

    std::atexit([]() { std::cout << rang::style::reset; });
    try {
      app.parse(argc, argv);
    }
    catch (const CLI::ParseError& e) {
      // Stupid trick to avoid repetition of error messages in parallel
      parallel::Messenger::create(argc, argv);
      parallel::Messenger::destroy();
      std::exit(app.exit(e, std::cout, std::cerr));
    }

    if (app.count("--mpi-split-color") > 0) {
      CommunicatorManager::setSplitColor(mpi_split_color);
    }

    ResumingManager::getInstance().setIsResuming(is_resuming);
    ResumingManager::getInstance().setFilename(filename);

    ExecutionStatManager::getInstance().setPrint(print_exec_stat);
    BacktraceManager::setShow(show_backtrace);
    ConsoleManager::setShowPreamble(show_preamble);
    ConsoleManager::init(enable_color);
    SignalManager::setPauseForDebug(pause_on_error);
    ReproducibleSumManager::setReproducibleSums(reproducible_sums);
  }

  parallel::Messenger::create(argc, argv, parallel_output);

  PETScWrapper::initialize(argc, argv);
  SLEPcWrapper::initialize(argc, argv);

  FPEManager::init(enable_fpe);
  SignalManager::init(enable_signals);

  setDefaultOMPEnvironment();
  {
    Kokkos::InitializationSettings args;
    args.set_num_threads(nb_threads);
    args.set_device_id(-1);
    args.set_disable_warnings(true);

    Kokkos::initialize(args);
  }

  ParallelChecker::instance().setMode(pc_mode);
  ParallelChecker::instance().setFilename(pc_filename);

  if (ConsoleManager::showPreamble()) {
    std::cout << "----------------- " << rang::fg::green << "pugs exec info" << rang::fg::reset
              << " ----------------------" << '\n';

    std::cout << rang::style::bold;
#ifdef PUGS_HAS_MPI
    if (CommunicatorManager::hasSplitColor()) {
      std::cout << "MPI number of global ranks " << parallel::Messenger::getInstance().globalSize() << '\n';
      std::cout << "MPI local pugs communication world\n";
      std::cout << " - number of local ranks  " << parallel::size() << '\n';
      std::cout << " - local comm world color " << CommunicatorManager::splitColor() << '\n';
    } else {
      std::cout << "MPI number of ranks " << parallel::size() << '\n';
    }
#else    // PUGS_HAS_MPI
    std::cout << "Sequential build\n";
#endif   // PUGS_HAS_MPI
    std::cout << "Number of threads " << Kokkos::DefaultHostExecutionSpace::concurrency() << " / "
              << std::max(std::thread::hardware_concurrency(), 1u) << '\n';
    std::cout << rang::style::reset;
    std::cout << "-------------------------------------------------------\n";
  }

  return filename;
}

// LCOV_EXCL_STOP

// LCOV_EXCL_START
// This function cannot be unit-tested: run once when pugs stops

void
finalize()
{
  Kokkos::finalize();
  SLEPcWrapper::finalize();
  PETScWrapper::finalize();
  parallel::Messenger::destroy();
}

// LCOV_EXCL_STOP