// =============================================
//  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)
{
    WARN_IF(!IsPowerOfTwo(extent.width) || !IsPowerOfTwo(extent.width), "Image {2} is {0}x{1} (Non Power of Two)",
            extent.width, extent.height, name ? name : "<unnamed>");

    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 = ToExtent3D(extent, 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 = {},
        .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 = imageCreateInfo.extent;
    m_MipLevels_ = mipLevels | OWNED_BIT | VALID_BIT;

    device->SetName(m_Image, name);
}

/*
Cube map Faces info.

TODO: Correct this based on the actual layout for upside down viewport.

| Axis | Layer | Up |
|:----:|:-----:|:--:|
|  +x  |   0   | -y |
|  -x  |   1   | -y |
|  +y  |   2   | +z |
|  -y  |   3   | -z |
|  +z  |   4   | -y |
|  -z  |   5   | -y |

Remember, we use upside down viewport.

*/

void
TextureCube::Init(const Device *device, u32 cubeSide, vk::Format imageFormat, bool isMipMapped, cstr name)
{
    WARN_IF(!IsPowerOfTwo(cubeSide), "Image {1} is {0}x{0} (Non Power of Two)", cubeSide, name ? name : "<unnamed>");

    const u32 mipLevels = isMipMapped ? 1 + Cast<u32>(floor(log2(cubeSide))) : 1;
    assert(mipLevels <= MIP_MASK);

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

    vk::Extent3D extent = {.width = cubeSide, .height = cubeSide, .depth = 1};

    vk::ImageCreateInfo imageCreateInfo = {
        .flags = vk::ImageCreateFlagBits::eCubeCompatible,
        .imageType = vk::ImageType::e2D,
        .format = imageFormat,
        .extent = extent,
        .mipLevels = mipLevels,
        .arrayLayers = 6,
        .samples = vk::SampleCountFlagBits::e1,
        .tiling = vk::ImageTiling::eOptimal,
        .usage = usage,
        .sharingMode = vk::SharingMode::eExclusive,
        .initialLayout = vk::ImageLayout::eUndefined,
    };
    constexpr VmaAllocationCreateInfo allocationCreateInfo = {
        .flags = {},
        .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::eCube,
        .format = imageFormat,
        .components = {},
        .subresourceRange =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseMipLevel = 0,
                .levelCount = mipLevels,
                .baseArrayLayer = 0,
                .layerCount = 6,
            },
    };
    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;
    m_MipLevels_ = mipLevels | OWNED_BIT | VALID_BIT;

    device->SetName(m_Image, name);
}

void
AttachmentImage::Init(const Device *device, vk::Extent2D extent, vk::Format imageFormat, cstr name)
{
    vk::ImageCreateInfo imageCreateInfo = {
        .imageType = vk::ImageType::e2D,
        .format = imageFormat,
        .extent = ToExtent3D(extent, 1),
        .mipLevels = 1,
        .arrayLayers = 1,
        .samples = vk::SampleCountFlagBits::e1,
        .tiling = vk::ImageTiling::eOptimal,
        .usage = vk::ImageUsageFlagBits::eColorAttachment | vk::ImageUsageFlagBits::eTransferSrc,
        .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::eColor,
                .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 = imageCreateInfo.extent;
    m_MipLevels_ = 1 | 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 = ToExtent3D(extent, 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 = imageCreateInfo.extent;
    m_MipLevels_ = 1 | OWNED_BIT | VALID_BIT;

    device->SetName(m_Image, name);
}

void
StorageTexture::Init(const Device *device, vk::Extent2D extent, const vk::Format imageFormat, const bool isSampled,
                     cstr name)
{

    // Reasoning:
    // Transfer Src and Dst to copy to and from the buffer since Storage will often be loaded with info, and read for
    // results.
    auto usage =
        vk::ImageUsageFlagBits::eStorage | vk::ImageUsageFlagBits::eTransferSrc | vk::ImageUsageFlagBits::eTransferDst;
    if (isSampled)
    {
        WARN_IF(!IsPowerOfTwo(extent.width) || !IsPowerOfTwo(extent.width), "Image {2} is {0}x{1} (Non Power of Two)",
                extent.width, extent.height, name ? name : "<unnamed>");
        usage |= vk::ImageUsageFlagBits::eSampled;
    }

    vk::ImageCreateInfo imageCreateInfo = {
        .imageType = vk::ImageType::e2D,
        .format = imageFormat,
        .extent = ToExtent3D(extent, 1),
        .mipLevels = 1,
        .arrayLayers = 1,
        .samples = vk::SampleCountFlagBits::e1,
        .tiling = vk::ImageTiling::eOptimal,
        .usage = usage,
        .sharingMode = vk::SharingMode::eExclusive,
        .initialLayout = vk::ImageLayout::eUndefined,
    };
    constexpr VmaAllocationCreateInfo allocationCreateInfo = {
        .flags = {},
        .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;
    const vk::ImageViewCreateInfo imageViewCreateInfo = {
        .image = image,
        .viewType = vk::ImageViewType::e2D,
        .format = imageFormat,
        .components = {},
        .subresourceRange =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseMipLevel = 0,
                .levelCount = 1,
                .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 = imageCreateInfo.extent;
    m_MipLevels_ = 1 | OWNED_BIT | VALID_BIT;

    device->SetName(m_Image, name);
}