#include "vulkan_image.h"

#include <SDL3/SDL_log.h>
#include <stdint.h>

#include "src/graphics_context.h"
#include "utilities/commons.h"
#include "utilities/errors_common.h"
#include "utilities/file_utils.h"
#include "vulkan/vulkan_core.h"
#include "vulkan_buffers.h"
#include "vulkan_commands.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define IMAGE_TAKEN 1U
#define IMAGE_FREE 0U

rse_err_t sampler_create(struct graphics_context_t* context, VkSampler* sampler)
{
    VkSamplerCreateInfo create_info = {0};

    create_info.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;
    create_info.pNext = NULL;
    create_info.flags = 0U;
    create_info.magFilter = VK_FILTER_LINEAR;
    create_info.minFilter = VK_FILTER_LINEAR;
    create_info.mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR;
    create_info.addressModeU = VK_SAMPLER_ADDRESS_MODE_REPEAT;
    create_info.addressModeV = VK_SAMPLER_ADDRESS_MODE_REPEAT;
    create_info.addressModeW = VK_SAMPLER_ADDRESS_MODE_REPEAT;
    create_info.mipLodBias = 0.0f;
    create_info.anisotropyEnable = VK_FALSE;  // TODO: Enable this
    create_info.maxAnisotropy = 1.0f;
    create_info.compareEnable = VK_FALSE;
    create_info.compareOp = VK_COMPARE_OP_ALWAYS;
    create_info.minLod = 0.0f;
    create_info.maxLod = 0.0f;
    create_info.borderColor = VK_BORDER_COLOR_INT_OPAQUE_BLACK;
    create_info.unnormalizedCoordinates = VK_FALSE;

    if (VK_SUCCESS != vkCreateSampler(context->vulkan_handles.device, &create_info, NULL, sampler)) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed to create image view");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    return RSE_ERROR_NO_ERROR;
}

static rse_err_t find_supported_format(struct graphics_context_t* context,
                                       const VkFormat* candidates,
                                       size_t candidates_count,
                                       VkImageTiling tiling,
                                       VkFormatFeatureFlags features,
                                       VkFormat* found_format)
{
    size_t i = 0U;

    for (i = 0U; i < candidates_count; ++i) {
        VkFormatProperties props;
        vkGetPhysicalDeviceFormatProperties(context->vulkan_handles.physical_device, candidates[i], &props);

        if (tiling == VK_IMAGE_TILING_LINEAR && (props.linearTilingFeatures & features) == features) {
            *found_format = candidates[i];
            return RSE_ERROR_NO_ERROR;
        } else if (tiling == VK_IMAGE_TILING_OPTIMAL && (props.optimalTilingFeatures & features) == features) {
            *found_format = candidates[i];
            return RSE_ERROR_NO_ERROR;
        }
    }

    SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed to find correct format");
    return RSE_ERROR_INTERNAL_ERROR;
}

static rse_err_t create_image(struct graphics_context_t* context,
                              struct vulkan_image_t* image,
                              uint32_t width,
                              uint32_t height,
                              VkFormat format,
                              VkSampleCountFlagBits maa_samples,
                              VkImageUsageFlags usage,
                              VmaAllocationCreateFlags allocationFlags)
{
    VkResult result;
    VkExtent3D image_extent = {0};
    VkImageCreateInfo image_info = {0};
    VmaAllocationCreateInfo alloc_create_info = {0};

    image_extent.width = width;
    image_extent.height = height;
    image_extent.depth = 1U;

    image_info.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
    image_info.pNext = NULL;
    image_info.flags = 0U;
    image_info.imageType = VK_IMAGE_TYPE_2D;
    image_info.format = format;
    image_info.extent = image_extent;
    image_info.mipLevels = 1U;
    image_info.arrayLayers = 1U;
    image_info.samples = maa_samples;
    image_info.tiling = VK_IMAGE_TILING_OPTIMAL;
    image_info.usage = usage;
    image_info.sharingMode = VK_SHARING_MODE_EXCLUSIVE;  // TODO: Change if used by more than one
                                                         // queue family
    image_info.queueFamilyIndexCount = 0U;
    image_info.pQueueFamilyIndices = NULL;
    image_info.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;

    alloc_create_info.flags = allocationFlags;
    alloc_create_info.usage = VMA_MEMORY_USAGE_AUTO;
    alloc_create_info.requiredFlags = 0U;
    alloc_create_info.preferredFlags = 0U;
    alloc_create_info.memoryTypeBits = 0U;
    alloc_create_info.pool = NULL;
    alloc_create_info.pUserData = NULL;
    alloc_create_info.priority = 1.0f;

    result = vmaCreateImage(context->vulkan_handles.allocator,
                            &image_info,
                            &alloc_create_info,
                            &image->image,
                            &image->allocation,
                            &image->allocation_info);
    if (VK_SUCCESS != result) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed to create image");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    return RSE_ERROR_NO_ERROR;
}

static rse_err_t create_image_view(struct graphics_context_t* context,
                                   struct vulkan_image_t* image,
                                   VkFormat format,
                                   VkImageAspectFlags aspectFlags)
{
    VkImageViewCreateInfo image_view_create_info = {0};

    image_view_create_info.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
    image_view_create_info.pNext = NULL;
    image_view_create_info.flags = 0U;
    image_view_create_info.image = image->image;
    image_view_create_info.viewType = VK_IMAGE_VIEW_TYPE_2D; /* 2D Texture */
    image_view_create_info.format = format;
    image_view_create_info.components.a = VK_COMPONENT_SWIZZLE_IDENTITY;
    image_view_create_info.components.r = VK_COMPONENT_SWIZZLE_IDENTITY;
    image_view_create_info.components.g = VK_COMPONENT_SWIZZLE_IDENTITY;
    image_view_create_info.components.b = VK_COMPONENT_SWIZZLE_IDENTITY;
    image_view_create_info.subresourceRange.aspectMask = aspectFlags;
    image_view_create_info.subresourceRange.baseMipLevel = 0U;
    image_view_create_info.subresourceRange.levelCount = 1U;
    image_view_create_info.subresourceRange.baseArrayLayer = 0U;
    image_view_create_info.subresourceRange.layerCount = 1U;

    if (VK_SUCCESS !=
        vkCreateImageView(context->vulkan_handles.device, &image_view_create_info, NULL, &image->image_view)) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed to create swapchain image view");
        return RSE_ERROR_INTERNAL_ERROR;
    }
    return RSE_ERROR_NO_ERROR;
}

rse_err_t create_color_resource(struct graphics_context_t* context)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    VkFormat color_format = IMAGE_FORMAT;
    VkSurfaceCapabilitiesKHR physical_device_surface_capabilities;

    vkGetPhysicalDeviceSurfaceCapabilitiesKHR(context->vulkan_handles.physical_device,
                                              context->vulkan_handles.surface,
                                              &physical_device_surface_capabilities);

    STATUS_CHECK(create_image(context,
                              &context->color_image,
                              physical_device_surface_capabilities.currentExtent.width,
                              physical_device_surface_capabilities.currentExtent.height,
                              color_format,
                              VK_SAMPLE_COUNT_8_BIT,
                              VK_IMAGE_USAGE_TRANSIENT_ATTACHMENT_BIT | VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
                              0U));
    STATUS_CHECK(create_image_view(context, &context->color_image, color_format, VK_IMAGE_ASPECT_COLOR_BIT));

    return status;
}

rse_err_t create_depth_resources(struct graphics_context_t* context)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    VkFormat depth_format = {0};
    VkSurfaceCapabilitiesKHR physical_device_surface_capabilities;

    STATUS_CHECK(find_depth_format(context, &depth_format));
    vkGetPhysicalDeviceSurfaceCapabilitiesKHR(context->vulkan_handles.physical_device,
                                              context->vulkan_handles.surface,
                                              &physical_device_surface_capabilities);

    STATUS_CHECK(create_image(context,
                              &context->depth_image,
                              physical_device_surface_capabilities.currentExtent.width,
                              physical_device_surface_capabilities.currentExtent.height,
                              depth_format,
                              VK_SAMPLE_COUNT_8_BIT,
                              VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT,
                              0U));
    STATUS_CHECK(create_image_view(context, &context->depth_image, depth_format, VK_IMAGE_ASPECT_DEPTH_BIT));

    return status;
}

/**
 * @brief Translate vulkan format to pixel width
 *
 * @param[int] format Format to translate
 * @param[out] format_size Format size in bytes to return
 * @return RSE_ERROR_NO_ERROR on success, error code otherwise
 */
static rse_err_t format_to_pixel_size(VkFormat format, VkDeviceSize* format_size)
{
    switch (format) {
        case VK_FORMAT_R8_SRGB:
        case VK_FORMAT_R8_SNORM:
        case VK_FORMAT_R8_UNORM:
        case VK_FORMAT_R8_USCALED:
        case VK_FORMAT_R8_SSCALED:
        case VK_FORMAT_R8_UINT:
            *format_size = 1;
            break;
        case VK_FORMAT_R8G8B8_UINT:
        case VK_FORMAT_R8G8B8_UNORM:
            *format_size = 3;
            break;
        case VK_FORMAT_R8G8B8A8_SRGB:
        case VK_FORMAT_R8G8B8A8_UNORM:
            *format_size = 4;
            break;

        default:
            SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed translate vulkan format to pixel width");
            break;
    }

    return RSE_ERROR_NO_ERROR;
}

static rse_err_t transition_image_layout(struct graphics_context_t* context,
                                         struct vulkan_image_t* image,
                                         VkImageLayout old_layout,
                                         VkImageLayout new_layout)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    VkImageSubresourceRange range = {0};
    VkImageMemoryBarrier image_barrier = {0};
    VkCommandBuffer command_buffer = {0};
    VkPipelineStageFlags source_stage = {0};
    VkPipelineStageFlags destination_stage = {0};

    /* Put image into correct layout to copy pixels from buffer to image memory
     */
    range.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    range.baseMipLevel = 0;
    range.levelCount = 1;
    range.baseArrayLayer = 0;
    range.layerCount = 1;

    image_barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
    image_barrier.pNext = NULL;
    image_barrier.oldLayout = old_layout;
    image_barrier.newLayout = new_layout;
    image_barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    image_barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    image_barrier.image = image->image;
    image_barrier.subresourceRange = range;

    if (old_layout == VK_IMAGE_LAYOUT_UNDEFINED && new_layout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL) {
        image_barrier.srcAccessMask = 0;
        image_barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;

        source_stage = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
        destination_stage = VK_PIPELINE_STAGE_TRANSFER_BIT;
    } else if (old_layout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL &&
               new_layout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL) {
        image_barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
        image_barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;

        source_stage = VK_PIPELINE_STAGE_TRANSFER_BIT;
        destination_stage = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
    } else {
        return RSE_ERROR_INTERNAL_ERROR;
    }

    STATUS_CHECK(begin_single_time_command(context, &command_buffer));
    vkCmdPipelineBarrier(command_buffer, source_stage, destination_stage, 0, 0, NULL, 0, NULL, 1, &image_barrier);
    STATUS_CHECK(end_single_time_comands(context, command_buffer));

    return RSE_ERROR_NO_ERROR;
}

static rse_err_t create_textured_image(struct graphics_context_t* context,
                                       uint32_t width,
                                       uint32_t height,
                                       VkFormat format,
                                       const unsigned char* pixels,
                                       uint16_t* texture_id)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    struct vulkan_buffer_t staging_buffer = {0};
    VkDeviceSize staging_buffer_size = 0U;
    VkCommandBuffer command_buffer = {0};
    VkBufferImageCopy copy_region = {0};
    VkExtent3D image_extent = {0};
    struct vulkan_image_t* free_texture_image = NULL;

    image_extent.width = width;
    image_extent.height = height;
    image_extent.depth = 1U;

    *texture_id = 0U;

    while (free_texture_image == NULL && *texture_id < RSE_MAX_IMAGE_COUNT) {
        if (context->texture_images[*texture_id].id_taken == IMAGE_FREE) {
            free_texture_image = &context->texture_images[*texture_id];
        } else {
            *texture_id += 1;
        }
    }

    /* Check if we found free image handle */
    if (free_texture_image == NULL) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed to find free texture image handle");
        *texture_id = -1;
        return RSE_ERROR_INTERNAL_ERROR;
    }

    STATUS_CHECK(format_to_pixel_size(format, &staging_buffer_size));

    staging_buffer_size = staging_buffer_size * width * height;

    STATUS_CHECK(
        create_image(context,
                     free_texture_image,
                     width,
                     height,
                     format,
                     VK_SAMPLE_COUNT_1_BIT,
                     VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
                     VMA_ALLOCATION_CREATE_DEDICATED_MEMORY_BIT));

    transition_image_layout(context,
                            free_texture_image,
                            VK_IMAGE_LAYOUT_UNDEFINED,
                            VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);

    /* Copy pixel data to GPU memory */
    STATUS_CHECK(
        create_buffer(context,
                      staging_buffer_size,
                      VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
                      VMA_MEMORY_USAGE_AUTO_PREFER_HOST,
                      VMA_ALLOCATION_CREATE_MAPPED_BIT | VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT,
                      &staging_buffer));

    memcpy(staging_buffer.allocation_info.pMappedData, pixels, staging_buffer_size);

    copy_region.bufferOffset = 0U;
    copy_region.bufferRowLength = 0U;
    copy_region.bufferImageHeight = 0U;
    copy_region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    copy_region.imageSubresource.mipLevel = 0U;
    copy_region.imageSubresource.baseArrayLayer = 0U;
    copy_region.imageSubresource.layerCount = 1U;
    copy_region.imageOffset.x = 0;
    copy_region.imageOffset.y = 0;
    copy_region.imageOffset.z = 0;
    copy_region.imageExtent = image_extent;

    STATUS_CHECK(begin_single_time_command(context, &command_buffer));
    vkCmdCopyBufferToImage(command_buffer,
                           staging_buffer.buffer,
                           free_texture_image->image,
                           VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                           1,
                           &copy_region);
    end_single_time_comands(context, command_buffer);
    destroy_buffer(context, &staging_buffer);

    transition_image_layout(context,
                            free_texture_image,
                            VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                            VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
    STATUS_CHECK(create_image_view(context, free_texture_image, format, VK_IMAGE_ASPECT_COLOR_BIT));

    free_texture_image->id_taken = IMAGE_TAKEN;
    return RSE_ERROR_NO_ERROR;
}

rse_err_t init_vulkan_images(struct graphics_context_t* context)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;

    STATUS_CHECK(create_depth_resources(context));
    STATUS_CHECK(create_color_resource(context));

    return status;
}

rse_err_t texture_update_image(struct graphics_context_t* context,
                               const unsigned char* pixels,
                               int width,
                               int height,
                               uint16_t texture_id,
                               VkFormat format)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    VkDeviceSize staging_buffer_size = 0U;
    VkBufferImageCopy copy_region = {0};
    struct vulkan_buffer_t staging_buffer = {0};
    VkExtent3D image_extent = {0};
    VkCommandBuffer command_buffer = {0};

    image_extent.width = width;
    image_extent.height = height;
    image_extent.depth = 1U;

    STATUS_CHECK(format_to_pixel_size(format, &staging_buffer_size));

    staging_buffer_size = staging_buffer_size * width * height;

    transition_image_layout(context,
                            &context->texture_images[texture_id],
                            VK_IMAGE_LAYOUT_UNDEFINED,
                            VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);

    /* Copy pixel data to GPU memory */
    STATUS_CHECK(
        create_buffer(context,
                      staging_buffer_size,
                      VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
                      VMA_MEMORY_USAGE_AUTO_PREFER_HOST,
                      VMA_ALLOCATION_CREATE_MAPPED_BIT | VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT,
                      &staging_buffer));

    memcpy(staging_buffer.allocation_info.pMappedData, pixels, staging_buffer_size);

    copy_region.bufferOffset = 0U;
    copy_region.bufferRowLength = 0U;
    copy_region.bufferImageHeight = 0U;
    copy_region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    copy_region.imageSubresource.mipLevel = 0U;
    copy_region.imageSubresource.baseArrayLayer = 0U;
    copy_region.imageSubresource.layerCount = 1U;
    copy_region.imageOffset.x = 0;
    copy_region.imageOffset.y = 0;
    copy_region.imageOffset.z = 0;
    copy_region.imageExtent = image_extent;

    STATUS_CHECK(begin_single_time_command(context, &command_buffer));
    vkCmdCopyBufferToImage(command_buffer,
                           staging_buffer.buffer,
                           context->texture_images[texture_id].image,
                           VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                           1,
                           &copy_region);
    end_single_time_comands(context, command_buffer);
    destroy_buffer(context, &staging_buffer);

    transition_image_layout(context,
                            &context->texture_images[texture_id],
                            VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                            VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

    return RSE_ERROR_NO_ERROR;
}

rse_err_t load_texture_from_bitmat(struct graphics_context_t* context,
                                   const unsigned char* pixel_buffer,
                                   int width,
                                   int height,
                                   uint16_t* texture_id,
                                   VkFormat color_format)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;

    /* Create actual Vulkan Image */
    STATUS_CHECK(create_textured_image(context, width, height, color_format, pixel_buffer, texture_id));

    return status;
}

rse_err_t load_texture_from_file(struct graphics_context_t* context, const char* file_path, uint16_t* texture_id)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    int width = 0;
    int height = 0;
    size_t buffer_size = 0U;
    unsigned char* pixel_buffer = NULL;

    STATUS_CHECK(file_load_pixels(file_path, NULL, &buffer_size, &width, &height));
    rse_malloc(pixel_buffer, buffer_size);
    if (file_load_pixels(file_path, pixel_buffer, &buffer_size, &width, &height) != RSE_ERROR_NO_ERROR) {
        status = RSE_ERROR_INTERNAL_ERROR;

        goto mem_free;
    }

    if (load_texture_from_bitmat(context, pixel_buffer, width, height, texture_id, VK_FORMAT_R8G8B8A8_SRGB) != RSE_ERROR_NO_ERROR) {
        status = RSE_ERROR_INTERNAL_ERROR;

        goto mem_free;
    }

mem_free:

    rse_free(pixel_buffer);

    return status;
}

void destroy_textures(struct graphics_context_t* context)
{
    size_t i = 0U;

    for (i = 0U; i < RSE_MAX_IMAGE_COUNT; ++i) {
        if (context->texture_images[i].id_taken == IMAGE_TAKEN) {
            vkDestroyImageView(context->vulkan_handles.device, context->texture_images[i].image_view, NULL);
            vmaDestroyImage(context->vulkan_handles.allocator,
                            context->texture_images[i].image,
                            context->texture_images[i].allocation);
        }
    }

    vkDestroySampler(context->vulkan_handles.device, context->vulkan_handles.sampler, NULL);
    vkDestroySampler(context->vulkan_handles.device, context->debug_overlay.sampler, NULL);
}

void destroy_depth_resource(struct graphics_context_t* context)
{
    vkDestroyImageView(context->vulkan_handles.device, context->depth_image.image_view, NULL);
    vmaDestroyImage(context->vulkan_handles.allocator, context->depth_image.image, context->depth_image.allocation);
}

void destroy_color_resource(struct graphics_context_t* context)
{
    vkDestroyImageView(context->vulkan_handles.device, context->color_image.image_view, NULL);
    vmaDestroyImage(context->vulkan_handles.allocator, context->color_image.image, context->color_image.allocation);
}

uint8_t image_exists(struct graphics_context_t* context, uint8_t image_id)
{
    return context->texture_images[image_id].id_taken == IMAGE_TAKEN;
}

size_t get_textures_count(struct graphics_context_t* context)
{
    size_t count = 0U;
    size_t i = 0U;

    for (i = 0U; i < RSE_MAX_IMAGE_COUNT; ++i) {
        if (context->texture_images[i].id_taken == IMAGE_TAKEN) {
            count++;
        }
    }

    return count;
}

rse_err_t find_depth_format(struct graphics_context_t* context, VkFormat* found_format)
{
#define formats_count 3
    rse_err_t status;
    VkFormat formats[formats_count] = {VK_FORMAT_D32_SFLOAT};

    STATUS_CHECK(find_supported_format(context,
                                       formats,
                                       formats_count,
                                       VK_IMAGE_TILING_OPTIMAL,
                                       VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT,
                                       found_format));

    return status;
#undef formats_count
}