export module renderer.backend.vk.swapchain;
import renderer.backend.vk.library_wrapper;
import renderer.backend.vk.surface;
import renderer.backend.vk.device;
import renderer.backend.vk.instance;
import renderer.backend.vk.gpu;
import renderer.frontend;
import math.vec2;
import memory.null_on_move;
import logger;
import std;

namespace lt::renderer::vkb {

export class Swapchain: public ISwapchain
{
public:
	Swapchain(ISurface *surface, IGpu *gpu, IDevice *device);

	[[nodiscard]] auto vk() -> vk::Swapchain &
	{
		return m_swapchain;
	}

	[[nodiscard]] auto get_resolution() const -> math::uvec2
	{
		return m_resolution;
	}

	[[nodiscard]] auto get_format() const -> vk::Format
	{
		return m_format;
	}

	[[nodiscard]] auto get_image_count() const -> std::size_t
	{
		return m_images.size();
	}

	[[nodiscard]] auto get_image_view(std::uint32_t idx) -> vk::ImageView &
	{
		return m_image_views[idx];
	}

	[[nodiscard]] auto get_image(std::uint32_t idx) -> vk::Image &
	{
		return m_images[idx];
	}

private:
	[[nodiscard]] auto get_optimal_image_count(
	    vk::Surface::Capabilities capabilities,
	    std::uint32_t desired_image_count
	) const -> std::uint32_t;

	Gpu *m_gpu;

	Surface *m_surface {};

	Device *m_device;

	vk::Swapchain m_swapchain;

	std::vector<vk::Image> m_images;

	std::vector<vk::ImageView> m_image_views;

	math::uvec2 m_resolution {};

	vk::Format m_format {};
};

} // namespace lt::renderer::vkb


module :private;
using namespace lt::renderer;
using namespace lt::renderer::vkb;

Swapchain::Swapchain(ISurface *surface, IGpu *gpu, IDevice *device)
    : m_surface(static_cast<Surface *>(surface))
    , m_gpu(static_cast<Gpu *>(gpu))
    , m_device(static_cast<Device *>(device))
{
	static auto idx = 0u;

	const auto capabilities = m_gpu->vk().get_surface_capabilities(m_surface->vk());
	const auto formats = m_gpu->vk().get_surface_formats(m_surface->vk());

	// TODO(Light): parameterize
	constexpr auto desired_image_count = std::uint32_t { 3u };
	const auto surface_format = formats.front();
	m_format = surface_format.format;

	m_swapchain = vk::Swapchain(
	    m_device->vk(),
	    m_surface->vk(),
	    vk::Swapchain::CreateInfo {

	        .format = surface_format.format,
	        .color_space = surface_format.color_space,
	        .extent = capabilities.current_extent,
	        .min_image_count = get_optimal_image_count(capabilities, desired_image_count),
	        .queue_family_indices = m_device->get_family_indices(),
	        .present_mode = vk::Swapchain::PresentMode::immediate,
	        .pre_transform = capabilities.current_transform,
	    }
	);
	m_resolution = capabilities.current_extent;
	m_device->vk().name(m_swapchain, "swapchain {}", idx++);
	m_device->vk().wait_idle();

	m_images = m_swapchain.get_images();
	for (auto idx = 0u; auto &image : m_images)
	{
		m_image_views.emplace_back(
		    vk::ImageView {
		        m_device->vk(),
		        image,
		        vk::ImageView::CreateInfo {
                .type = vk::ImageView::Type::_2d,
                .format = surface_format.format,
                .components = {
                    vk::ImageView::Swizzle::identity,
                    vk::ImageView::Swizzle::identity,
                    vk::ImageView::Swizzle::identity,
                    vk::ImageView::Swizzle::identity,
                },
                .range = {
		            .aspect_flags = vk::Image::AspectFlags::color_bit,
		            .base_mip_level = 0u,
		            .level_count = 1u,
		            .base_array_layer = 0u,
		            .layer_count = 1u,
                },

                .debug_name = std::format("swapchain image {}", idx++),
		        },
		    }
		);
	}
	m_device->vk().wait_idle();
}

[[nodiscard]] auto Swapchain::get_optimal_image_count(
    vk::Surface::Capabilities capabilities,
    std::uint32_t desired_image_count
) const -> std::uint32_t
{
	const auto min_image_count = capabilities.min_image_count;
	const auto max_image_count = capabilities.max_image_count;

	const auto has_max_limit = max_image_count != 0;

	// Desired image count is in range
	if ((!has_max_limit || max_image_count >= desired_image_count)
	    && min_image_count <= desired_image_count)
	{
		return desired_image_count;
	}

	// Fall-back to 2 if in range
	if (min_image_count <= 2 && max_image_count >= 2)
	{
		return 2;
	}

	// Fall-back to min_image_count
	return min_image_count;
}