//
// Created by atdunbg on 2026/9/20.
//

#include <Engine/Vulkan/Swapchain.h>
#include <Engine/Vulkan/Device.h>
#include <Engine/Core/Window.h>
#include <Engine/Core/Log.h>

#include <algorithm>
#include <limits>

#include "Engine/Vulkan/MemoryAllocator.h"

namespace Engine
{
    Swapchain::Swapchain(Device& device, MemoryAllocator& allocator, Window& window)
        : m_Device(device), m_Allocator(allocator)
    {
        Create(window);
    }

    Swapchain::~Swapchain()
    {
        Destroy();
    }

    void Swapchain::Create(Window& window)
    {
        const auto& vkDevice = m_Device.GetDevice();
        auto surface = m_Device.GetSurface();
        auto physicalDevice = m_Device.GetPhysicalDevice();

        // 查询支持
        auto caps = physicalDevice.getSurfaceCapabilitiesKHR(surface);
        auto formats = physicalDevice.getSurfaceFormatsKHR(surface);
        auto modes = physicalDevice.getSurfacePresentModesKHR(surface);

        auto surfaceFormat = ChooseSurfaceFormat(formats);
        auto presentMode = ChoosePresentMode(modes);
        auto extent = ChooseExtent(caps, window);

        m_Format = surfaceFormat.format;
        m_Extent = extent;

        // 交换链 image 数量
        uint32_t imageCount = caps.minImageCount + 1;
        if (caps.maxImageCount > 0 && imageCount > caps.maxImageCount)
        {
            imageCount = caps.maxImageCount;
        }

        vk::SwapchainCreateInfoKHR createInfo;
        createInfo.setSurface(surface)
                  .setMinImageCount(imageCount)
                  .setImageFormat(surfaceFormat.format)
                  .setImageColorSpace(surfaceFormat.colorSpace)
                  .setImageExtent(extent)
                  .setImageArrayLayers(1)
                  .setImageUsage(vk::ImageUsageFlagBits::eColorAttachment)
                  .setImageSharingMode(vk::SharingMode::eExclusive)
                  .setPreTransform(caps.currentTransform)
                  .setCompositeAlpha(vk::CompositeAlphaFlagBitsKHR::eOpaque)
                  .setPresentMode(presentMode)
                  .setClipped(VK_TRUE)
                  .setOldSwapchain(nullptr);

        // graphics 和 present 队列族不同时，用 concurrent
        uint32_t indices[] = {m_Device.GetGraphicsFamily(), m_Device.GetPresentFamily()};
        if (indices[0] != indices[1])
        {
            createInfo.setImageSharingMode(vk::SharingMode::eConcurrent)
                      .setQueueFamilyIndices(indices);
        }

        m_Swapchain = vkDevice.createSwapchainKHR(createInfo);

        // 拿 image
        m_Images = vkDevice.getSwapchainImagesKHR(m_Swapchain);

        // 创建 image view
        m_ImageViews.resize(m_Images.size());
        for (size_t i = 0; i < m_Images.size(); ++i)
        {
            vk::ImageViewCreateInfo viewInfo;
            viewInfo.setImage(m_Images[i])
                    .setViewType(vk::ImageViewType::e2D)
                    .setFormat(m_Format)
                    .setSubresourceRange({
                        vk::ImageAspectFlagBits::eColor,
                        0, 1,
                        0, 1
                    });

            m_ImageViews[i] = vkDevice.createImageView(viewInfo);
        }

        ENGINE_INFO("Swapchain created: {}x{}, {} images, format={}",
                    m_Extent.width, m_Extent.height,
                    m_Images.size(),
                    vk::to_string(m_Format));

        CreateDepthResources();
    }

    void Swapchain::Destroy()
    {
        DestroyDepthResources();

        const auto& vkDevice = m_Device.GetDevice();

        for (auto view : m_ImageViews)
        {
            vkDevice.destroyImageView(view);
        }
        m_ImageViews.clear();
        m_Images.clear();

        if (m_Swapchain)
        {
            vkDevice.destroySwapchainKHR(m_Swapchain);
            m_Swapchain = nullptr;
        }
    }

    void Swapchain::Recreate(Window& window)
    {
        const auto& vkDevice = m_Device.GetDevice();
        vkDevice.waitIdle();

        Destroy();
        Create(window);

        ENGINE_INFO("Swapchain recreated: {}x{}", m_Extent.width, m_Extent.height);
    }

    // ---------- 选择函数 ----------

    vk::SurfaceFormatKHR Swapchain::ChooseSurfaceFormat(const std::vector<vk::SurfaceFormatKHR>& formats)
    {
        for (const auto& f : formats)
        {
            if (f.format == vk::Format::eB8G8R8A8Srgb &&
                f.colorSpace == vk::ColorSpaceKHR::eSrgbNonlinear)
            {
                return f;
            }
        }
        return formats[0];
    }

    vk::PresentModeKHR Swapchain::ChoosePresentMode(const std::vector<vk::PresentModeKHR>& modes)
    {
        for (const auto& m : modes)
        {
            if (m == vk::PresentModeKHR::eMailbox)
            {
                return m;
            }
        }
        return vk::PresentModeKHR::eFifo;
    }

    vk::Extent2D Swapchain::ChooseExtent(const vk::SurfaceCapabilitiesKHR& caps, Window& window)
    {
        if (caps.currentExtent.width != std::numeric_limits<uint32_t>::max())
        {
            return caps.currentExtent;
        }

        vk::Extent2D extent = {
            static_cast<uint32_t>(window.GetWidth()),
            static_cast<uint32_t>(window.GetHeight())
        };

        extent.width = std::clamp(extent.width, caps.minImageExtent.width, caps.maxImageExtent.width);
        extent.height = std::clamp(extent.height, caps.minImageExtent.height, caps.maxImageExtent.height);

        return extent;
    }

    vk::Format Swapchain::FindDepthFormat()
    {
        // 优先 D32_SFLOAT，退回 D32_SFLOAT_S8_UINT，再退回 D24_UNORM_S8_UINT
        std::array<vk::Format, 3> candidates = {
            vk::Format::eD32Sfloat,
            vk::Format::eD32SfloatS8Uint,
            vk::Format::eD24UnormS8Uint
        };

        for (auto format : candidates)
        {
            auto props = m_Device.GetPhysicalDevice().getFormatProperties(format);
            if (props.optimalTilingFeatures & vk::FormatFeatureFlagBits::eDepthStencilAttachment)
            {
                return format;
            }
        }

        ENGINE_FATAL("No depth format supported");
        return vk::Format::eD32Sfloat;
    }

    void Swapchain::CreateDepthResources()
    {
        m_DepthFormat = FindDepthFormat();

        vk::ImageCreateInfo imageInfo;
        imageInfo.setImageType(vk::ImageType::e2D)
                 .setFormat(m_DepthFormat)
                 .setExtent({m_Extent.width, m_Extent.height, 1})
                 .setMipLevels(1)
                 .setArrayLayers(1)
                 .setSamples(vk::SampleCountFlagBits::e1)
                 .setTiling(vk::ImageTiling::eOptimal)
                 .setUsage(vk::ImageUsageFlagBits::eDepthStencilAttachment)
                 .setSharingMode(vk::SharingMode::eExclusive)
                 .setInitialLayout(vk::ImageLayout::eUndefined);

        m_DepthImage = m_Device.GetDevice().createImage(imageInfo);

        auto reqs = m_Device.GetDevice().getImageMemoryRequirements(m_DepthImage);
        auto alloc = m_Allocator.Allocate(reqs, vk::MemoryPropertyFlagBits::eDeviceLocal);

        m_DepthMemory = alloc.Memory;
        m_Device.GetDevice().bindImageMemory(m_DepthImage, m_DepthMemory, alloc.Offset);

        vk::ImageViewCreateInfo viewInfo;
        viewInfo.setImage(m_DepthImage)
                .setViewType(vk::ImageViewType::e2D)
                .setFormat(m_DepthFormat)
                .setSubresourceRange({
                    vk::ImageAspectFlagBits::eDepth,
                    0, 1, 0, 1
                });

        m_DepthImageView = m_Device.GetDevice().createImageView(viewInfo);
    }

    void Swapchain::DestroyDepthResources()
    {
        const auto& vkDevice = m_Device.GetDevice();
        if (m_DepthImageView)
        {
            vkDevice.destroyImageView(m_DepthImageView);
            m_DepthImageView = nullptr;
        }
        if (m_DepthImage)
        {
            vkDevice.destroyImage(m_DepthImage);
            m_DepthImage = nullptr;
        }
        // 内存由 MemoryAllocator 统一释放
    }
} // namespace Engine