// =============================================
//  Aster: image.cpp
//  Copyright (c) 2020-2024 Anish Bhobe
// =============================================

#include "image.h"

#include "device.h"

void
Image::Destroy(const Device *device)
{
    if (!IsValid() || !IsOwned())
        return;

    device->m_Device.destroy(m_View, nullptr);
    vmaDestroyImage(device->m_Allocator, m_Image, m_Allocation);
    m_MipLevels_ = 0;
}

void
Texture::Init(const Device *device, const vk::Extent2D extent, vk::Format imageFormat, const bool isMipmapped,
              const cstr name)
{
    const u32 mipLevels = isMipmapped ? 1 + Cast<u32>(floor(log2(eastl::max(extent.width, extent.height)))) : 1;
    assert(mipLevels <= MIP_MASK);

    auto usage = vk::ImageUsageFlagBits::eSampled | vk::ImageUsageFlagBits::eTransferDst;
    if (isMipmapped)
    {
        usage |= vk::ImageUsageFlagBits::eTransferSrc;
    }

    vk::ImageCreateInfo imageCreateInfo = {
        .imageType = vk::ImageType::e2D,
        .format = imageFormat,
        .extent = {.width = extent.width, .height = extent.height, .depth = 1},
        .mipLevels = mipLevels,
        .arrayLayers = 1,
        .samples = vk::SampleCountFlagBits::e1,
        .tiling = vk::ImageTiling::eOptimal,
        .usage = usage,
        .sharingMode = vk::SharingMode::eExclusive,
        .initialLayout = vk::ImageLayout::eUndefined,
    };
    constexpr VmaAllocationCreateInfo allocationCreateInfo = {
        .flags = VMA_ALLOCATION_CREATE_DEDICATED_MEMORY_BIT,
        .usage = VMA_MEMORY_USAGE_AUTO,
    };

    VkImage image;
    VmaAllocation allocation;
    auto result = Cast<vk::Result>(vmaCreateImage(device->m_Allocator, Recast<VkImageCreateInfo *>(&imageCreateInfo),
                                                  &allocationCreateInfo, &image, &allocation, nullptr));
    ERROR_IF(Failed(result), "Could not allocate image {}. Cause: {}", name, result) THEN_ABORT(result);

    vk::ImageView view;
    vk::ImageViewCreateInfo imageViewCreateInfo = {
        .image = image,
        .viewType = vk::ImageViewType::e2D,
        .format = imageFormat,
        .components = {},
        .subresourceRange =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseMipLevel = 0,
                .levelCount = mipLevels,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
    };
    result = device->m_Device.createImageView(&imageViewCreateInfo, nullptr, &view);
    ERROR_IF(Failed(result), "Could not create image view {}. Cause: {}", name, result) THEN_ABORT(result);

    m_Image = image;
    m_View = view;
    m_Allocation = allocation;
    m_Extent = {extent.width, extent.height, 1};
    m_MipLevels_ = mipLevels | OWNED_BIT | VALID_BIT;

    device->SetName(m_Image, name);
}

void
DepthImage::Init(const Device *device, vk::Extent2D extent, cstr name)
{
    constexpr vk::Format imageFormat = vk::Format::eD24UnormS8Uint;
    vk::ImageCreateInfo imageCreateInfo = {
        .imageType = vk::ImageType::e2D,
        .format = imageFormat,
        .extent = {.width = extent.width, .height = extent.height, .depth = 1},
        .mipLevels = 1,
        .arrayLayers = 1,
        .samples = vk::SampleCountFlagBits::e1,
        .tiling = vk::ImageTiling::eOptimal,
        .usage = vk::ImageUsageFlagBits::eDepthStencilAttachment,
        .sharingMode = vk::SharingMode::eExclusive,
        .initialLayout = vk::ImageLayout::eUndefined,
    };
    constexpr VmaAllocationCreateInfo allocationCreateInfo = {
        .flags = VMA_ALLOCATION_CREATE_DEDICATED_MEMORY_BIT,
        .usage = VMA_MEMORY_USAGE_AUTO,
    };

    VkImage image;
    VmaAllocation allocation;
    auto result = Cast<vk::Result>(vmaCreateImage(device->m_Allocator, Recast<VkImageCreateInfo *>(&imageCreateInfo),
                                                  &allocationCreateInfo, &image, &allocation, nullptr));
    ERROR_IF(Failed(result), "Could not allocate depth buffer. Cause: {}", result) THEN_ABORT(result);

    vk::ImageView view;
    vk::ImageViewCreateInfo imageViewCreateInfo = {
        .image = image,
        .viewType = vk::ImageViewType::e2D,
        .format = imageFormat,
        .components = {},
        .subresourceRange =
            {
                .aspectMask = vk::ImageAspectFlagBits::eDepth,
                .baseMipLevel = 0,
                .levelCount = 1,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
    };
    result = device->m_Device.createImageView(&imageViewCreateInfo, nullptr, &view);
    ERROR_IF(Failed(result), "Could not create depth image view {}. Cause: {}", name, result) THEN_ABORT(result);

    m_Image = image;
    m_View = view;
    m_Allocation = allocation;
    m_Extent = {extent.width, extent.height, 1};
    m_MipLevels_ = 1 | OWNED_BIT | VALID_BIT;

    device->SetName(m_Image, name);
}