#include "imgui_renderer.hpp"

#include <imgui.h>
#include <imgui_impl_glfw.h>
#include <imgui_impl_vulkan.h>
#include <imgui_internal.h>  // DockBuilder API

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) {
  // 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() {
  ImGui_ImplVulkan_Shutdown();
  ImGui_ImplGlfw_Shutdown();
  ImGui::DestroyContext();
}

void ImGuiRenderer::beginFrame() {
  ImGui_ImplVulkan_NewFrame();
  ImGui_ImplGlfw_NewFrame();
  ImGui::NewFrame();

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

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) {
  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