/*
 * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <array>
#include <cassert>
#include <chrono>
#include <cstdint>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <memory>
#include <random>
#include <regex>
#include <sstream>
#include <string>
#include <vector>

#include "jit_utils.cuh"
#include "nvrtc.h"
#include "runtime.cuh"
#include "scheduler.cuh"
#include "tensorrt_llm/common/assert.h"
#include "tensorrt_llm/common/cudaUtils.h"
#include "tensorrt_llm/common/logger.h"

#ifdef _WIN32
#include <windows.h>
#endif

namespace deep_gemm::jit {

// Generate a unique ID for temporary directories to avoid collisions
inline std::string generateUniqueId() {
  // Use current time and random number to generate a unique ID
  static std::mt19937 gen(std::random_device{}());
  static std::uniform_int_distribution<> distrib(0, 999999);

  auto now = std::chrono::system_clock::now();
  auto now_ms = std::chrono::time_point_cast<std::chrono::milliseconds>(now);
  auto value = now_ms.time_since_epoch().count();

  // Use the static random generator
  int random_value = distrib(gen);

  return std::to_string(value) + "_" + std::to_string(random_value);
}

inline std::filesystem::path getDefaultUserDir() {
  static std::filesystem::path userDir;
  if (userDir.empty()) {
    char const* cacheDir = getenv("TRTLLM_DG_CACHE_DIR");
    if (cacheDir) {
      userDir = cacheDir;
      std::filesystem::create_directories(userDir);
    } else {
#ifdef _WIN32
      char const* appData = getenv("APPDATA");
      if (appData) {
        userDir = std::filesystem::path(appData) / "tensorrt_llm";
      } else {
        userDir = std::filesystem::temp_directory_path() / "tensorrt_llm";
      }
#else
      char const* homeDir = getenv("HOME");
      if (homeDir) {
        userDir = std::filesystem::path(homeDir) / ".tensorrt_llm";
      } else {
        userDir = std::filesystem::temp_directory_path() / "tensorrt_llm";
      }
#endif
    }
  }
  return userDir;
}

inline std::filesystem::path getTmpDir() { return getDefaultUserDir() / "tmp"; }

inline std::filesystem::path getCacheDir() { return getDefaultUserDir() / "cache"; }

inline std::string getNvccCompiler() {
  static std::string compiler;
  if (compiler.empty()) {
    // Check environment variable
    char const* envCompiler = getenv("TRTLLM_DG_NVCC_COMPILER");
    if (envCompiler) {
      compiler = envCompiler;
    } else {
      // Check CUDA_HOME
      char const* cudaHome = getenv("CUDA_HOME");
      if (cudaHome) {
        std::filesystem::path cudaPath(cudaHome);
#ifdef _WIN32
        compiler = (cudaPath / "bin" / "nvcc.exe").string();
#else
        compiler = (cudaPath / "bin" / "nvcc").string();
#endif
      } else {
// Default to system nvcc
#ifdef _WIN32
        compiler = "nvcc.exe";
#else
        compiler = "nvcc";
#endif
      }
    }
  }
  return compiler;
}

inline std::vector<std::filesystem::path>& getJitIncludeDirs() {
  static std::vector<std::filesystem::path> includeDirs;
  return includeDirs;
}

inline void setJitIncludeDirs(std::vector<std::filesystem::path> const& dirs) {
  static std::vector<std::filesystem::path>& includeDirs = getJitIncludeDirs();
  includeDirs = dirs;
}

inline std::string generateKernel(uint32_t const shape_n, uint32_t const shape_k,
                                  uint32_t const block_m, uint32_t const block_n,
                                  uint32_t const block_k, uint32_t const num_groups,
                                  uint32_t const num_stages, uint32_t const num_tma_multicast,
                                  deep_gemm::GemmType const gemm_type, bool swapAB = false) {
  constexpr uint32_t kNumTMAThreads = 128;
  constexpr uint32_t kNumMathThreadsPerGroup = 128;

  std::string input_type;
  if (!swapAB) {
    switch (gemm_type) {
      case deep_gemm::GemmType::Normal:
        input_type = "NormalSchedulerInput";
        break;
      case deep_gemm::GemmType::GroupedContiguous:
        input_type = "GroupedContiguousSchedulerInput";
        break;
      case deep_gemm::GemmType::GroupedMasked:
        input_type = "GroupedMaskedSchedulerInput";
        break;
      case deep_gemm::GemmType::GroupedWithOffset:
        input_type = "GroupedWithOffsetSchedulerInput";
        break;
      case deep_gemm::GemmType::StridedBatched:
        input_type = "StridedBatchedSchedulerInput";
        break;
      default:
        throw std::runtime_error("Unsupported gemm type");
    }
  } else {
    switch (gemm_type) {
      case deep_gemm::GemmType::Normal:
        input_type = "NormalSchedulerInputSwapAB";
        break;
      case deep_gemm::GemmType::GroupedWithOffset:
        input_type = "GroupedWithOffsetSchedulerInputSwapAB";
        break;
      default:
        throw std::runtime_error("Unsupported gemm type");
    }
  }

  // Modify kernel name based on swapAB to determine which kernel function to use
  std::string kernel_name = swapAB ? "fp8_gemm_kernel_swapAB" : "fp8_gemm_kernel";
  std::string scheduler_name = swapAB ? "SchedulerSelectorSwapAB" : "SchedulerSelector";

  // Create the kernel source code using raw string literal
  std::string code = R"(
#ifdef __CUDACC_RTC__
#ifndef NVRTC_JIT_COMPILATION
#define NVRTC_JIT_COMPILATION
#endif

#include <deep_gemm/nvrtc_std.cuh>

#else

#include <string>
#include <cuda.h>

#endif

#include <cuda_bf16.h>
#include <cuda_fp8.h>
#include <deep_gemm/nvrtc_cutlass.cuh>
#include <deep_gemm/fp8_gemm_impl.cuh>

using namespace deep_gemm;

using SchedulerType =
typename )" + scheduler_name +
                     R"(<GemmType::)" + gemm_type_to_string(gemm_type) + R"(, )" +
                     std::to_string(shape_n) + R"(, )" + std::to_string(shape_k) + R"(, )" +
                     std::to_string(block_m) + R"(, )" + std::to_string(block_n) + R"(, )" +
                     std::to_string(block_k) + R"(, )" + std::to_string(num_groups) + R"(, )" +
                     std::to_string(num_tma_multicast) + R"(>::type;

__global__ void dummy_kernel() {
  void *ptr = (void *)&)" +
                     kernel_name + R"(<)" + std::to_string(shape_n) + R"(, )" +
                     std::to_string(shape_k) + R"(, )" + std::to_string(block_m) + R"(, )" +
                     std::to_string(block_n) + R"(, )" + std::to_string(block_k) + R"(, )" +
                     std::to_string(num_groups) + R"(, )" + std::to_string(num_stages) + R"(, )" +
                     std::to_string(kNumTMAThreads) + R"(, )" +
                     std::to_string(kNumMathThreadsPerGroup) + R"(, )" +
                     std::to_string(num_tma_multicast) + R"(, SchedulerType, )" + input_type + R"(>;
}
)";

  return code;
}

/**
 * C++ implementation of the Compiler class
 * Compiles CUDA code into CUBINs
 */
class Compiler {
 public:
  // Get singleton instance
  static Compiler& getInstance() {
    static Compiler instance;
    return instance;
  }

  [[nodiscard]] bool isValid() const { return !getJitIncludeDirs().empty(); }

  // Set include directories before the singleton is initialized
  static void setIncludeDirs(std::vector<std::filesystem::path> const& dirs) {
    setJitIncludeDirs(dirs);
  }

  // Build function
  Runtime* build(uint32_t const shape_n, uint32_t const shape_k, uint32_t const block_m,
                 uint32_t const block_n, uint32_t const block_k, uint32_t const num_groups,
                 uint32_t const num_stages, uint32_t const num_tma_multicast,
                 deep_gemm::GemmType const gemm_type, bool swapAB = false) {
    int sm_version = tensorrt_llm::common::getSMVersion();
    if (sm_version != 90) {
      TLLM_THROW(
          "DeepGEMM only supports Hopper (SM90) architectures, but current device compute "
          "capability is %d.",
          sm_version);
    }

    // Build signature - simplified, no MD5 calculation
    std::string name = std::string(swapAB ? "gemm_swapAB_" : "gemm_") + std::to_string(shape_n) +
                       "_" + std::to_string(shape_k) + "_" + std::to_string(block_m) + "_" +
                       std::to_string(block_n) + "_" + std::to_string(block_k) + "_" +
                       std::to_string(num_groups) + "_" + std::to_string(num_stages) +
                       std::to_string(num_groups) + "_" + std::to_string(num_stages) + "_" +
                       std::to_string(num_tma_multicast) + "_" + gemm_type_to_string(gemm_type);
    std::filesystem::path path = getCacheDir() / name;

    // Check runtime cache or file system hit
    auto& runtimeCache = getGlobalRuntimeCache();
    Runtime* cachedRuntime = runtimeCache[path.string()];
    if (cachedRuntime != nullptr) {
      if (kJitDebugging) {
        TLLM_LOG_INFO("Using cached JIT runtime %s during build", name.c_str());
      }
      return cachedRuntime;
    }

    // Compiler flags
    std::vector<std::string> flags = {"-std=c++17",
                                      "--gpu-architecture=sm_90a",
                                      "--ptxas-options=-allow-expensive-optimizations=true",
                                      "--ptxas-options=--register-usage-level=10",
                                      "--diag-suppress=161,174,177,940",
                                      "-D__FORCE_INCLUDE_CUDA_FP16_HPP_FROM_FP16_H__=1",
                                      "-D__FORCE_INCLUDE_CUDA_BF16_HPP_FROM_BF16_H__=1"};

    if (kJitUseNvcc) {
      flags.push_back("-O3");
      flags.push_back("-cubin");
      flags.push_back("--expt-relaxed-constexpr");
      flags.push_back("--expt-extended-lambda");

      std::vector<std::string> cxxFlags = {"-fPIC", "-O3", "-Wno-deprecated-declarations",
                                           "-Wno-abi"};
      std::string cxxFlagsStr = "--compiler-options=";
      for (size_t i = 0; i < cxxFlags.size(); ++i) {
        cxxFlagsStr += cxxFlags[i];
        if (i < cxxFlags.size() - 1) {
          cxxFlagsStr += ",";
        }
      }
      flags.push_back(cxxFlagsStr);
    } else {
      flags.push_back("-default-device");
    }

    std::filesystem::path tmpPath = getTmpDir() / (name + "_" + generateUniqueId());
    std::filesystem::path cubinPath = path / kKernelName;
    std::filesystem::path tmpCubinPath = tmpPath / kKernelName;

    // Create the target directory if it doesn't exist
    if (kJitUseNvcc || kJitDumpCubin) {
      std::filesystem::create_directories(tmpPath);
      std::filesystem::create_directories(path);
    }

    for (auto const& dir : getJitIncludeDirs()) {
      flags.push_back("-I" + dir.string());
    }

    // Print options if debug enabled
    if (kJitDebugging) {
      TLLM_LOG_INFO("Compiling JIT runtime %s with options: ", name.c_str());
      for (auto const& flag : flags) {
        TLLM_LOG_INFO("%s ", flag.c_str());
      }
      TLLM_LOG_INFO("\n");
    }

    std::string code = generateKernel(shape_n, shape_k, block_m, block_n, block_k, num_groups,
                                      num_stages, num_tma_multicast, gemm_type, swapAB);

    if (kJitDebugging) {
      TLLM_LOG_INFO("Generated kernel code:\n%s", code.c_str());
    }

    if (kJitUseNvcc) {
      std::filesystem::path tmpSrcPath = tmpPath / "kernel.cu";

      // Write files
      std::ofstream srcFile(tmpSrcPath);
      srcFile << code;
      srcFile.close();

      // Build command
      std::vector<std::string> command = {getNvccCompiler(), tmpSrcPath.string(), "-o",
                                          tmpCubinPath.string()};
      command.insert(command.end(), flags.begin(), flags.end());

      // Execute command
      std::string cmd;
      for (auto const& arg : command) {
        cmd += arg + " ";
      }

      // Buffer to store the output
      std::array<char, 128> buffer;
      std::string result;

      // Time the compilation
      auto start = std::chrono::high_resolution_clock::now();

      // Open pipe to command
#ifdef _MSC_VER
      FILE* pipe = _popen(cmd.c_str(), "r");
#else
      FILE* pipe = popen(cmd.c_str(), "r");
#endif

      if (pipe) {
        // Read the output
        while (fgets(buffer.data(), buffer.size(), pipe) != nullptr) {
          result += buffer.data();
        }

// Close the pipe
#ifdef _MSC_VER
        _pclose(pipe);
#else
        pclose(pipe);
#endif

        // Output result if debug enabled
        if (kJitDebugging) {
          auto end = std::chrono::high_resolution_clock::now();
          auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
          TLLM_LOG_INFO("NVCC compilation took %d ms", duration.count());
          TLLM_LOG_INFO("Compilation log:\n%s", result.c_str());
        }
      }
    } else {
      nvrtcProgram prog;
      CHECK_NVRTC(nvrtcCreateProgram(&prog, code.c_str(), "kernel.cu", 0, nullptr, nullptr));

      std::vector<char const*> options;
      for (auto const& flag : flags) {
        options.push_back(flag.c_str());
      }

      // Time the compilation
      auto start = std::chrono::high_resolution_clock::now();
      nvrtcResult compileResult = nvrtcCompileProgram(prog, options.size(), options.data());

      if (kJitDebugging) {
        auto end = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        TLLM_LOG_INFO("NVRTC compilation took %d ms", duration.count());

        size_t logSize;
        CHECK_NVRTC(nvrtcGetProgramLogSize(prog, &logSize));
        std::vector<char> log(logSize);
        CHECK_NVRTC(nvrtcGetProgramLog(prog, log.data()));
        TLLM_LOG_INFO("Compilation log:\n%s", log.data());
      }

      // Check if compilation succeeded
      if (compileResult != NVRTC_SUCCESS) {
        TLLM_LOG_ERROR("NVRTC compilation failed");
        CHECK_NVRTC(nvrtcDestroyProgram(&prog));
        throw std::runtime_error("NVRTC compilation failed");
      }

      // Save CUBIN to a file
      size_t cubinSize;
      CHECK_NVRTC(nvrtcGetCUBINSize(prog, &cubinSize));
      std::vector<char> cubin(cubinSize);
      CHECK_NVRTC(nvrtcGetCUBIN(prog, cubin.data()));

      // Cache the runtime in memory by default
      if (!kJitDumpCubin) {
        auto runtime = std::make_unique<Runtime>(path.string(), cubin, gemm_type);
        Runtime* result = runtime.get();
        runtimeCache.set(path.string(), std::move(runtime));
        if (kJitDebugging) {
          TLLM_LOG_INFO("Successfully cached JIT runtime %s in memory", name.c_str());
        }
        return result;
      }

      std::ofstream cubinFile(tmpCubinPath.string(), std::ios::binary);
      cubinFile.write(cubin.data(), static_cast<std::streamsize>(cubinSize));
      cubinFile.close();
      CHECK_NVRTC(nvrtcDestroyProgram(&prog));
    }

    // Copy the source and compiled files to the cache directory
    try {
      // Rename (atomic operation) to final locations
      std::filesystem::rename(tmpCubinPath, cubinPath);
      if (kJitDebugging) {
        TLLM_LOG_INFO("Successfully copied kernel files to cache directory: %s",
                      path.string().c_str());
      }
    } catch (std::exception const& e) {
      TLLM_LOG_ERROR("Warning: Failed to copy kernel files to cache: %s", e.what());
    }

    // Clean up temporary directory after successful compilation
    try {
      std::filesystem::remove_all(tmpPath);
    } catch (std::exception const& e) {
      TLLM_LOG_ERROR("Warning: Failed to clean up temporary directory: %s", e.what());
    }

    // Create runtime and cache it
    auto runtime = std::make_unique<Runtime>(path.string(), std::vector<char>(), gemm_type);
    Runtime* result = runtime.get();
    runtimeCache.set(path.string(), std::move(runtime));
    return result;
  }

 private:
  // Private constructor for singleton pattern
  Compiler() {
    // Create necessary directories
    if (kJitUseNvcc || kJitDumpCubin) {
      std::filesystem::create_directories(getTmpDir());
      std::filesystem::create_directories(getCacheDir());
    }
  }

  // Delete copy constructor and assignment operator
  Compiler(Compiler const&) = delete;
  Compiler& operator=(Compiler const&) = delete;
};

// Global function to access the Compiler singleton
inline Compiler& getGlobalCompiler() { return Compiler::getInstance(); }

}  // namespace deep_gemm::jit
