/*
 * Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 *
 * SPDX-License-Identifier: MIT
 */

#include "rocprogram.hpp"

#include "utils/options.hpp"
#include "rockernel.hpp"
#include "rocurilocator.hpp"

#include <string>
#include <vector>
#include <cstring>
#include <fstream>
#include <sstream>
#include <iostream>
#include <iterator>

namespace amd::roc {

static inline const char* hsa_strerror(hsa_status_t status) {
  const char* str = nullptr;
  if (Hsa::status_string(status, &str) == HSA_STATUS_SUCCESS) {
    return str;
  }
  return "Unknown error";
}

Program::~Program() {
  // Destroy the executable.
  if (hsaExecutable_.handle != 0) {
    Hsa::executable_destroy(hsaExecutable_);
  }
  if (hsaCodeObjectReader_.handle != 0) {
    Hsa::code_object_reader_destroy(hsaCodeObjectReader_);
  }
  releaseClBinary();
}

Program::Program(roc::NullDevice& device, amd::Program& owner) : device::Program(device, owner) {
  hsaExecutable_.handle = 0;
  hsaCodeObjectReader_.handle = 0;
  isHIP_ = (owner.language() == amd::Program::HIP);
}

bool Program::initClBinary(char* binaryIn, size_t size) {
  // Save the original binary that isn't owned by ClBinary
  clBinary()->saveOrigBinary(binaryIn, size);

  char* bin = binaryIn;
  size_t sz = size;

  int encryptCode;

  char* decryptedBin;
  size_t decryptedSize;
  if (!clBinary()->decryptElf(binaryIn, size, &decryptedBin, &decryptedSize, &encryptCode)) {
    buildLog_ += "Decrypting ELF Failed\n";
    return false;
  }
  if (decryptedBin != nullptr) {
    // It is decrypted binary.
    bin = decryptedBin;
    sz = decryptedSize;
  }

  // Both 32-bit and 64-bit are allowed!
  if (!amd::Elf::isElfMagic(bin)) {
    // Invalid binary.
    delete[] decryptedBin;
    buildLog_ += "Elf Magic failed\n";
    return false;
  }

  clBinary()->setFlags(encryptCode);

  return clBinary()->setBinary(bin, sz, (decryptedBin != nullptr));
}


bool Program::defineGlobalVar(const char* name, void* dptr) {
  if (!device().isOnline()) {
    return false;
  }

  hsa_agent_t hsa_device = rocDevice().getBackendDevice();

  hsa_status_t status =
      Hsa::executable_agent_global_variable_define(hsaExecutable_, hsa_device, name, dptr);
  if (status != HSA_STATUS_SUCCESS) {
    buildLog_ += "Error: Could not define global variable : ";
    buildLog_ += hsa_strerror(status);
    buildLog_ += "\n";
  }

  return (status == HSA_STATUS_SUCCESS);
}

bool Program::createGlobalVarObj(amd::Memory** amd_mem_obj, void** device_pptr, size_t* bytes,
                                 const char* global_name) const {
  if (!device().isOnline()) {
    return false;
  }

  hsa_status_t status = HSA_STATUS_SUCCESS;
  const roc::Device* roc_device = nullptr;
  hsa_agent_t hsa_device;
  hsa_symbol_kind_t sym_type;
  hsa_executable_symbol_t global_symbol;

  if (amd_mem_obj == nullptr) {
    buildLog_ += "amd_mem_obj is null";
    buildLog_ += "\n";
    return false;
  }

  hsa_device = rocDevice().getBackendDevice();

  /* Find HSA Symbol by name */
  status =
      Hsa::executable_get_symbol_by_name(hsaExecutable_, global_name, &hsa_device, &global_symbol);
  if (status != HSA_STATUS_SUCCESS) {
    buildLog_ += "Error: Failed to find the Symbol by Name: ";
    buildLog_ += hsa_strerror(status);
    buildLog_ += "\n";
    return false;
  }

  /* Find HSA Symbol Type */
  status =
      Hsa::executable_symbol_get_info(global_symbol, HSA_EXECUTABLE_SYMBOL_INFO_TYPE, &sym_type);
  if (status != HSA_STATUS_SUCCESS) {
    buildLog_ += "Error: Failed to find the Symbol Type : ";
    buildLog_ += hsa_strerror(status);
    buildLog_ += "\n";
    return false;
  }

  /* Make sure symbol type is VARIABLE */
  if (sym_type != HSA_SYMBOL_KIND_VARIABLE) {
    buildLog_ += "Error: Symbol is not of type VARIABLE : ";
    buildLog_ += hsa_strerror(status);
    buildLog_ += "\n";
    return false;
  }

  /* Retrieve the size of the variable */
  status = Hsa::executable_symbol_get_info(global_symbol, HSA_EXECUTABLE_SYMBOL_INFO_VARIABLE_SIZE,
                                          bytes);

  if (status != HSA_STATUS_SUCCESS) {
    buildLog_ += "Error: Failed to retrieve the Symbol Size : ";
    buildLog_ += hsa_strerror(status);
    buildLog_ += "\n";
    return false;
  }

  // Handle size 0 symbols
  if (*bytes != 0) {
    // Find HSA Symbol Address
    status = Hsa::executable_symbol_get_info(
        global_symbol, HSA_EXECUTABLE_SYMBOL_INFO_VARIABLE_ADDRESS, device_pptr);
    if (status != HSA_STATUS_SUCCESS) {
      buildLog_ += "Error: Failed to find the Symbol Address : ";
      buildLog_ += hsa_strerror(status);
      buildLog_ += "\n";
      return false;
    }

    roc_device = &(rocDevice());
    *amd_mem_obj = new (roc_device->context())
        amd::Buffer(roc_device->context(), ROCCLR_MEM_INTERNAL_MEMORY, *bytes, *device_pptr);

    if (*amd_mem_obj == nullptr) {
      buildLog_ += "[OCL] Failed to create a mem object!";
      buildLog_ += "\n";
      return false;
    }

    if (!((*amd_mem_obj)->create(nullptr))) {
      buildLog_ += "[OCL] failed to create a svm hidden buffer!";
      buildLog_ += "\n";
      (*amd_mem_obj)->release();
      return false;
    }
  }
  return true;
}

bool Program::createBinary(amd::option::Options* options) {
  if (!clBinary()->createElfBinary(options->oVariables->BinEncrypt, type())) {
    LogError("Failed to create ELF binary image!");
    return false;
  }
  return true;
}

bool Program::createKernels(void* binary, size_t binSize, bool useUniformWorkGroupSize,
                            bool internalKernel) {
  // Find the size of global variables from the binary
  if (!FindGlobalVarSize(binary, binSize)) {
    buildLog_ += "Error: Cannot Find Global Var Sizes\n";
    return false;
  }

  for (const auto& kernelMeta : kernelMetadataMap_) {
    const std::string kernelName = kernelMeta.first;
    Kernel* aKernel = new roc::Kernel(kernelName, this);
    if (!aKernel->init()) {
      return false;
    }
    if (codeObjectVer() < 5) {
      aKernel->setUniformWorkGroupSize(useUniformWorkGroupSize);
    }
    aKernel->setInternalKernelFlag(internalKernel);
    addKernel(aKernel);
  }
  return true;
}

bool Program::setKernels(void* binary, size_t binSize, amd::Os::FileDesc fdesc,
                         size_t foffset, std::string uri) {
  // Stop compilation if it is an offline device - HSA runtime does not
  // support ISA compiled offline
  if (!device().isOnline()) {
    return true;
  }

  hsa_agent_t agent = rocDevice().getBackendDevice();
  hsa_status_t status;

  status = Hsa::executable_create_alt(HSA_PROFILE_FULL, HSA_DEFAULT_FLOAT_ROUNDING_MODE_DEFAULT,
                                     nullptr, &hsaExecutable_);
  if (status != HSA_STATUS_SUCCESS) {
    buildLog_ += "Error: Executable for AMD HSA Code Object isn't created: ";
    buildLog_ += hsa_strerror(status);
    buildLog_ += "\n";
    return false;
  }

  // If we received a file descriptor, hand the code object's file region to
  // rocr via the vendor extension so it can mmap it itself (and resolve the
  // URI directly from the fd). Fall back to from_memory when the fd is
  // invalid or the vendor extension isn't available in the linked rocr.
  // The fd is owned by setKernels in all cases: closed before returning.
  auto* create_from_file =
      Device::loaderExtensionTable()
          .hsa_ven_amd_loader_code_object_reader_create_from_file_with_offset_size;
  if (fdesc != amd::Os::FDescInit() && create_from_file != nullptr) {
    status = create_from_file(fdesc, foffset, binSize, &hsaCodeObjectReader_);
  } else {
    status = Hsa::code_object_reader_create_from_memory(binary, binSize, &hsaCodeObjectReader_);
  }
  if (fdesc != amd::Os::FDescInit()) {
    amd::Os::CloseFileHandle(fdesc);
  }
  if (status != HSA_STATUS_SUCCESS) {
    buildLog_ += "Error: AMD HSA Code Object Reader create failed: ";
    buildLog_ += hsa_strerror(status);
    buildLog_ += "\n";
    return false;
  }

  status = Hsa::executable_load_agent_code_object(hsaExecutable_, agent, hsaCodeObjectReader_,
                                                 nullptr, nullptr);
  if (status != HSA_STATUS_SUCCESS) {
    buildLog_ += "Error: AMD HSA Code Object loading failed: ";
    buildLog_ += hsa_strerror(status);
    buildLog_ += "\n";
    return false;
  }

  // Freeze the executable.
  status = Hsa::executable_freeze(hsaExecutable_, nullptr);
  if (status != HSA_STATUS_SUCCESS) {
    buildLog_ += "Error: Freezing the executable failed: ";
    buildLog_ += hsa_strerror(status);
    buildLog_ += "\n";
    return false;
  }

#if defined(__linux__) && defined(__clang__)
#if __has_feature(address_sanitizer)
  UriLocator::recordCodeObjects(hsaExecutable_);
#endif
#endif

  for (auto& kit : kernels()) {
    Kernel* kernel = static_cast<Kernel*>(kit.second);
    if (!kernel->postLoad()) {
      return false;
    }
  }
  return true;
}

}  // namespace amd::roc