#include "shader_compiler.hpp"

#include <cstring>
#include <filesystem>
#include <format>
#include <string_view>
#include <vector>

#include <slang-com-ptr.h>
#include <slang.h>

namespace trb::render {

void shader_compiler::destroy_slang(ISlangUnknown *object) {
  object->release();
}

shader_compiler::shader_compiler(std::filesystem::path root) : root(root) {
  const auto slang_global_session_description = SlangGlobalSessionDesc{};
  slang::IGlobalSession *global_session_ptr = nullptr;
  [[maybe_unused]] const auto result = slang::createGlobalSession(
      &slang_global_session_description, &global_session_ptr);
  global_session = slang_unique_ptr<slang::IGlobalSession>(global_session_ptr,
                                                           destroy_slang);

  auto paths_as_c_strings = std::vector<const char *>();
  for (const auto &path : search_paths) {
    paths_as_c_strings.push_back(path.generic_string().c_str());
  }

  auto profile = global_session->findProfile("spirv_1_4");

  auto target = slang::TargetDesc{
      .format = SlangCompileTarget::SLANG_SPIRV,
      .profile = profile,
      .flags = 0,
  };

  auto optimisation_level = slang::CompilerOptionEntry{
      .name = slang::CompilerOptionName::Optimization,
      .value = {
          .intValue0 = SlangOptimizationLevel::SLANG_OPTIMIZATION_LEVEL_MAXIMAL,
      }};
  auto debug_info = slang::CompilerOptionEntry{
      .name = slang::CompilerOptionName::DebugInformation,
      .value = {
          .intValue0 = SLANG_DEBUG_INFO_LEVEL_MAXIMAL,
      }};
  auto line_directive = slang::CompilerOptionEntry{
      .name = slang::CompilerOptionName::LineDirectiveMode,
      .value = {
          .intValue0 = SLANG_LINE_DIRECTIVE_MODE_SOURCE_MAP,
      }};
  auto name_main = slang::CompilerOptionEntry{
      .name = slang::CompilerOptionName::VulkanUseEntryPointName,
      .value = {
          .intValue0 = 1,
      }};

  auto infos =
      std::array{optimisation_level, debug_info, line_directive, name_main};

  auto slang_session_descriptor = slang::SessionDesc{};
  slang_session_descriptor.searchPaths = paths_as_c_strings.data();
  slang_session_descriptor.searchPathCount =
      static_cast<int>(paths_as_c_strings.size());
  slang_session_descriptor.targets = &target;
  slang_session_descriptor.targetCount = 1;
  slang_session_descriptor.defaultMatrixLayoutMode =
      SLANG_MATRIX_LAYOUT_COLUMN_MAJOR;
  slang_session_descriptor.compilerOptionEntries = infos.data();
  slang_session_descriptor.compilerOptionEntryCount =
      static_cast<std::int32_t>(infos.size());

  slang::ISession *session_ptr = nullptr;
  global_session->createSession(slang_session_descriptor, &session_ptr);
  session = slang_unique_ptr<slang::ISession>(session_ptr, destroy_slang);
}

std::vector<std::uint32_t>
shader_compiler::compile(const std::string_view module_name,
                         const std::string_view entry_name) {
  const auto initial_path = root / module_name;

  const auto module = load_module(initial_path.generic_string());
  const auto entry_point = find_entry_point(module, entry_name);
  const auto composed_program = create_composed_program(module, entry_point);
  const auto spirv_blob = compile_to_spirv(composed_program.get());

  auto data = std::vector<std::uint32_t>(spirv_blob->getBufferSize() / 4);
  std::memcpy(data.data(), spirv_blob->getBufferPointer(), data.size() * 4);

  return data;
}

slang::IModule *
shader_compiler::load_module(const std::string_view module_name) {
  slang::IBlob *module_blob_ptr = nullptr;
  auto module = session->loadModule(module_name.data(), &module_blob_ptr);
  auto module_blob =
      slang_unique_ptr<slang::IBlob>(module_blob_ptr, destroy_slang);

  if (!module) {
    throw_error("loading module", module_blob.get());
  }

  return module;
}

void shader_compiler::throw_error(const std::string_view context,
                                  slang::IBlob *diagnostic_blob) {
  const auto char_ptr =
      reinterpret_cast<const char *>(diagnostic_blob->getBufferPointer());

  const auto message = std::format("{}: {}", context, char_ptr);
  throw std::runtime_error(message);
}
slang::IEntryPoint *
shader_compiler::find_entry_point(slang::IModule *module,
                                  const std::string_view entry_point_name) {
  slang::IEntryPoint *entry_point;
  if (module->findEntryPointByName(entry_point_name.data(), &entry_point) !=
      SLANG_OK) {
    throw_error(std::format("finding entry point {}", entry_point_name));
  }
  return entry_point;
}

void shader_compiler::throw_error(const std::string_view context) {
  throw std::runtime_error(std::string(context));
}
shader_compiler::slang_unique_ptr<slang::IComponentType>
shader_compiler::create_composed_program(slang::IModule *module,
                                         slang::IEntryPoint *entry_point) {
  slang::IBlob *diagnostic_blob_ptr = nullptr;

  auto component_types = std::vector<slang::IComponentType *>();
  component_types.push_back(module);
  component_types.push_back(entry_point);

  slang::IComponentType *composed_program = nullptr;

  if (session->createCompositeComponentType(
          component_types.data(), component_types.size(), &composed_program,
          &diagnostic_blob_ptr) != SLANG_OK) {
    throw_error("creating composed program", diagnostic_blob_ptr);
  }

  return slang_unique_ptr<slang::IComponentType>(composed_program,
                                                 destroy_slang);
}
shader_compiler::slang_unique_ptr<slang::IBlob>
shader_compiler::compile_to_spirv(slang::IComponentType *composed_program) {
  slang::IBlob *diagnostic_blob_ptr = nullptr;

  slang::IBlob *spirv_blob = nullptr;

  composed_program->getEntryPointCode(0, 0, &spirv_blob, &diagnostic_blob_ptr);
  if (!spirv_blob) {
    throw_error("compiling to spirv", diagnostic_blob_ptr);
  }

  return slang_unique_ptr<slang::IBlob>(spirv_blob, destroy_slang);
}
void shader_compiler::reload_session() {
  auto paths_as_c_strings = std::vector<const char *>();
  for (const auto &path : search_paths) {
    paths_as_c_strings.push_back(path.generic_string().c_str());
  }

  auto profile = global_session->findProfile("spirv_1_4");

  auto target = slang::TargetDesc{
      .format = SlangCompileTarget::SLANG_SPIRV,
      .profile = profile,
      .flags = 0,
  };

  auto optimisation_level = slang::CompilerOptionEntry{
      .name = slang::CompilerOptionName::Optimization,
      .value = {
          .intValue0 = SlangOptimizationLevel::SLANG_OPTIMIZATION_LEVEL_NONE,
      }};
  auto debug_info = slang::CompilerOptionEntry{
      .name = slang::CompilerOptionName::DebugInformation,
      .value = {
          .intValue0 = SLANG_DEBUG_INFO_LEVEL_MAXIMAL,
      }};
  auto line_directive = slang::CompilerOptionEntry{
      .name = slang::CompilerOptionName::LineDirectiveMode,
      .value = {
          .intValue0 = SLANG_LINE_DIRECTIVE_MODE_SOURCE_MAP,
      }};
  auto name_main = slang::CompilerOptionEntry{
      .name = slang::CompilerOptionName::VulkanUseEntryPointName,
      .value = {
          .intValue0 = 1,
      }};

  auto infos =
      std::array{optimisation_level, debug_info, line_directive, name_main};

  auto slang_session_descriptor = slang::SessionDesc{};
  slang_session_descriptor.searchPaths = paths_as_c_strings.data();
  slang_session_descriptor.searchPathCount =
      static_cast<int>(paths_as_c_strings.size());
  slang_session_descriptor.targets = &target;
  slang_session_descriptor.targetCount = 1;
  slang_session_descriptor.defaultMatrixLayoutMode =
      SLANG_MATRIX_LAYOUT_COLUMN_MAJOR;
  slang_session_descriptor.compilerOptionEntries = infos.data();
  slang_session_descriptor.compilerOptionEntryCount =
      static_cast<std::int32_t>(infos.size());

  slang::ISession *session_ptr = nullptr;
  global_session->createSession(slang_session_descriptor, &session_ptr);
  session = slang_unique_ptr<slang::ISession>(session_ptr, destroy_slang);
}

} // namespace trb::render