// Copyright © 2024-2026 Apple Inc.

#include "mlx/backend/cpu/jit_compiler.h"
#include "mlx/backend/common/utils.h"
#include "mlx/backend/cpu/compiled_preamble.h"

#include <algorithm>
#include <cstdlib>
#include <sstream>
#include <vector>

#include <fmt/format.h>

#if defined(_MSC_VER) && \
    (defined(_M_X64) || defined(_M_IX86) || defined(_M_AMD64))
#include <intrin.h>
#endif

namespace mlx::core {

namespace {

#if defined(_MSC_VER)

// Split string into array.
std::vector<std::string> str_split(const std::string& str, char delimiter) {
  std::vector<std::string> tokens;
  std::string token;
  std::istringstream tokenStream(str);
  while (std::getline(tokenStream, token, delimiter)) {
    tokens.push_back(token);
  }
  return tokens;
}

// Get path information about MSVC.
struct VisualStudioInfo {
  VisualStudioInfo() {
#ifdef _M_ARM64
    arch = "arm64";
#else
    arch = "x64";
#endif
    // Get path of Visual Studio.
    // Use -latest to get only the most recent installation when multiple
    // versions are installed, avoiding path concatenation issues.
    auto pf86 = std::getenv("ProgramFiles(x86)");
    if (!pf86) {
      throw std::runtime_error(
          "ProgramFiles(x86) environment variable not set.");
    }
    std::string vs_path = JitCompiler::exec(
        fmt::format(
            "\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\""
            " -latest -property installationPath 2>&1",
            pf86));
    if (vs_path.empty()) {
      throw std::runtime_error("Can not find Visual Studio.");
    }
    // Trim any trailing whitespace/newlines from the path
    vs_path.erase(
        std::find_if(
            vs_path.rbegin(),
            vs_path.rend(),
            [](unsigned char ch) { return !std::isspace(ch); })
            .base(),
        vs_path.end());
    // Read the envs from vcvarsall.
    std::string envs = JitCompiler::exec(
        fmt::format(
            "\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL 2>&1 && set",
            vs_path,
            arch));
    for (const std::string& line : str_split(envs, '\n')) {
      // Each line is in the format "ENV_NAME=values".
      auto pos = line.find_first_of('=');
      if (pos == std::string::npos || pos == 0 || pos == line.size() - 1)
        continue;
      std::string name = line.substr(0, pos);
      std::string value = line.substr(pos + 1);
      if (name == "INCLUDE") {
        includepaths = str_split(value, ';');
      } else if (name == "LIB") {
        libpaths = str_split(value, ';');
      } else if (name == "VCToolsInstallDir" || name == "VCTOOLSINSTALLDIR") {
        msvc_cl = fmt::format("{0}\\bin\\Host{1}\\{1}\\cl.exe", value, arch);
      }
    }

    // Check for clang-cl bundled with Visual Studio.
    std::string clang_cl_path = fmt::format(
        "{0}\\VC\\Tools\\Llvm\\{1}\\bin\\clang-cl.exe", vs_path, arch);
    if (std::filesystem::exists(clang_cl_path)) {
      clang_cl = clang_cl_path;
    }
  }

  std::string compiler(bool use_include) const {
    if (use_include) {
      // With installed headers, either cl.exe or clang-cl can parse the source.
      // Prefer cl.exe since it is part of the standard Visual Studio C++
      // toolchain.
      return !msvc_cl.empty() ? msvc_cl : clang_cl;
    }

    // Without installed headers, fall back to the embedded prebuilt preamble.
    // It was generated by the build compiler and may contain compiler-specific
    // builtins, so prefer the same compiler family when possible.
#ifdef __clang__
    return !clang_cl.empty() ? clang_cl : msvc_cl;
#else
    return !msvc_cl.empty() ? msvc_cl : clang_cl;
#endif
  }

  std::string arch;
  std::string msvc_cl;
  std::string clang_cl;
  std::vector<std::string> includepaths;
  std::vector<std::string> libpaths;
};

const VisualStudioInfo& GetVisualStudioInfo() {
  static VisualStudioInfo info;
  return info;
}

#endif // defined(_MSC_VER)

bool supports_avx2_target() {
#if defined(_MSC_VER)

#if defined(_M_X64) || defined(_M_IX86) || defined(_M_AMD64)
  int info[4];
  __cpuid(info, 0);
  if (info[0] < 7) {
    return false;
  }
  __cpuid(info, 1);
  bool os_avx = (info[2] & (1 << 27)) && (info[2] & (1 << 28));
  bool fma = info[2] & (1 << 12);
  bool f16c = info[2] & (1 << 29);
  if (!os_avx || !fma || !f16c) {
    return false;
  }
  unsigned long long xcr0 = _xgetbv(0);
  if ((xcr0 & 0x6) != 0x6) {
    return false;
  }
  __cpuidex(info, 7, 0);
  return info[1] & (1 << 5);
#else
  return false;
#endif // defined(_M_X64) || defined(_M_IX86) || defined(_M_AMD64)

#elif defined(__GNUC__) || defined(__clang__)

#if defined(__x86_64__) || defined(__i386__) || defined(__amd64__)
  __builtin_cpu_init();
  return __builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma") &&
      __builtin_cpu_supports("f16c");
#else
  return false;
#endif

#else
  return false;
#endif // defined(_MSC_VER)
}

} // namespace

const std::tuple<bool, std::string, std::string>& JitCompiler::get_preamble() {
  static auto preamble = []() -> std::tuple<bool, std::string, std::string> {
    // Check whether the headers are shipped with the binary, if so use the
    // preamble from the headers, otherwise use the prebuilt one embedded in
    // binary, which may not work with all compilers.
    auto root_dir = current_binary_dir();
#if !defined(_WIN32)
    root_dir = root_dir.parent_path();
#endif
    auto include_dir = root_dir / "include";
    if (std::filesystem::exists(
            include_dir / "mlx/backend/cpu/compiled_preamble.h")) {
      return std::make_tuple(
          true,
          include_dir.string(),
          "#include \"mlx/backend/cpu/compiled_preamble.h\"\n");
    } else {
      return std::make_tuple(false, "", get_prebuilt_preamble());
    }
  }();
  return preamble;
}

bool JitCompiler::available() {
#ifdef _MSC_VER
  static bool result = [] {
    try {
      bool use_include = std::get<0>(get_preamble());
      const auto& info = GetVisualStudioInfo();
      return !info.compiler(use_include).empty();
    } catch (...) {
      return false;
    }
  }();
  return result;
#else
#ifdef _WIN32
  static int result = std::system("g++ --version > NUL 2>&1");
#else
  static int result = std::system("g++ --version > /dev/null 2>&1");
#endif
  return result == 0;
#endif
}

std::string JitCompiler::build_command(
    const std::filesystem::path& dir,
    const std::string& source_file_name,
    const std::string& shared_lib_name) {
  auto& [use_include, include_dir, preamble] = get_preamble();
#ifdef _MSC_VER
  std::string compiler_flags;
  if (use_include) {
    compiler_flags += fmt::format(" /I \"{}\"", include_dir);
  }
  const VisualStudioInfo& info = GetVisualStudioInfo();
  std::string compiler = info.compiler(use_include);
  for (const std::string& include : info.includepaths) {
    compiler_flags += fmt::format(" /I \"{}\"", include);
  }
  std::string libpaths;
  for (const std::string& lib : info.libpaths) {
    libpaths += fmt::format(" /libpath:\"{}\"", lib);
  }
  // clang-cl accepts the same flags as cl.exe (/LD, /EHsc, etc.)
  // but we add -Wno-everything to suppress warnings from system and shipped
  // MLX headers.
  if (!info.clang_cl.empty() && compiler == info.clang_cl) {
    compiler_flags += " -Wno-everything";
  }
#ifdef __AVX2__
  if (supports_avx2_target()) {
    compiler_flags += " /arch:AVX2";
  }
#endif
  return fmt::format(
      "\""
      "cd /D \"{0}\" && "
      "\"{1}\" /LD /EHsc /MD /Ox /nologo /std:c++17{2} \"{3}\" "
      "/link /out:\"{4}\" {5} 2>&1"
      "\"",
      dir.string(),
      compiler,
      compiler_flags,
      source_file_name,
      shared_lib_name,
      libpaths);
#else
  std::string extra_flags;
  if (use_include) {
    extra_flags = fmt::format("-I \"{}\"", include_dir);
  }
#ifdef __AVX2__
  if (supports_avx2_target()) {
    if (!extra_flags.empty()) {
      extra_flags += " ";
    }
    extra_flags += "-mavx2 -mfma -mf16c";
  }
#endif
  return fmt::format(
      "g++ -std=c++17 -O3 -Wall -fPIC -shared {} \"{}\" -o \"{}\" 2>&1",
      extra_flags,
      (dir / source_file_name).string(),
      (dir / shared_lib_name).string());
#endif
}

std::string JitCompiler::exec(const std::string& cmd) {
#ifdef _MSC_VER
  FILE* pipe = _popen(cmd.c_str(), "r");
#else
  FILE* pipe = popen(cmd.c_str(), "r");
#endif
  if (!pipe) {
    throw std::runtime_error("popen() failed.");
  }
  char buffer[128];
  std::string ret;
  while (fgets(buffer, sizeof(buffer), pipe)) {
    ret += buffer;
  }
  // Trim trailing whitespace.
  ret.erase(
      std::find_if(
          ret.rbegin(),
          ret.rend(),
          [](unsigned char ch) { return !std::isspace(ch); })
          .base(),
      ret.end());

#ifdef _MSC_VER
  int status = _pclose(pipe);
#else
  int status = pclose(pipe);
#endif
  if (status == -1) {
    throw std::runtime_error("pclose() failed.");
  }
#if defined(_WIN32) || defined(__FreeBSD__)
  int code = status;
#else
  int code = WEXITSTATUS(status);
#endif
  if (code != 0) {
    throw std::runtime_error(
        fmt::format(
            "Failed to execute command with return code {0}: \"{1}\", "
            "the output is: {2}",
            code,
            cmd,
            ret));
  }
  return ret;
}

} // namespace mlx::core
