/* SPDX-License-Identifier: Apache-2.0
 * Copyright 2017-2022 Blender Foundation */

#include <stdint.h>
#include <stdio.h>

#include <string>
#include <vector>

#include <OpenImageIO/argparse.h>
#include <OpenImageIO/filesystem.h>

#include "cuew.h"

#ifdef _MSC_VER
#  include <Windows.h>
#endif

using std::string;
using std::vector;

namespace std {
template<typename T> std::string to_string(const T &n)
{
  std::ostringstream s;
  s << n;
  return s.str();
}
}  // namespace std

class CompilationSettings {
 public:
  CompilationSettings()
      : target_arch(0), bits(64), verbose(false), fast_math(false), ptx_only(false)
  {
  }

  string cuda_toolkit_dir;
  string input_file;
  string output_file;
  string ptx_file;
  vector<string> defines;
  vector<string> includes;
  int target_arch;
  int bits;
  bool verbose;
  bool fast_math;
  bool ptx_only;
};

static bool compile_cuda(CompilationSettings &settings)
{
  const char *headers[] = {"stdlib.h", "float.h", "math.h", "stdio.h", "stddef.h"};
  const char *header_content[] = {"\n", "\n", "\n", "\n", "\n"};

  printf("Building %s\n", settings.input_file.c_str());

  string code;
  if (!OIIO::Filesystem::read_text_file(settings.input_file, code)) {
    fprintf(stderr, "Error: unable to read %s\n", settings.input_file.c_str());
    return false;
  }

  vector<string> options;
  for (size_t i = 0; i < settings.includes.size(); i++) {
    options.push_back("-I" + settings.includes[i]);
  }

  for (size_t i = 0; i < settings.defines.size(); i++) {
    options.push_back("-D" + settings.defines[i]);
  }
  options.push_back("-D__KERNEL_CUDA_VERSION__=" + std::to_string(cuewNvrtcVersion()));
  options.push_back("-arch=compute_" + std::to_string(settings.target_arch));
  options.push_back("--device-as-default-execution-space");
  options.push_back("-DCYCLES_CUBIN_CC");
  options.push_back("--std=c++11");
  if (settings.fast_math)
    options.push_back("--use_fast_math");

  nvrtcProgram prog;
  nvrtcResult result = nvrtcCreateProgram(&prog,
                                          code.c_str(),                      // buffer
                                          NULL,                              // name
                                          sizeof(headers) / sizeof(void *),  // numHeaders
                                          header_content,                    // headers
                                          headers);                          // includeNames

  if (result != NVRTC_SUCCESS) {
    fprintf(stderr, "Error: nvrtcCreateProgram failed (%d)\n\n", (int)result);
    return false;
  }

  /* Transfer options to a classic C array. */
  vector<const char *> opts(options.size());
  for (size_t i = 0; i < options.size(); i++) {
    opts[i] = options[i].c_str();
  }

  result = nvrtcCompileProgram(prog, options.size(), &opts[0]);

  if (result != NVRTC_SUCCESS) {
    fprintf(stderr, "Error: nvrtcCompileProgram failed (%d)\n\n", (int)result);

    size_t log_size;
    nvrtcGetProgramLogSize(prog, &log_size);

    vector<char> log(log_size);
    nvrtcGetProgramLog(prog, &log[0]);
    fprintf(stderr, "%s\n", &log[0]);

    return false;
  }

  /* Retrieve the ptx code. */
  size_t ptx_size;
  result = nvrtcGetPTXSize(prog, &ptx_size);
  if (result != NVRTC_SUCCESS) {
    fprintf(stderr, "Error: nvrtcGetPTXSize failed (%d)\n\n", (int)result);
    return false;
  }

  vector<char> ptx_code(ptx_size);
  result = nvrtcGetPTX(prog, &ptx_code[0]);
  if (result != NVRTC_SUCCESS) {
    fprintf(stderr, "Error: nvrtcGetPTX failed (%d)\n\n", (int)result);
    return false;
  }
  if (settings.ptx_only) {
    settings.ptx_file = settings.output_file;
  }
  else {
    /* Write a file in the temp folder with the ptx code. */
    settings.ptx_file = OIIO::Filesystem::temp_directory_path() + "/" +
                        OIIO::Filesystem::unique_path();
  }
  FILE *f = fopen(settings.ptx_file.c_str(), "wb");
  fwrite(&ptx_code[0], 1, ptx_size, f);
  fclose(f);

  return true;
}

static bool link_ptxas(CompilationSettings &settings)
{
  string cudapath = "";
  if (settings.cuda_toolkit_dir.size())
    cudapath = settings.cuda_toolkit_dir + "/bin/";

  string ptx = "\"" + cudapath + "ptxas\" " + settings.ptx_file + " -o " + settings.output_file +
               " --gpu-name sm_" + std::to_string(settings.target_arch) + " -m" +
               std::to_string(settings.bits);

  if (settings.verbose) {
    ptx += " --verbose";
    printf("%s\n", ptx.c_str());
  }

  int pxresult = system(ptx.c_str());
  if (pxresult) {
    fprintf(stderr, "Error: ptxas failed (%d)\n\n", pxresult);
    return false;
  }

  if (!OIIO::Filesystem::remove(settings.ptx_file)) {
    fprintf(stderr, "Error: removing %s\n\n", settings.ptx_file.c_str());
  }

  return true;
}

static bool init(CompilationSettings &settings)
{
#ifdef _MSC_VER
  if (settings.cuda_toolkit_dir.size()) {
    SetDllDirectory((settings.cuda_toolkit_dir + "/bin").c_str());
  }
#else
  (void)settings;
#endif

  int cuewresult = cuewInit(CUEW_INIT_NVRTC);
  if (cuewresult != CUEW_SUCCESS) {
    fprintf(stderr, "Error: cuew init fialed (0x%d)\n\n", cuewresult);
    return false;
  }

  if (cuewNvrtcVersion() < 80) {
    fprintf(stderr, "Error: only cuda 8 and higher is supported, %d\n\n", cuewCompilerVersion());
    return false;
  }

  if (!nvrtcCreateProgram) {
    fprintf(stderr, "Error: nvrtcCreateProgram not resolved\n");
    return false;
  }

  if (!nvrtcCompileProgram) {
    fprintf(stderr, "Error: nvrtcCompileProgram not resolved\n");
    return false;
  }

  if (!nvrtcGetProgramLogSize) {
    fprintf(stderr, "Error: nvrtcGetProgramLogSize not resolved\n");
    return false;
  }

  if (!nvrtcGetProgramLog) {
    fprintf(stderr, "Error: nvrtcGetProgramLog not resolved\n");
    return false;
  }

  if (!nvrtcGetPTXSize) {
    fprintf(stderr, "Error: nvrtcGetPTXSize not resolved\n");
    return false;
  }

  if (!nvrtcGetPTX) {
    fprintf(stderr, "Error: nvrtcGetPTX not resolved\n");
    return false;
  }

  return true;
}

static bool parse_parameters(int argc, const char **argv, CompilationSettings &settings)
{
  OIIO::ArgParse ap;
  ap.options("Usage: cycles_cubin_cc [options]",
             "-target %d",
             &settings.target_arch,
             "target shader model",
             "-m %d",
             &settings.bits,
             "Cuda architecture bits",
             "-i %s",
             &settings.input_file,
             "Input source filename",
             "-o %s",
             &settings.output_file,
             "Output cubin filename",
             "-I %L",
             &settings.includes,
             "Add additional includepath",
             "-D %L",
             &settings.defines,
             "Add additional defines",
             "-ptx",
             &settings.ptx_only,
             "emit PTX code",
             "-v",
             &settings.verbose,
             "Use verbose logging",
             "--use_fast_math",
             &settings.fast_math,
             "Use fast math",
             "-cuda-toolkit-dir %s",
             &settings.cuda_toolkit_dir,
             "path to the cuda toolkit binary directory",
             NULL);

  if (ap.parse(argc, argv) < 0) {
    fprintf(stderr, "%s\n", ap.geterror().c_str());
    ap.usage();
    return false;
  }

  if (!settings.output_file.size()) {
    fprintf(stderr, "Error: Output file not set(-o), required\n\n");
    return false;
  }

  if (!settings.input_file.size()) {
    fprintf(stderr, "Error: Input file not set(-i, required\n\n");
    return false;
  }

  if (!settings.target_arch) {
    fprintf(stderr, "Error: target shader model not set (-target), required\n\n");
    return false;
  }

  return true;
}

int main(int argc, const char **argv)
{
  CompilationSettings settings;

  if (!parse_parameters(argc, argv, settings)) {
    fprintf(stderr, "Error: invalid parameters, exiting\n");
    exit(EXIT_FAILURE);
  }

  if (!init(settings)) {
    fprintf(stderr, "Error: initialization error, exiting\n");
    exit(EXIT_FAILURE);
  }

  if (!compile_cuda(settings)) {
    fprintf(stderr, "Error: compilation error, exiting\n");
    exit(EXIT_FAILURE);
  }

  if (!settings.ptx_only) {
    if (!link_ptxas(settings)) {
      exit(EXIT_FAILURE);
    }
  }

  return 0;
}