/// @author F.c.o Javier guinot Almenar <[email protected]>
#include "jvlke/ui/imgui.hpp"

#include "jvlke/render/swapchain.hpp"
#include "jvlke/gpu/device.hpp"
#include "jvlke/utils/logger.hpp"

#include "imgui/imgui.h"
#include "imgui/imgui_impl_glfw.h"
#include "imgui/imgui_impl_vulkan.h"

namespace jvlke
{
  namespace ui
  {
    struct Imgui::Data
    {
      VkDescriptorPool descriptor_pool_ = VK_NULL_HANDLE;
      VkRenderPass render_pass_ = VK_NULL_HANDLE;

      ImGuiContext *context_ = nullptr;
      ImGuiIO *io_ = nullptr;

      GLFWwindow *window_ = nullptr;

      std::weak_ptr<gpu::Device> device_ = {};
    };

    Imgui::Imgui() {}

    void Imgui::init(std::weak_ptr<render::Viewer> viewer, std::weak_ptr<gpu::Device> device, VkInstance instance)
    {
      if (data_)
        return;

      LOG_INFO("Making imgui context");

      data_ = std::make_shared<Imgui::Data>();
      data_->device_ = device;
      data_->window_ = viewer.lock()->getWindow().lock()->getWindowPtr();

      // Crear contexto ImGui
      /////////////////////////////////////////////////////////////////////////
      IMGUI_CHECKVERSION();
      data_->context_ = ImGui::CreateContext();
      ImGui::SetCurrentContext(data_->context_);
      data_->io_ = &ImGui::GetIO();

      // Inicializar backend de GLFW
      // Todos los callbacks se setean en la clase window
      if (!ImGui_ImplGlfw_InitForVulkan(data_->window_, false))
      {
        LOG_ERROR("Imgui GLFW initialization failed");
        this->free();
        return;
      }
      /////////////////////////////////////////////////////////////////////////

      // Crear descriptor pool para este contexto
      /////////////////////////////////////////////////////////////////////////
      VkDescriptorPoolSize pool_sizes[1] = {
          {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1},
      };
      VkDescriptorPoolCreateInfo pool_info = {};
      pool_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
      pool_info.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT;
      pool_info.maxSets = 1;
      pool_info.poolSizeCount = (uint32_t)(sizeof(pool_sizes) / sizeof(*pool_sizes));
      pool_info.pPoolSizes = pool_sizes;

      if (vkCreateDescriptorPool(data_->device_.lock()->getDevice(), &pool_info, nullptr, &data_->descriptor_pool_) != VK_SUCCESS)
      {
        LOG_ERROR("Failed to create descriptor pool for Imgui");
        this->free();
        return;
      }
      /////////////////////////////////////////////////////////////////////////

      // Crear render pass para este contexto
      /////////////////////////////////////////////////////////////////////////
      VkAttachmentDescription colorAttachment = {};
      colorAttachment.format = viewer.lock()->getRender().lock()->getSwapChainFormat();
      colorAttachment.samples = VK_SAMPLE_COUNT_1_BIT;
      colorAttachment.loadOp = VK_ATTACHMENT_LOAD_OP_LOAD;
      colorAttachment.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
      colorAttachment.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
      colorAttachment.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
      colorAttachment.initialLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
      colorAttachment.finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;

      // Configurar referencia para el color attachment
      VkAttachmentReference colorAttachmentRef = {};
      colorAttachmentRef.attachment = 0;
      colorAttachmentRef.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

      // Configurar descripción del depth attachment
      VkAttachmentDescription depthAttachment = {};
      depthAttachment.format = viewer.lock()->getRender().lock()->getSwapChainDepthFormat();
      depthAttachment.samples = VK_SAMPLE_COUNT_1_BIT;
      depthAttachment.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
      depthAttachment.storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
      depthAttachment.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
      depthAttachment.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
      depthAttachment.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
      depthAttachment.finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;

      // Configurar referencia para el depth attachment
      VkAttachmentReference depthAttachmentRef = {};
      depthAttachmentRef.attachment = 1;
      depthAttachmentRef.layout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;

      // Configurar subpass para usar el depth attachment
      VkSubpassDescription subpass = {};
      subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
      subpass.colorAttachmentCount = 1;
      subpass.pColorAttachments = &colorAttachmentRef;
      subpass.pDepthStencilAttachment = &depthAttachmentRef;

      // Configurar dependencias para garantizar las transiciones de layouts
      VkSubpassDependency dependency = {};
      dependency.srcSubpass = VK_SUBPASS_EXTERNAL;
      dependency.dstSubpass = 0;
      dependency.srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT | VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT;
      dependency.srcAccessMask = 0;
      dependency.dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT | VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT;
      dependency.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT | VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;

      // Crear el render pass con color y depth attachments
      std::vector<VkAttachmentDescription> attachments = {colorAttachment, depthAttachment};

      VkRenderPassCreateInfo renderPassInfo = {};
      renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
      renderPassInfo.attachmentCount = static_cast<uint32_t>(attachments.size());
      renderPassInfo.pAttachments = attachments.data();
      renderPassInfo.subpassCount = 1;
      renderPassInfo.pSubpasses = &subpass;
      renderPassInfo.dependencyCount = 1;
      renderPassInfo.pDependencies = &dependency;

      // Crear el render pass
      if (vkCreateRenderPass(device.lock()->getDevice(), &renderPassInfo, nullptr, &data_->render_pass_) != VK_SUCCESS)
      {
        LOG_ERROR("Failed to create ImGui render pass with depth attachment");
        this->free();
        return;
      }
      /////////////////////////////////////////////////////////////////////////

      // Configuración del backend (Vulkan)
      /////////////////////////////////////////////////////////////////////////
      ImGui_ImplVulkan_InitInfo init_info = {};
      init_info.Instance = instance;
      init_info.PhysicalDevice = device.lock()->getPhysicalDevice();
      init_info.Device = device.lock()->getDevice();
      init_info.QueueFamily = device.lock()->getQueueFamilyIndices().graphics_family_;
      init_info.Queue = device.lock()->getGraphicsQueue();
      init_info.DescriptorPool = data_->descriptor_pool_;
      init_info.RenderPass = data_->render_pass_;
      init_info.MinImageCount = render::SwapChain::MAX_FRAMES_IN_FLIGHT;
      init_info.ImageCount = (uint32_t)(viewer.lock()->getRender().lock()->getImageCount());
      init_info.CheckVkResultFn = [](VkResult err)
      {
        if (err != VK_SUCCESS)
          LOG_ERROR(("Vulkan error: " + std::to_string(err)).c_str());
      };

      if (!ImGui_ImplVulkan_Init(&init_info))
      {
        LOG_ERROR("ImGui Vulkan initialization failed");
        this->free();
        return;
      }
      /////////////////////////////////////////////////////////////////////////

      LOG_INFO("Imgui context created");
    }

    void Imgui::free()
    {
      if (!data_)
        return;

      LOG_INFO("Cleaning imgui context");

      ImGui::SetCurrentContext(data_->context_);
      ImGui_ImplGlfw_Shutdown();
      ImGui_ImplVulkan_Shutdown();
      ImGui::DestroyContext(data_->context_);

      if (data_->descriptor_pool_ != VK_NULL_HANDLE)
        vkDestroyDescriptorPool(data_->device_.lock()->getDevice(), data_->descriptor_pool_, nullptr);

      if (data_->render_pass_ != VK_NULL_HANDLE)
        vkDestroyRenderPass(data_->device_.lock()->getDevice(), data_->render_pass_, nullptr);

      data_.reset();
      LOG_INFO("Imgui context cleaned");
    }

    Imgui::~Imgui()
    {
      if (data_)
      {
        LOG_WARNING("Calling ui::Imgui::free in ui::Imgui::destructor, some error could occur");
        this->free();
      }
    }

    void *Imgui::getContext() const { return (void *)(data_->context_); }

    void Imgui::frameStart()
    {
      if (!data_ || !data_->context_)
        return;

      ImGui::SetCurrentContext(data_->context_);
      ImGui_ImplGlfw_NewFrame();
      ImGui_ImplVulkan_NewFrame();
      ImGui::NewFrame();
    }

    void Imgui::render(VkCommandBuffer command_buffer)
    {
      if (!data_ || !data_->context_)
        return;

      ImGui::SetCurrentContext(data_->context_);
      ImGui::Render();

      ImDrawData *draw_data = ImGui::GetDrawData();
      if (draw_data && draw_data->DisplaySize.x > 0 && draw_data->DisplaySize.y > 0)
        ImGui_ImplVulkan_RenderDrawData(draw_data, command_buffer);
    }

  } /* ui */
} /* jvlke */