#include "runtime/concurrency_controller.h"

#include <algorithm>
#include <cstring>
#include <limits>
#include <mutex>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#ifdef _WIN32
  #define WIN32_LEAN_AND_MEAN
  #define NOMINMAX
  #include <windows.h>
#else
  #include <cstdlib>
  #include <dlfcn.h>
#endif

namespace cf {

namespace {

struct EnvVarSnapshot {
  std::string name;
  bool was_set = false;
  std::string value;
};

std::optional<std::string> get_env_var(const char* name) {
#ifdef _WIN32
  DWORD needed = GetEnvironmentVariableA(name, nullptr, 0);
  if (needed == 0) {
    return std::nullopt;
  }
  std::string value;
  value.resize(needed - 1);
  if (needed > 1) {
    GetEnvironmentVariableA(name, value.data(), needed);
  }
  return value;
#else
  const char* v = std::getenv(name);
  if (!v) {
    return std::nullopt;
  }
  return std::string(v);
#endif
}

bool set_env_var(const char* name, const std::optional<std::string>& value) {
#ifdef _WIN32
  if (value.has_value()) {
    return SetEnvironmentVariableA(name, value->c_str()) != 0;
  }
  return SetEnvironmentVariableA(name, nullptr) != 0;
#else
  if (value.has_value()) {
    return ::setenv(name, value->c_str(), 1) == 0;
  }
  return ::unsetenv(name) == 0;
#endif
}

struct EnvVarGuard {
  explicit EnvVarGuard(const char* name, std::optional<std::string> new_value) : name_(name) {
    auto prev = get_env_var(name);
    snapshot_.name = name_;
    snapshot_.was_set = prev.has_value();
    snapshot_.value = prev.value_or(std::string());
    set_env_var(name, new_value);
  }

  EnvVarGuard(EnvVarGuard&& other) noexcept
      : name_(std::move(other.name_)), snapshot_(std::move(other.snapshot_)), active_(other.active_) {
    other.active_ = false;
  }

  EnvVarGuard& operator=(EnvVarGuard&& other) noexcept {
    if (this == &other) return *this;
    restore();
    name_ = std::move(other.name_);
    snapshot_ = std::move(other.snapshot_);
    active_ = other.active_;
    other.active_ = false;
    return *this;
  }

  ~EnvVarGuard() { restore(); }

  EnvVarGuard(const EnvVarGuard&) = delete;
  EnvVarGuard& operator=(const EnvVarGuard&) = delete;

 private:
  void restore() {
    if (!active_) return;
    if (snapshot_.was_set) {
      set_env_var(name_.c_str(), snapshot_.value);
    } else {
      set_env_var(name_.c_str(), std::nullopt);
    }
    active_ = false;
  }

  std::string name_;
  EnvVarSnapshot snapshot_;
  bool active_ = true;
};

struct EnvVarBundleGuard {
  void set(const char* name, const std::string& value) {
    guards_.emplace_back(name, value);
  }

  void unset(const char* name) {
    guards_.emplace_back(name, std::nullopt);
  }

 private:
  std::vector<EnvVarGuard> guards_;
};

using omp_get_max_threads_fn = int (*)();
using omp_set_num_threads_fn = void (*)(int);
using omp_get_dynamic_fn = int (*)();
using omp_set_dynamic_fn = void (*)(int);
using omp_get_nested_fn = int (*)();
using omp_set_nested_fn = void (*)(int);
using omp_get_max_active_levels_fn = int (*)();
using omp_set_max_active_levels_fn = void (*)(int);

struct OpenMpApi {
  omp_get_max_threads_fn get_max_threads = nullptr;
  omp_set_num_threads_fn set_num_threads = nullptr;
  omp_get_dynamic_fn get_dynamic = nullptr;
  omp_set_dynamic_fn set_dynamic = nullptr;
  omp_get_nested_fn get_nested = nullptr;
  omp_set_nested_fn set_nested = nullptr;
  omp_get_max_active_levels_fn get_max_active_levels = nullptr;
  omp_set_max_active_levels_fn set_max_active_levels = nullptr;

  bool available() const { return get_max_threads && set_num_threads; }
};

OpenMpApi load_openmp_api() {
  OpenMpApi api;

#ifdef _WIN32
  auto resolve_from = [&](HMODULE module, const char* symbol) -> void* {
    if (!module) return nullptr;
    return reinterpret_cast<void*>(GetProcAddress(module, symbol));
  };

  auto try_module = [&](const char* module_name) -> HMODULE {
    HMODULE mod = GetModuleHandleA(module_name);
    if (mod) return mod;
    return LoadLibraryA(module_name);
  };

  const char* candidates[] = {
    "vcomp140.dll",
    "vcomp120.dll",
    "vcomp110.dll",
    "vcomp100.dll",
    "libiomp5md.dll",
    "libiomp5.dll",
    "libomp.dll",
    "libgomp-1.dll",
  };

  for (const char* dll : candidates) {
    HMODULE mod = try_module(dll);
    if (!mod) continue;

    api.get_max_threads = reinterpret_cast<omp_get_max_threads_fn>(resolve_from(mod, "omp_get_max_threads"));
    api.set_num_threads = reinterpret_cast<omp_set_num_threads_fn>(resolve_from(mod, "omp_set_num_threads"));
    api.get_dynamic = reinterpret_cast<omp_get_dynamic_fn>(resolve_from(mod, "omp_get_dynamic"));
    api.set_dynamic = reinterpret_cast<omp_set_dynamic_fn>(resolve_from(mod, "omp_set_dynamic"));
    api.get_nested = reinterpret_cast<omp_get_nested_fn>(resolve_from(mod, "omp_get_nested"));
    api.set_nested = reinterpret_cast<omp_set_nested_fn>(resolve_from(mod, "omp_set_nested"));
    api.get_max_active_levels =
      reinterpret_cast<omp_get_max_active_levels_fn>(resolve_from(mod, "omp_get_max_active_levels"));
    api.set_max_active_levels =
      reinterpret_cast<omp_set_max_active_levels_fn>(resolve_from(mod, "omp_set_max_active_levels"));

    if (api.available()) {
      return api;
    }

    api = OpenMpApi{};
  }

  return api;
#else
  auto resolve = [&](const char* symbol) -> void* {
    return dlsym(RTLD_DEFAULT, symbol);
  };

  api.get_max_threads = reinterpret_cast<omp_get_max_threads_fn>(resolve("omp_get_max_threads"));
  api.set_num_threads = reinterpret_cast<omp_set_num_threads_fn>(resolve("omp_set_num_threads"));
  api.get_dynamic = reinterpret_cast<omp_get_dynamic_fn>(resolve("omp_get_dynamic"));
  api.set_dynamic = reinterpret_cast<omp_set_dynamic_fn>(resolve("omp_set_dynamic"));
  api.get_nested = reinterpret_cast<omp_get_nested_fn>(resolve("omp_get_nested"));
  api.set_nested = reinterpret_cast<omp_set_nested_fn>(resolve("omp_set_nested"));
  api.get_max_active_levels = reinterpret_cast<omp_get_max_active_levels_fn>(resolve("omp_get_max_active_levels"));
  api.set_max_active_levels = reinterpret_cast<omp_set_max_active_levels_fn>(resolve("omp_set_max_active_levels"));
  return api;
#endif
}

ParallelismModel normalized_model(const StepParallelism& p) {
  if (p.model == ParallelismModel::kUnknown) {
    return ParallelismModel::kExternal;
  }
  return p.model;
}

size_t compute_node_budget(const ConcurrencyController& controller, const PlanNode& node) {
  const auto& settings = controller.settings();
  if (!settings.respect_step_parallelism) {
    return 1;
  }

  const StepParallelism& p = node.step_def.parallelism;
  ParallelismModel model = normalized_model(p);
  if (model == ParallelismModel::kNone) {
    return 1;
  }

  size_t requested = settings.default_step_threads > 0 ? settings.default_step_threads : 1;
  if (p.preferred_threads.has_value() && p.preferred_threads.value() > 0) {
    requested = static_cast<size_t>(p.preferred_threads.value());
  }
  if (p.max_threads.has_value() && p.max_threads.value() > 0) {
    requested = std::min(requested, static_cast<size_t>(p.max_threads.value()));
  }

  requested = std::max<size_t>(1, requested);
  requested = std::min(requested, settings.max_total_threads);

  if (settings.deterministic) {
    if (!p.deterministic.has_value() || !p.deterministic.value()) {
      return 1;
    }
  }

  return requested;
}

}  // namespace

struct ConcurrencyController::ProcessScope::Impl {
  std::mutex* env_mutex = nullptr;
  std::unique_lock<std::mutex> lock;
  EnvVarBundleGuard env;

  explicit Impl(std::mutex& m) : env_mutex(&m), lock(m) {}

  ~Impl() {
    // Ensure env var restoration runs while holding the mutex (member destructors run after this body).
    if (env_mutex && !lock.owns_lock()) {
      lock.lock();
    }
  }
};

struct ConcurrencyController::StepScope::Impl {
  std::mutex* env_mutex = nullptr;
  std::unique_lock<std::mutex> lock;
  EnvVarBundleGuard env;

  CfRuntimeHintsV1 hints{};

  OpenMpApi omp{};
  bool used_omp = false;
  int prev_omp_max_threads = 0;
  int prev_omp_dynamic = 0;
  int prev_omp_nested = 0;
  int prev_omp_max_active_levels = 0;

  explicit Impl(std::mutex& m) : env_mutex(&m), lock(m, std::defer_lock) {}
};

namespace {
std::mutex& env_mutex_singleton() {
  static std::mutex m;
  return m;
}

OpenMpApi& openmp_singleton() {
  static OpenMpApi api = load_openmp_api();
  return api;
}
}  // namespace

ConcurrencyController::ConcurrencyController(const ResolvedConcurrencySettings& settings) : settings_(settings) {}

ResolvedConcurrencySettings ConcurrencyController::resolve(const ConcurrencyOptions& options, size_t hardware_threads) {
  ResolvedConcurrencySettings out;
  out.engine_threads = options.engine_threads > 0 ? options.engine_threads
                                                  : (hardware_threads > 0 ? hardware_threads : 1);
  out.max_total_threads = options.max_total_threads > 0 ? options.max_total_threads : out.engine_threads;
  out.default_step_threads = options.default_step_threads > 0 ? options.default_step_threads : 1;
  out.respect_step_parallelism = options.respect_step_parallelism;
  out.deterministic = options.deterministic;

  out.engine_threads = std::max<size_t>(1, out.engine_threads);
  out.max_total_threads = std::max<size_t>(1, out.max_total_threads);
  out.default_step_threads = std::max<size_t>(1, out.default_step_threads);
  return out;
}

ConcurrencyController::ProcessScope::ProcessScope(Impl* impl) : impl_(impl) {}

ConcurrencyController::ProcessScope::ProcessScope(ProcessScope&& other) noexcept : impl_(other.impl_) {
  other.impl_ = nullptr;
}

ConcurrencyController::ProcessScope& ConcurrencyController::ProcessScope::operator=(ProcessScope&& other) noexcept {
  if (this == &other) return *this;
  delete impl_;
  impl_ = other.impl_;
  other.impl_ = nullptr;
  return *this;
}

ConcurrencyController::ProcessScope::~ProcessScope() {
  delete impl_;
}

ConcurrencyController::StepScope::StepScope(Impl* impl) : impl_(impl) {}

ConcurrencyController::StepScope::StepScope(StepScope&& other) noexcept : impl_(other.impl_) {
  other.impl_ = nullptr;
}

ConcurrencyController::StepScope& ConcurrencyController::StepScope::operator=(StepScope&& other) noexcept {
  if (this == &other) return *this;
  delete impl_;
  impl_ = other.impl_;
  other.impl_ = nullptr;
  return *this;
}

ConcurrencyController::StepScope::~StepScope() {
  if (!impl_) {
    return;
  }

  if (impl_->used_omp && impl_->omp.available()) {
    if (impl_->omp.set_num_threads) {
      impl_->omp.set_num_threads(impl_->prev_omp_max_threads);
    }
    if (impl_->omp.set_dynamic) {
      impl_->omp.set_dynamic(impl_->prev_omp_dynamic);
    }
    if (impl_->omp.set_nested) {
      impl_->omp.set_nested(impl_->prev_omp_nested);
    }
    if (impl_->omp.set_max_active_levels) {
      impl_->omp.set_max_active_levels(impl_->prev_omp_max_active_levels);
    }
  }

  delete impl_;
}

ConcurrencyController::ProcessScope ConcurrencyController::enter_process_scope() const {
  auto* impl = new ProcessScope::Impl(env_mutex_singleton());

  // Safe default: constrain common nested parallel frameworks unless explicitly enabled.
  impl->env.set("OMP_NUM_THREADS", "1");
  impl->env.set("OMP_DYNAMIC", "FALSE");
  impl->env.set("MKL_NUM_THREADS", "1");
  impl->env.set("MKL_DYNAMIC", "FALSE");
  impl->env.set("OPENBLAS_NUM_THREADS", "1");
  impl->env.set("VECLIB_MAXIMUM_THREADS", "1");
  impl->env.set("NUMEXPR_NUM_THREADS", "1");

  if (settings_.deterministic) {
    // Conservative affinity defaults for determinism; frameworks may ignore if unsupported.
    impl->env.set("OMP_PROC_BIND", "TRUE");
    impl->env.set("OMP_PLACES", "cores");
  }

  impl->lock.unlock();
  return ProcessScope(impl);
}

void ConcurrencyController::assign_plan_budgets(ExecutionPlan& plan) const {
  OpenMpApi& omp = openmp_singleton();

  for (auto& node : plan.nodes) {
    node.thread_budget = compute_node_budget(*this, node);
    node.scheduler_cost = node.thread_budget;

    if (!settings_.respect_step_parallelism) {
      node.thread_budget = 1;
      node.scheduler_cost = 1;
      continue;
    }

    ParallelismModel model = normalized_model(node.step_def.parallelism);

    // If we must rely on process-global env vars for thread control, the safest option is
    // to run the step alone (scheduler_cost == max_total_threads).
    const bool wants_parallel = node.thread_budget > 1;
    const bool needs_env_exclusive =
      wants_parallel &&
      ((model == ParallelismModel::kOmp && !omp.available()) ||
       (model == ParallelismModel::kBlas) ||
       (model == ParallelismModel::kMkl));

    if (needs_env_exclusive) {
      node.scheduler_cost = settings_.max_total_threads;
    }

    node.thread_budget = std::max<size_t>(1, std::min(node.thread_budget, settings_.max_total_threads));
    node.scheduler_cost = std::max<size_t>(1, std::min(node.scheduler_cost, settings_.max_total_threads));
  }
}

ConcurrencyController::StepScope ConcurrencyController::enter_step_scope(const PlanNode& node, CfCallContext& ctx) const {
  auto* impl = new StepScope::Impl(env_mutex_singleton());

  impl->hints.version = 1;
  impl->hints.flags = settings_.deterministic ? CF_RUNTIME_HINT_DETERMINISTIC : 0u;
  impl->hints.thread_budget = static_cast<uint32_t>(
    std::min<size_t>(node.thread_budget, std::numeric_limits<uint32_t>::max()));
  impl->hints.reserved = 0u;
  ctx.user_data = &impl->hints;

  const StepParallelism& p = node.step_def.parallelism;
  ParallelismModel model = normalized_model(p);
  const size_t threads = std::max<size_t>(1, node.thread_budget);

  // When a step requires env var overrides, the scheduler must have granted exclusivity.
  const bool env_exclusive = (node.scheduler_cost == settings_.max_total_threads) && threads > 1;

  if (env_exclusive) {
    impl->lock.lock();

    const std::string t = std::to_string(threads);
    if (model == ParallelismModel::kOmp) {
      impl->env.set("OMP_NUM_THREADS", t);
      impl->env.set("OMP_DYNAMIC", "FALSE");
    } else if (model == ParallelismModel::kMkl) {
      impl->env.set("MKL_NUM_THREADS", t);
      impl->env.set("MKL_DYNAMIC", "FALSE");
      impl->env.set("OMP_NUM_THREADS", t);
      impl->env.set("OMP_DYNAMIC", "FALSE");
    } else if (model == ParallelismModel::kBlas) {
      impl->env.set("MKL_NUM_THREADS", t);
      impl->env.set("MKL_DYNAMIC", "FALSE");
      impl->env.set("OPENBLAS_NUM_THREADS", t);
      impl->env.set("VECLIB_MAXIMUM_THREADS", t);
      impl->env.set("NUMEXPR_NUM_THREADS", t);
      impl->env.set("OMP_NUM_THREADS", t);
      impl->env.set("OMP_DYNAMIC", "FALSE");
    }

    if (settings_.deterministic) {
      impl->env.set("OMP_PROC_BIND", "TRUE");
      impl->env.set("OMP_PLACES", "cores");
    }
  }

  // Thread-local OpenMP control (preferred), safe under concurrent engine execution.
  OpenMpApi& omp = openmp_singleton();
  if (model == ParallelismModel::kOmp && omp.available()) {
    impl->omp = omp;
    impl->used_omp = true;
    impl->prev_omp_max_threads = omp.get_max_threads ? omp.get_max_threads() : 0;
    impl->prev_omp_dynamic = omp.get_dynamic ? omp.get_dynamic() : 0;
    impl->prev_omp_nested = omp.get_nested ? omp.get_nested() : 0;
    impl->prev_omp_max_active_levels = omp.get_max_active_levels ? omp.get_max_active_levels() : 0;

    omp.set_num_threads(static_cast<int>(threads));
    if (omp.set_dynamic) {
      omp.set_dynamic(0);
    }
    if (omp.set_nested) {
      omp.set_nested(0);
    }
    if (omp.set_max_active_levels) {
      omp.set_max_active_levels(1);
    }
  }

  return StepScope(impl);
}

}  // namespace cf
