#include "rse_vulkan_image.h"
#include <stdint.h>

#include "src/vulkan_pipeline.h"
#include "utilities/file_utils.h"
#include "utilities/logger.h"
#include "utilities/rse_commons.h"
#include "utilities/rse_errors_common.h"
#include "locale_vulkan.h"
#include "vk_mem_alloc.h"
#include "vulkan/vulkan_core.h"
#include "vulkan_errors.h"
#include "rse_vulkan_commands.h"
#include "vulkan_buffers.h"

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#include "vma/vk_mem_alloc.h"

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

#define IMAGE_TAKEN 1U
#define IMAGE_FREE 0U

/**
 * @brief Wrapper for Vulkan buffer. Holds allocation data;
 *
 */
struct rse_vulkan_image_t
{
    uint8_t id_taken;
    size_t allocated_size;
    VkImage image;
    VkImageView image_view;
    VmaAllocation allocation;
    VmaAllocationInfo allocation_info;
};

/* For now we will only have one image view per image */
extern struct rse_vulkan_state_t g_vulkan_state;

struct rse_vulkan_image_t g_texture_images[RSE_MAX_IMAGE_COUNT];

struct rse_vulkan_image_t g_depth_image;
struct rse_vulkan_image_t g_color_image;
VkSampler g_sampler;


static rse_err_t create_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(g_vulkan_state.device,
                                   &create_info,
                                   NULL,
                                   &g_sampler)) {
        LOGF("Failed to create image view");
        return 1;
    }

    return RSE_ERROR_NO_ERROR;
}

static rse_err_t find_supported_format(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(g_vulkan_state.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;
        }
    }

    LOGE(vulkan_messages[VULKAN_FORMAT_NOT_FOUND]);
    return VULKAN_ERROR_FORMAT_NOT_FOUND;
}

static rse_err_t create_image(struct rse_vulkan_image_t* image,
                              uint32_t width,
                              uint32_t height,
                              VkFormat format,
                              VkSampleCountFlagBits maa_samples,
                              VkImageUsageFlags usage,
                              VmaAllocationCreateFlags allocationFlags)
{
    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;

    if (VK_SUCCESS != vmaCreateImage(g_vulkan_state.allocator,
                                     &image_info,
                                     &alloc_create_info,
                                     &image->image,
                                     &image->allocation,
                                     &image->allocation_info)) {
        LOGF(vulkan_messages[VULKAN_IMAGE_NOT_CREATED]);
        return VULKAN_ERROR_IMAGE_NOT_CREATED;
    }

    return RSE_ERROR_NO_ERROR;
}

static rse_err_t create_image_view(struct rse_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(g_vulkan_state.device,
                                        &image_view_create_info,
                                        NULL,
                                        &image->image_view)) {
        LOGF(vulkan_messages[VULKAN_SWAPCHAIN_IMVIEW_CREATION_FAILED]);
        return VULKAN_ERROR_SWAPCHAIN_IMVIEW_CREATION_FAILED;
    }
    return RSE_ERROR_NO_ERROR;
}

rse_err_t rse_create_color_resource()
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    VkFormat color_format = IMAGE_FORMAT;
    VkSurfaceCapabilitiesKHR physical_device_surface_capabilities;

    vkGetPhysicalDeviceSurfaceCapabilitiesKHR(g_vulkan_state.physical_device,
                                                g_vulkan_state.surface,
                                                &physical_device_surface_capabilities);

    STATUS_CHECK(create_image(&g_color_image,
                 physical_device_surface_capabilities.currentExtent.width,
                 physical_device_surface_capabilities.currentExtent.height,
                 color_format,
                 VK_SAMPLE_COUNT_16_BIT,
                 VK_IMAGE_USAGE_TRANSIENT_ATTACHMENT_BIT | VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT, 0U));
    STATUS_CHECK(create_image_view(&g_color_image, color_format, VK_IMAGE_ASPECT_COLOR_BIT));

    return status;
}


rse_err_t rse_create_depth_resources()
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    VkFormat depth_format = {0};
    VkSurfaceCapabilitiesKHR physical_device_surface_capabilities;


    STATUS_CHECK(find_depth_format(&depth_format));
    vkGetPhysicalDeviceSurfaceCapabilitiesKHR(g_vulkan_state.physical_device,
                                                g_vulkan_state.surface,
                                                &physical_device_surface_capabilities);

    STATUS_CHECK(create_image(&g_depth_image,
                              physical_device_surface_capabilities.currentExtent.width,
                              physical_device_surface_capabilities.currentExtent.height,
                              depth_format,
                              VK_SAMPLE_COUNT_16_BIT,
                              VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT,
                              0U));
    STATUS_CHECK(create_image_view(&g_depth_image, depth_format, VK_IMAGE_ASPECT_DEPTH_BIT));

    return status;
}

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_R8G8B8A8_SRGB:
        *format_size = 4;
        break;

    default:
        LOGE("Failed translate vulkan format to pixel width");
        break;
    }

    return RSE_ERROR_NO_ERROR;
}

static rse_err_t transition_image_layout(struct rse_vulkan_image_t* image, VkImageLayout old_layout, VkImageLayout new_layout)
{
    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 VULKAN_ERROR_TRANSITION_IMAGE_LAYOUT_FAILED;
    }

    command_buffer = rse_begin_single_time_command();
    vkCmdPipelineBarrier(command_buffer,
                        source_stage,
                        destination_stage,
                        0,
                        0,
                        NULL,
                        0,
                        NULL,
                        1,
                        &image_barrier);
    rse_end_single_time_comands(command_buffer);

    return RSE_ERROR_NO_ERROR;
}


static rse_err_t create_textured_image(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 rse_vulkan_buffer_t staging_buffer = {0};
    VkDeviceSize staging_buffer_size = 0U;
    VkCommandBuffer command_buffer = {0};
    VkBufferImageCopy copy_region = {0};
    VkExtent3D image_extent = {0};
    struct rse_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(g_texture_images[*texture_id].id_taken == IMAGE_FREE) {
            free_texture_image = &g_texture_images[*texture_id];
        } else {
            *texture_id += 1;
        }
    }

    /* Check if we found free image handle */
    if (free_texture_image == NULL) {
        LOGE(vulkan_messages[VULKAN_MAX_TEXTURE_COUNT_REACHED]);
        *texture_id = -1;
        return VULKAN_ERROR_MAX_TEXTURE_COUNT_REACHED;
    }

    STATUS_CHECK(format_to_pixel_size(format, &staging_buffer_size));

    staging_buffer_size = staging_buffer_size * width * height;

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

    // /* 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.srcAccessMask = 0U;
    // image_barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
    // image_barrier.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED;
    // image_barrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
    // image_barrier.srcQueueFamilyIndex = 0U;
    // image_barrier.dstQueueFamilyIndex = 0U;
    // image_barrier.image = free_texture_image->image;
    // image_barrier.subresourceRange = range;
    //
    // command_buffer = rse_begin_single_time_command();
    // vkCmdPipelineBarrier(command_buffer,
    //                     VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
    //                     VK_PIPELINE_STAGE_TRANSFER_BIT,
    //                     0,
    //                     0,
    //                     NULL,
    //                     0,
    //                     NULL,
    //                     1,
    //                     &image_barrier);
    // rse_end_single_time_comands(command_buffer);

    transition_image_layout(free_texture_image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
    /* Copy pixel data to GPU memory */
    STATUS_CHECK(create_buffer(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;

    command_buffer = rse_begin_single_time_command();
    vkCmdCopyBufferToImage(command_buffer,
                            staging_buffer.buffer,
                            free_texture_image->image,
                            VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                            1,
                            &copy_region);
    rse_end_single_time_comands(command_buffer);
    destroy_buffer(&staging_buffer);


    transition_image_layout(free_texture_image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

    // /* Change image layout for shaders to be able to read it */
    // image_barrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
    // image_barrier.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    // image_barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
    // image_barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
    //
    // command_buffer = rse_begin_single_time_command();
    // vkCmdPipelineBarrier(command_buffer,
    //                     VK_PIPELINE_STAGE_TRANSFER_BIT,
    //                     VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
    //                     0,
    //                     0,
    //                     NULL,
    //                     0,
    //                     NULL,
    //                     1,
    //                     &image_barrier);
    // rse_end_single_time_comands(command_buffer);

    STATUS_CHECK(create_image_view(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()
{
    rse_err_t status = RSE_ERROR_NO_ERROR;

    STATUS_CHECK(create_sampler());
    STATUS_CHECK(rse_create_depth_resources());
    STATUS_CHECK(rse_create_color_resource());

    return status;
}

rse_err_t rse_load_texture_from_bitmat(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(width,
                                        height,
                                        color_format,
                                        pixel_buffer, texture_id));

    return status;
}

rse_err_t rse_load_texture_from_file(const char* file_path, uint16_t* texture_id)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    int width = 0;
    int height = 0;
    int channels = 0;
    stbi_uc* pixel_buffer = NULL;

    pixel_buffer = stbi_load(file_path, &width, &height, &channels, STBI_rgb_alpha);

    if (pixel_buffer == NULL) {
        LOGE("Failed to load texture from file: %s", file_path);
        return -1;
    }

    STATUS_CHECK(rse_load_texture_from_bitmat(pixel_buffer, width, height, texture_id, VK_FORMAT_R8G8B8A8_SRGB));

    stbi_image_free(pixel_buffer);

    return status;
}

void rse_destroy_textures()
{
    size_t i = 0U;

    vkDestroySampler(g_vulkan_state.device, g_sampler, NULL);

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

void rse_destroy_depth_resource()
{
    vkDestroyImageView(g_vulkan_state.device, g_depth_image.image_view, NULL);
    vmaDestroyImage(g_vulkan_state.allocator, g_depth_image.image, g_depth_image.allocation);
}

void rse_destroy_color_resource()
{
    vkDestroyImageView(g_vulkan_state.device, g_color_image.image_view, NULL);
    vmaDestroyImage(g_vulkan_state.allocator, g_color_image.image, g_color_image.allocation);
}

uint8_t image_exists(uint8_t image_id) {
    return g_texture_images[image_id].id_taken == IMAGE_TAKEN;
}

VkImageView rse_get_image_view(size_t image_id)
{
    return g_texture_images[image_id].image_view;
}

VkImageView rse_get_depth_image_view()
{
    return g_depth_image.image_view;
}

VkImageView rse_get_color_image_view()
{
    return g_color_image.image_view;
}

VkSampler rse_get_texture_sampler()
{
    return g_sampler;
}

size_t get_textures_count()
{
    size_t count = 0U;
    size_t i = 0U;

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

    return count;
}

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

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

    return status;
    #undef formats_count
}