#include "renderer/imgui_renderer.hpp"

#include <imgui.h>
#include <imgui_impl_glfw.h>
#include <imgui_impl_vulkan.h>
#include <imgui_internal.h>

#include <tracy/Tracy.hpp>

#include "resources/render_target.hpp"

namespace sophia {

namespace {
constexpr ImGuiDockNodeFlags DOCKSPACE_FLAGS =
    ImGuiDockNodeFlags_PassthruCentralNode;

void dockAll(const std::vector<std::string>& windows, ImGuiID nodeId) {
  for (const std::string& name : windows) {
    ImGui::DockBuilderDockWindow(name.c_str(), nodeId);
  }
}
}  // namespace

ImGuiRenderer::ImGuiRenderer(Window& window, Device& device,
                             vk::Format colorAttachmentFormat,
                             u32 minImageCount, ResourceRegistry& resources)
    : resources{resources} {
  // 1. Construct ImGui context
  IMGUI_CHECKVERSION();
  ImGui::CreateContext();

  ImGui_ImplGlfw_InitForVulkan(window.getGLFWwindow(), true);

  ImGui_ImplVulkan_InitInfo initInfo = {};
  device.populateImGuiInitInfo(initInfo);
  initInfo.DescriptorPoolSize =
      IMGUI_IMPL_VULKAN_MINIMUM_SAMPLED_IMAGE_POOL_SIZE + 5;
  initInfo.MinImageCount = minImageCount;
  initInfo.ImageCount = minImageCount;

  VkFormat rawColorAttachmentFormat =
      static_cast<VkFormat>(colorAttachmentFormat);
  initInfo.UseDynamicRendering = true;
  initInfo.PipelineInfoMain.PipelineRenderingCreateInfo.sType =
      VK_STRUCTURE_TYPE_PIPELINE_RENDERING_CREATE_INFO_KHR;
  initInfo.PipelineInfoMain.PipelineRenderingCreateInfo.colorAttachmentCount =
      1;
  initInfo.PipelineInfoMain.PipelineRenderingCreateInfo
      .pColorAttachmentFormats = &rawColorAttachmentFormat;

  ImGui_ImplVulkan_Init(&initInfo);

  // 2. Enable dark theme
  ImGui::StyleColorsDark();

  // 3. Enable docking / keyboard
  ImGuiIO& io = ImGui::GetIO();
  io.ConfigFlags |= ImGuiConfigFlags_DockingEnable;
  io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;
  io.ConfigFlags |= ImGuiConfigFlags_NavEnableGamepad;

  /**
   * When viewports are enabled we tweak WindowRounding/WindowBg so
   * platform windows can look identical to regular ones.
   */
  ImGuiStyle& style = ImGui::GetStyle();
  if (ImGui::GetIO().ConfigFlags & ImGuiConfigFlags_ViewportsEnable) {
    style.WindowRounding = 0.0f;
    style.Colors[ImGuiCol_WindowBg].w = 1.0f;
  }
}

ImGuiRenderer::~ImGuiRenderer() {
  for (const auto& [index, texture] : textures) {
    ImGui_ImplVulkan_RemoveTexture(texture.descriptorSet);
  }
  ImGui_ImplVulkan_Shutdown();
  ImGui_ImplGlfw_Shutdown();
  ImGui::DestroyContext();
}

void ImGuiRenderer::beginFrame() {
  ZoneScoped;
  releaseStaleTextures();

  ImGui_ImplVulkan_NewFrame();
  ImGui_ImplGlfw_NewFrame();
  ImGui::NewFrame();

  ImGuiID dockspaceId = ImGui::DockSpaceOverViewport(
      0, ImGui::GetMainViewport(), DOCKSPACE_FLAGS);
  applyLayoutIfNeeded(dockspaceId, DOCKSPACE_FLAGS);
}

ImTextureID ImGuiRenderer::getTextureId(RenderTargetHandle handle) {
  const RenderTarget& target = resources.renderTargetPool.get(handle);
  vk::ImageView view = target.getImage().getView();

  auto it = textures.find(handle.index);
  if (it != textures.end()) {
    if (it->second.generation == handle.generation &&
        it->second.version == target.getVersion()) {
      return (ImTextureID)it->second.descriptorSet;
    }
    ImGui_ImplVulkan_RemoveTexture(it->second.descriptorSet);
    textures.erase(it);
  }

  // ImGui owns the sampler; it only needs the view and the layout the image is
  // left in once the scene pass has finished.
  VkDescriptorSet descriptorSet = ImGui_ImplVulkan_AddTexture(
      static_cast<VkImageView>(view), VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
  textures[handle.index] = {handle.generation, target.getVersion(),
                            descriptorSet};
  return (ImTextureID)descriptorSet;
}

// Resizing and releasing a target both wait for the GPU before destroying its
// image, so nothing in flight can still be using a stale texture.
void ImGuiRenderer::releaseStaleTextures() {
  for (auto it = textures.begin(); it != textures.end();) {
    RenderTargetHandle handle{it->first, it->second.generation};
    RenderTargetPool& renderTargetPool = resources.renderTargetPool;

    bool stale =
        !renderTargetPool.isAlive(handle) ||
        renderTargetPool.get(handle).getVersion() != it->second.version;
    if (stale) {
      ImGui_ImplVulkan_RemoveTexture(it->second.descriptorSet);
      it = textures.erase(it);
    } else {
      ++it;
    }
  }
}

void ImGuiRenderer::setDefaultLayout(DockLayout layout) {
  defaultLayout = std::move(layout);
  seedIfUnsaved = true;
}

void ImGuiRenderer::resetLayout() { forceReset = true; }

void ImGuiRenderer::applyLayoutIfNeeded(ImGuiID dockspaceId,
                                        ImGuiDockNodeFlags flags) {
  if (defaultLayout.empty()) {
    return;
  }

  bool apply = forceReset;
  if (seedIfUnsaved) {
    // A layout restored from imgui.ini shows up as a split root node.
    ImGuiDockNode* root = ImGui::DockBuilderGetNode(dockspaceId);
    apply = apply || root == nullptr || !root->IsSplitNode();
    seedIfUnsaved = false;
  }
  if (!apply) {
    return;
  }
  forceReset = false;

  const ImVec2 size = ImGui::GetMainViewport()->Size;

  ImGui::DockBuilderRemoveNode(dockspaceId);
  ImGui::DockBuilderAddNode(dockspaceId, flags | ImGuiDockNodeFlags_DockSpace);
  ImGui::DockBuilderSetNodeSize(dockspaceId, size);

  ImGuiID center = dockspaceId;
  const ImGuiID left = ImGui::DockBuilderSplitNode(
      center, ImGuiDir_Left, defaultLayout.leftWidth, nullptr, &center);
  // The right split is taken from what's left after the left split.
  const float rightRatio =
      defaultLayout.rightWidth / (1.0f - defaultLayout.leftWidth);
  const ImGuiID right = ImGui::DockBuilderSplitNode(
      center, ImGuiDir_Right, rightRatio, nullptr, &center);
  const ImGuiID bottom = ImGui::DockBuilderSplitNode(
      center, ImGuiDir_Down, defaultLayout.bottomHeight, nullptr, &center);

  dockAll(defaultLayout.left, left);
  dockAll(defaultLayout.right, right);
  dockAll(defaultLayout.bottom, bottom);
  dockAll(defaultLayout.center, center);

  ImGui::DockBuilderFinish(dockspaceId);
}

void ImGuiRenderer::endFrame(vk::CommandBuffer commandBuffer) {
  ZoneScoped;
  ImGui::Render();
  ImDrawData* mainDrawData = ImGui::GetDrawData();

  ImGuiIO& io = ImGui::GetIO();
  if (io.ConfigFlags & ImGuiConfigFlags_ViewportsEnable) {
    ImGui::UpdatePlatformWindows();
    ImGui::RenderPlatformWindowsDefault();
  }

  ImGui_ImplVulkan_RenderDrawData(mainDrawData, commandBuffer);
}

}  // namespace sophia