#include "XCEngine/RHI/Vulkan/VulkanFramebuffer.h"

#include "XCEngine/RHI/Vulkan/VulkanRenderPass.h"
#include "XCEngine/RHI/Vulkan/VulkanResourceView.h"
#include "XCEngine/RHI/Vulkan/VulkanTexture.h"

#include <vector>

namespace XCEngine {
namespace RHI {

VulkanFramebuffer::~VulkanFramebuffer() {
    Shutdown();
}

bool VulkanFramebuffer::Initialize(VkDevice device, RHIRenderPass* renderPass, uint32_t width, uint32_t height,
                                   uint32_t colorAttachmentCount, RHIResourceView** colorAttachments,
                                   RHIResourceView* depthStencilAttachment) {
    if (device == VK_NULL_HANDLE || renderPass == nullptr || width == 0 || height == 0) {
        return false;
    }
    if (colorAttachmentCount > 0 && colorAttachments == nullptr) {
        return false;
    }

    auto* vulkanRenderPass = static_cast<VulkanRenderPass*>(renderPass);
    if (vulkanRenderPass == nullptr || vulkanRenderPass->GetRenderPass() == VK_NULL_HANDLE) {
        return false;
    }

    Shutdown();

    m_device = device;
    m_renderPass = vulkanRenderPass;
    m_width = width;
    m_height = height;
    if (colorAttachmentCount != vulkanRenderPass->GetColorAttachmentCount()) {
        Shutdown();
        return false;
    }
    if ((vulkanRenderPass->HasDepthStencil() && depthStencilAttachment == nullptr) ||
        (!vulkanRenderPass->HasDepthStencil() && depthStencilAttachment != nullptr)) {
        Shutdown();
        return false;
    }
    m_colorAttachmentViews.reserve(colorAttachmentCount);

    std::vector<VkImageView> attachments;
    attachments.reserve(colorAttachmentCount + (depthStencilAttachment != nullptr ? 1u : 0u));

    for (uint32_t i = 0; i < colorAttachmentCount; ++i) {
        auto* view = static_cast<VulkanResourceView*>(colorAttachments[i]);
        if (view == nullptr || view->GetImageView() == VK_NULL_HANDLE) {
            Shutdown();
            return false;
        }
        m_colorAttachmentViews.push_back(view);
        attachments.push_back(view->GetImageView());
    }

    m_depthStencilView = static_cast<VulkanResourceView*>(depthStencilAttachment);
    if (m_depthStencilView != nullptr) {
        if (m_depthStencilView->GetImageView() == VK_NULL_HANDLE) {
            Shutdown();
            return false;
        }
        attachments.push_back(m_depthStencilView->GetImageView());
    }

    VkFramebufferCreateInfo createInfo = {};
    createInfo.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
    createInfo.renderPass = vulkanRenderPass->GetRenderPass();
    createInfo.attachmentCount = static_cast<uint32_t>(attachments.size());
    createInfo.pAttachments = attachments.empty() ? nullptr : attachments.data();
    createInfo.width = width;
    createInfo.height = height;
    createInfo.layers = 1;

    if (vkCreateFramebuffer(device, &createInfo, nullptr, &m_framebuffer) != VK_SUCCESS) {
        Shutdown();
        return false;
    }

    return true;
}

bool VulkanFramebuffer::Initialize(RHIRenderPass* renderPass, uint32_t width, uint32_t height,
                                   uint32_t colorAttachmentCount, RHIResourceView** colorAttachments,
                                   RHIResourceView* depthStencilAttachment) {
    return Initialize(m_device, renderPass, width, height, colorAttachmentCount, colorAttachments, depthStencilAttachment);
}

void VulkanFramebuffer::Shutdown() {
    if (m_framebuffer != VK_NULL_HANDLE && m_device != VK_NULL_HANDLE) {
        vkDestroyFramebuffer(m_device, m_framebuffer, nullptr);
    }

    m_framebuffer = VK_NULL_HANDLE;
    m_device = VK_NULL_HANDLE;
    m_renderPass = nullptr;
    m_width = 0;
    m_height = 0;
    m_colorAttachmentViews.clear();
    m_depthStencilView = nullptr;
}

VulkanResourceView* VulkanFramebuffer::GetColorAttachmentView(uint32_t index) const {
    if (index >= m_colorAttachmentViews.size()) {
        return nullptr;
    }
    return m_colorAttachmentViews[index];
}

VulkanTexture* VulkanFramebuffer::GetColorAttachmentTexture(uint32_t index) const {
    VulkanResourceView* view = GetColorAttachmentView(index);
    return view != nullptr ? view->GetTexture() : nullptr;
}

VulkanTexture* VulkanFramebuffer::GetDepthStencilTexture() const {
    return m_depthStencilView != nullptr ? m_depthStencilView->GetTexture() : nullptr;
}

} // namespace RHI
} // namespace XCEngine