#include <runtime/config.h>
#include <stdlib.h>

#include "device.h"
#include "physical_resource_manager.h"

#include <unistd.h>

RT_CVAR_SZ(r_VkMaxResources,
           "Maximum number of simulaneously existing resources. (Default: 4096)",
           4096);

static VmaAllocator CreateAllocator(rt_vk_device *dev) {
#define SET_FNC(name)     fncs.name = name
#define SET_KHR_FNC(name) (fncs).name##KHR = name
    VmaVulkanFunctions fncs                = {NULL};
    SET_FNC(vkGetInstanceProcAddr);
    SET_FNC(vkGetDeviceProcAddr);
    SET_FNC(vkGetPhysicalDeviceProperties);
    SET_FNC(vkGetPhysicalDeviceMemoryProperties);
    SET_FNC(vkAllocateMemory);
    SET_FNC(vkFreeMemory);
    SET_FNC(vkMapMemory);
    SET_FNC(vkUnmapMemory);
    SET_FNC(vkFlushMappedMemoryRanges);
    SET_FNC(vkInvalidateMappedMemoryRanges);
    SET_FNC(vkBindBufferMemory);
    SET_FNC(vkBindImageMemory);
    SET_FNC(vkGetBufferMemoryRequirements);
    SET_FNC(vkGetImageMemoryRequirements);
    SET_FNC(vkCreateBuffer);
    SET_FNC(vkDestroyBuffer);
    SET_FNC(vkCreateImage);
    SET_FNC(vkDestroyImage);
    SET_FNC(vkCmdCopyBuffer);
    SET_KHR_FNC(vkGetBufferMemoryRequirements2);
    SET_KHR_FNC(vkGetImageMemoryRequirements2);
    SET_KHR_FNC(vkBindBufferMemory2);
    SET_KHR_FNC(vkBindImageMemory2);
    SET_KHR_FNC(vkGetPhysicalDeviceMemoryProperties2);
    SET_FNC(vkGetDeviceBufferMemoryRequirements);
    SET_FNC(vkGetDeviceImageMemoryRequirements);
#undef SET_FNC
#undef SET_KHR_FNC

    VmaAllocatorCreateInfo allocator_info = {
        .instance             = dev->instance,
        .physicalDevice       = dev->phys_device,
        .device               = dev->device,
        .pAllocationCallbacks = dev->alloc_cb,
        .vulkanApiVersion     = VK_API_VERSION_1_3,
        .pVulkanFunctions     = &fncs,
    };

    VmaAllocator allocator;
    if (vmaCreateAllocator(&allocator_info, &allocator) != VK_SUCCESS)
        return NULL;
    return allocator;
}

rt_create_vk_physical_resource_manager_result rtCreateVkPhysicalResourceManager(struct rt_vk_device *dev) {
    rt_vk_physical_resource_manager phys_res_mgr;
    phys_res_mgr.dev = dev;
    phys_res_mgr.allocator = CreateAllocator(dev);
    if (!phys_res_mgr.allocator) {
        return (rt_create_vk_physical_resource_manager_result){.result =  RT_UNKNOWN_ERROR};
    }

    rt_create_rwlock_result lock_res = rtCreateRWLock();
    if (!lock_res.ok) {
        vmaDestroyAllocator(phys_res_mgr.allocator);
        return (rt_create_vk_physical_resource_manager_result){.result =  RT_UNKNOWN_ERROR};
    }
    phys_res_mgr.lock = lock_res.lock;

    phys_res_mgr.resources = calloc(r_VkMaxResources.sz, sizeof(rt_vk_physical_resource));
    if (!phys_res_mgr.resources) {
        rtDestroyRWLock(&phys_res_mgr.lock);
        vmaDestroyAllocator(phys_res_mgr.allocator);
        return (rt_create_vk_physical_resource_manager_result){.result =  RT_OUT_OF_MEMORY};
    }
    void *lut_mem = malloc(RT_HASH_TABLE_MEMORY_REQUIRED(2 * r_VkMaxResources.sz));
    if (!lut_mem) {
        free(phys_res_mgr.resources);
        rtDestroyRWLock(&phys_res_mgr.lock);
        vmaDestroyAllocator(phys_res_mgr.allocator);
        return (rt_create_vk_physical_resource_manager_result){.result =  RT_OUT_OF_MEMORY};
    }
    phys_res_mgr.resource_lut = rtCreateHashtable(2 * r_VkMaxResources.sz, lut_mem, NULL, NULL);
    phys_res_mgr.lut_mem = lut_mem;

    phys_res_mgr.free_slots = calloc(r_VkMaxResources.sz, sizeof(uint32_t));
    if (!phys_res_mgr.free_slots) {
        free(phys_res_mgr.lut_mem);
        free(phys_res_mgr.resources);
        rtDestroyRWLock(&phys_res_mgr.lock);
        vmaDestroyAllocator(phys_res_mgr.allocator);
        return (rt_create_vk_physical_resource_manager_result){.result =  RT_OUT_OF_MEMORY};
    }
    for (uint32_t i = 0; i < r_VkMaxResources.sz; ++i) {
        phys_res_mgr.free_slots[i] = i;
    }
    phys_res_mgr.free_slot_count = r_VkMaxResources.sz;

    /* TODO: Create pools for transient resources. */

    return (rt_create_vk_physical_resource_manager_result){.result =  RT_SUCCESS,
        .phys_res_mgr = phys_res_mgr};
}

void rtDestroyVkPhysicalResourceManager(rt_vk_physical_resource_manager *phys_res_mgr) {
    if (phys_res_mgr->free_slot_count < r_VkMaxResources.sz) {
        rtReportError("VK", "Destroyed the physical resource manager, but there are still resources left.");
    }
    free(phys_res_mgr->free_slots);
    free(phys_res_mgr->lut_mem);
    free(phys_res_mgr->resources);
    vmaDestroyAllocator(phys_res_mgr->allocator);
}

bool rtVkPhysicalResourceManagerIsPresent(void *o, rt_render_resource_handle handle) {
    rt_vk_physical_resource_manager *phys_res_mgr = o;
    rtLockRead(&phys_res_mgr->lock);
    bool is_present = rtHashtableLookup(&phys_res_mgr->resource_lut, (uint64_t)handle.value, UINT64_MAX) != UINT64_MAX;
    rtUnlockRead(&phys_res_mgr->lock);
    return is_present;
}

void rtVkPhysicalResourceManagerDestroy(void *o, rt_render_resource_handle handle) {
    rt_vk_physical_resource_manager *phys_res_mgr = o;
    rtLockWrite(&phys_res_mgr->lock);
    uint32_t slot = rtHashtableLookup(&phys_res_mgr->resource_lut, handle.value, UINT32_MAX);
    if (slot != UINT32_MAX) {
        switch (phys_res_mgr->resources[slot].type) {
        case RT_RENDER_RESOURCE_TYPE_BUFFER: {
            vmaDestroyBuffer(phys_res_mgr->allocator, phys_res_mgr->resources[slot].buffer, phys_res_mgr->resources[slot].allocation);
        } break;
        case RT_RENDER_RESOURCE_TYPE_TEXTURE2D: {
            RT_NOT_IMPLEMENTED;
        } break;
        default: {
            rtReportError("VK", "Tried to destroy a resource, but the type is not valid.");
        }
        }
        phys_res_mgr->free_slots[phys_res_mgr->free_slot_count++] = slot;
    }
    rtUnlockWrite(&phys_res_mgr->lock);
}

rt_vk_physical_resource *rtGetVkPhysicalResource(rt_vk_physical_resource_manager *phys_res_mgr, rt_render_resource_handle handle) {
    rt_vk_physical_resource *resource = NULL;
    rtLockRead(&phys_res_mgr->lock);
    uint32_t slot = rtHashtableLookup(&phys_res_mgr->resource_lut, handle.value, UINT32_MAX);
    if (slot != UINT32_MAX) {
        resource = &phys_res_mgr->resources[slot];
    }
    rtUnlockRead(&phys_res_mgr->lock);
    return resource;
}

/* Call this with a held write lock! */
static uint32_t AllocStorageSlot(rt_vk_physical_resource_manager *phys_res_mgr, rt_render_resource_handle h) {
    if (phys_res_mgr->free_slot_count > 0u) {
        uint32_t slot = phys_res_mgr->free_slots[--phys_res_mgr->free_slot_count];
        /* The hashtable is large enough that this should never fail */
        rt_result insert_res = rtHashtableInsert(&phys_res_mgr->resource_lut, (uint64_t)h.value, (uint64_t)slot);
        RT_UNUSED(insert_res);
        RT_VERIFY(insert_res == RT_SUCCESS);
        return slot;
    }
    return UINT32_MAX;
}

rt_result rtVkPhysicalResourceManagerCreateBuffer(void *o, rt_render_resource_handle h, const rt_render_buffer_desc *desc) {
    rt_vk_physical_resource_manager *phys_res_mgr = o;

    VkBufferUsageFlags usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT;
    if ((desc->usage & RT_RENDER_BUFFER_USAGE_INDEX_BUFFER) != 0)
        usage |= VK_BUFFER_USAGE_INDEX_BUFFER_BIT;
    if ((desc->usage & RT_RENDER_BUFFER_USAGE_VERTEX_BUFFER) != 0)
        usage |= VK_BUFFER_USAGE_VERTEX_BUFFER_BIT;
    if ((desc->usage & RT_RENDER_BUFFER_USAGE_STORAGE_BUFFER) != 0)
        usage |= VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
    if ((desc->usage & RT_RENDER_BUFFER_USAGE_UPLOAD_BUFFER) != 0) {
        usage &= ~VK_BUFFER_USAGE_TRANSFER_DST_BIT;
        usage |= VK_BUFFER_USAGE_TRANSFER_SRC_BIT;
    }

    VkBufferCreateInfo buffer_info = {
        .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
        .size = desc->size,
        .usage = usage,
        .sharingMode = VK_SHARING_MODE_CONCURRENT,
        .pQueueFamilyIndices = phys_res_mgr->dev->unique_families,
        .queueFamilyIndexCount = phys_res_mgr->dev->unique_family_count,
    };

    VmaAllocationCreateFlags alloc_flags = 0;
    if (desc->access == RT_RENDER_BUFFER_ACCESS_CPU_AND_GPU)
        alloc_flags |= VMA_ALLOCATION_CREATE_HOST_ACCESS_RANDOM_BIT;

    VmaAllocationCreateInfo alloc_info = {
        .usage = VMA_MEMORY_USAGE_AUTO,
        .flags = alloc_flags,
    };

    VkBuffer buffer;
    VmaAllocation allocation;
    if (vmaCreateBuffer(phys_res_mgr->allocator, &buffer_info, &alloc_info, &buffer, &allocation, NULL) != VK_SUCCESS) {
        rtLog("VK", "Failed to create buffer.");
        return RT_UNKNOWN_ERROR;
    }

    /* Store */
    rt_result res = RT_SUCCESS;
    rtLockWrite(&phys_res_mgr->lock);
    uint32_t slot = AllocStorageSlot(phys_res_mgr, h);
    if (slot != UINT32_MAX) {
        phys_res_mgr->resources[slot].type = RT_RENDER_RESOURCE_TYPE_BUFFER;
        phys_res_mgr->resources[slot].buffer = buffer;
        phys_res_mgr->resources[slot].allocation = allocation;
    } else {
        res = RT_OUT_OF_MEMORY;
        vmaDestroyBuffer(phys_res_mgr->allocator, buffer, allocation);
        rtLog("VK","Could not create buffer because no storage space is available.");
    }
    rtUnlockWrite(&phys_res_mgr->lock);

    return res;
}

static VkFormat RtTextureFormatToVkFormat(rt_texture_format texture_format) {
    RT_ASSERT(texture_format < RT_TEXTURE_FORMAT_MAX, "Invalid format");
    VkFormat formats[RT_TEXTURE_FORMAT_MAX] = {
        VK_FORMAT_B8G8R8A8_SRGB, // RT_TEXTURE_FORMAT_B8G8R8A8_SRGB
    };
    return formats[texture_format];
}

static VkImageUsageFlagBits RtTextureUsageToVkImageUsage(rt_texture_usage_flags usage_flags) {
    VkImageUsageFlagBits usage = 0;
    if ((usage_flags & RT_TEXTURE_USAGE_COLOR_ATTACHMENT) != 0)
        usage |= VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;
    if ((usage_flags & RT_TEXTURE_USAGE_DEPTH_STENCIL_ATTACHMENT) != 0)
        usage |= VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT;
    if ((usage_flags & RT_TEXTURE_USAGE_SAMPLED_IMAGE) != 0)
        usage |= VK_IMAGE_USAGE_SAMPLED_BIT;
    if ((usage_flags & RT_TEXTURE_USAGE_STORAGE_IMAGE) != 0)
        usage |= VK_IMAGE_USAGE_STORAGE_BIT;
    return usage;
}

rt_result rtVkPhysicalResourceManagerCreateTexture2D(void *o, rt_render_resource_handle h, const rt_render_texture2d_desc *desc) {
    rt_vk_physical_resource_manager *phys_res_mgr = o;

    VkFormat format = RtTextureFormatToVkFormat(desc->format);
    VkExtent3D extent = {.width = desc->width, .height = desc->height, .depth = 1};
    VkImageUsageFlagBits usage = RtTextureUsageToVkImageUsage(desc->usage);

    VkImageCreateInfo image_info = {
        .sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
        .extent = extent,
        .format = format,
        .imageType = VK_IMAGE_TYPE_2D,
        .sharingMode = VK_SHARING_MODE_CONCURRENT,
        .queueFamilyIndexCount = phys_res_mgr->dev->unique_family_count,
        .pQueueFamilyIndices = &phys_res_mgr->dev->unique_families[0],
        .mipLevels = 1,
        .samples = VK_SAMPLE_COUNT_1_BIT,
        .arrayLayers = 1,
        .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED,
        .usage = usage,
        .tiling = VK_IMAGE_TILING_OPTIMAL,
    };

    VmaAllocationCreateInfo alloc_info = {
        .usage = VMA_MEMORY_USAGE_AUTO,
    };

    VkImage image;
    VmaAllocation allocation;
    if (vmaCreateImage(phys_res_mgr->allocator, &image_info, &alloc_info, &image, &allocation, NULL) != VK_SUCCESS) {
        rtLog("VK", "Failed to create image.");
        return RT_UNKNOWN_ERROR;
    }

    VkImageViewCreateInfo view_info ={
        .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
        .format = format,
        .viewType = VK_IMAGE_VIEW_TYPE_2D,
        .subresourceRange = {
            .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
            .levelCount = 1,
            .baseMipLevel = 0,
            .layerCount = 1,
            .baseArrayLayer = 0
        },
        .components = {
            .r = VK_COMPONENT_SWIZZLE_IDENTITY,
            .g = VK_COMPONENT_SWIZZLE_IDENTITY,
            .b = VK_COMPONENT_SWIZZLE_IDENTITY,
            .a = VK_COMPONENT_SWIZZLE_IDENTITY
        },
    };
    VkImageView view;
    if (vkCreateImageView(phys_res_mgr->dev->device, &view_info, phys_res_mgr->dev->alloc_cb, &view) != VK_SUCCESS) {
        vmaDestroyImage(phys_res_mgr->allocator, image, allocation);
        rtLog("VK","Failed to create image view.");
        return RT_UNKNOWN_ERROR;
    }

    /* Store */
    rt_result res = RT_SUCCESS;
    rtLockWrite(&phys_res_mgr->lock);
    uint32_t slot = AllocStorageSlot(phys_res_mgr, h);
    if (slot != UINT32_MAX) {
        phys_res_mgr->resources[slot].type = RT_RENDER_RESOURCE_TYPE_BUFFER;
        phys_res_mgr->resources[slot].texture2d.image = image;
        phys_res_mgr->resources[slot].texture2d.view = view;
        phys_res_mgr->resources[slot].allocation = allocation;
    } else {
        res = RT_OUT_OF_MEMORY;
        vkDestroyImageView(phys_res_mgr->dev->device, view, phys_res_mgr->dev->alloc_cb);
        vmaDestroyImage(phys_res_mgr->allocator, image, allocation);
        rtLog("VK","Could not create image because no storage space is available.");
    }
    rtUnlockWrite(&phys_res_mgr->lock);

    return res;
}