// SPDX-License-Identifier: Apache-2.0
// Copyright (C) 2024 Advanced Micro Devices, Inc. All rights reserved.

#include "hip/config.h"
#include "hip/hip_runtime_api.h"

#include "module.h"

#include <sstream>

namespace {
static xrt::elf
create_elf(void* data, size_t size)
{
  return xrt::elf{data, size};
}

static xrt::uuid
register_xclbin(std::shared_ptr<xrt::core::hip::context> ctx, const xrt::xclbin& xclbin)
{
  auto xrt_device = ctx->get_xrt_device();
  return xrt_device.register_xclbin(xclbin);
}
}

namespace xrt::core::hip {

module_xclbin::
module_xclbin(std::shared_ptr<context> ctx, const std::string& file_name)
  : module{std::move(ctx)}
  , m_xrt_xclbin{file_name}
  , m_xrt_hw_ctx{m_ctx->get_xrt_device(), register_xclbin(m_ctx, m_xrt_xclbin)}
{}

module_xclbin::
module_xclbin(std::shared_ptr<context> ctx, void* data, size_t size)
  : module{std::move(ctx)}
  , m_xrt_xclbin{std::vector<char>{static_cast<char*>(data), static_cast<char*>(data) + size}}
  , m_xrt_hw_ctx{m_ctx->get_xrt_device(), register_xclbin(m_ctx, m_xrt_xclbin)}
{}

module_xclbin::
module_xclbin(std::shared_ptr<context> ctx, const std::string& file_name,
	      const xrt::hw_context::cfg_param_type& cfg_param)
  : module{std::move(ctx)}
  , m_xrt_xclbin{file_name}
  , m_xrt_hw_ctx{m_ctx->get_xrt_device(), register_xclbin(m_ctx, m_xrt_xclbin), cfg_param}
{}

module_xclbin::
module_xclbin(std::shared_ptr<context> ctx, void* data, size_t size,
	      const xrt::hw_context::cfg_param_type& cfg_param)
  : module{std::move(ctx)}
  , m_xrt_xclbin{std::vector<char>{static_cast<char*>(data), static_cast<char*>(data) + size}}
  , m_xrt_hw_ctx{m_ctx->get_xrt_device(), register_xclbin(m_ctx, m_xrt_xclbin), cfg_param}
{}

module_elf::
module_elf(module_xclbin* xclbin_module, const std::string& file_name)
  : module{xclbin_module->get_context()}
  , m_xclbin_module{xclbin_module}
  , m_xrt_elf{file_name}
  , m_xrt_module{m_xrt_elf}
{}

module_elf::
module_elf(module_xclbin* xclbin_module, void* data, size_t size)
  : module{xclbin_module->get_context()}
  , m_xclbin_module{xclbin_module}
  , m_xrt_elf{create_elf(data, size)}
  , m_xrt_module{m_xrt_elf}
{}

module_full_elf::
module_full_elf(std::shared_ptr<context> ctx, const std::string& file_name)
  : module{std::move(ctx)}
  , m_xrt_elf{file_name}
  , m_xrt_hw_ctx{m_ctx->get_xrt_device(), m_xrt_elf}
{}

module_full_elf::
module_full_elf(std::shared_ptr<context> ctx, const void* data, size_t size)
  : module{std::move(ctx)}
  , m_xrt_elf{data, size}
  , m_xrt_hw_ctx{m_ctx->get_xrt_device(), m_xrt_elf}
{}

module_full_elf::
module_full_elf(std::shared_ptr<context> ctx, const std::string& file_name,
                const xrt::hw_context::cfg_param_type& cfg_param)
  : module{std::move(ctx)}
  , m_xrt_elf{file_name}
  , m_xrt_hw_ctx{m_ctx->get_xrt_device(), m_xrt_elf, cfg_param, xrt::hw_context::access_mode::shared}
{}

module_full_elf::
module_full_elf(std::shared_ptr<context> ctx, void* data, size_t size,
                const xrt::hw_context::cfg_param_type& cfg_param)
  : module{std::move(ctx)}
  , m_xrt_elf{data, size}
  , m_xrt_hw_ctx{m_ctx->get_xrt_device(), m_xrt_elf, cfg_param, xrt::hw_context::access_mode::shared}
{}
function_handle
module_elf::
add_function(const std::string& name)
{
  return insert_in_map(function_cache,
                       std::make_shared<function>(this, m_xrt_module, name));
}

function_handle
module_full_elf::
add_function(const std::string& name)
{
  return insert_in_map(function_cache,
                       std::make_shared<function>(this, name));
}

static xrt::kernel
create_kernel(module_elf* elf_mod, const xrt::module& xrt_mod, const std::string& name)
{
  auto xclbin_mod = elf_mod->get_xclbin_module();
  if (!xclbin_mod)
    throw_invalid_resource_if(!module_cache.count(xclbin_mod),
			      "corresponding xclbin module of elf module is not available");

  return xrt::ext::kernel{xclbin_mod->get_hw_context(), xrt_mod, name};
}

function::
function(module_elf* elf_mod_hdl, const xrt::module& xrt_module, const std::string& name)
  : m_elf_module{elf_mod_hdl}
  , m_func_name{name}
  , m_xrt_kernel{create_kernel(m_elf_module, xrt_module, name)}
{}

function::
function(module_full_elf* full_elf_mod_hdl, const std::string& name)
  : m_full_elf_module{full_elf_mod_hdl}
  , m_func_name{name}
  , m_xrt_kernel{xrt::ext::kernel{m_full_elf_module->get_hw_context(), name}}
{}

// Global map of modules
//we should override clang-tidy warning by adding NOLINT since module_cache is non-const parameter
xrt_core::handle_map<module_handle, std::shared_ptr<module>> module_cache; //NOLINT
}
