#include <cstdio>
#include <cstring>
#include <vulkan/vulkan.hpp>
#include <vulkan/vk_layer.h>
#include "dispatch_table.hpp"
#include "functions.hpp"
#include "record_list.hpp"
#include "vk_capture.h"

namespace vk_capture
{
inline constexpr std::uint32_t LOADER_LAYER_VERSION = 2;

PFN_vkVoidFunction getLayerFunctionPtr(const char* pName);

VkResult vkCreateInstance(const VkInstanceCreateInfo* pCreateInfo, const VkAllocationCallbacks* pAllocator, VkInstance* pInstance)
{
    const VkLayerInstanceCreateInfo* layerCreateInfo = static_cast<const VkLayerInstanceCreateInfo*>(pCreateInfo->pNext);

    // step through the chain of pNext until we get to the link info
    while(layerCreateInfo && (layerCreateInfo->sType != VK_STRUCTURE_TYPE_LOADER_INSTANCE_CREATE_INFO ||
                            layerCreateInfo->function != VK_LAYER_LINK_INFO))
    {
        layerCreateInfo = static_cast<const VkLayerInstanceCreateInfo*>(layerCreateInfo->pNext);
    }

    if(layerCreateInfo == nullptr)
    {
        // No loader instance create info
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    PFN_vkGetInstanceProcAddr gpa = layerCreateInfo->u.pLayerInfo->pfnNextGetInstanceProcAddr;
    // move chain on for next layer
    const_cast<VkLayerInstanceCreateInfo*>(layerCreateInfo)->u.pLayerInfo = layerCreateInfo->u.pLayerInfo->pNext;

    PFN_vkCreateInstance createFunc = reinterpret_cast<PFN_vkCreateInstance>(gpa(VK_NULL_HANDLE, "vkCreateInstance"));

    VkResult ret = createFunc(pCreateInfo, pAllocator, pInstance);
    if (ret != VK_SUCCESS)
    {
        return ret;
    }

    // fetch our own dispatch table for the functions we need, into the next layer
    layer_init_instance_dispatch_table(*pInstance, &vk_capture::g_instanceDispatchTable, gpa);

    return VK_SUCCESS;
}

VkResult vkCreateDevice(
    VkPhysicalDevice                            physicalDevice,
    const VkDeviceCreateInfo*                   pCreateInfo,
    const VkAllocationCallbacks*                pAllocator,
    VkDevice*                                   pDevice)
{
    const VkLayerDeviceCreateInfo* layerCreateInfo = static_cast<const VkLayerDeviceCreateInfo*>(pCreateInfo->pNext);

    // step through the chain of pNext until we get to the link info
    while(layerCreateInfo && (layerCreateInfo->sType != VK_STRUCTURE_TYPE_LOADER_DEVICE_CREATE_INFO ||
                            layerCreateInfo->function != VK_LAYER_LINK_INFO))
    {
        layerCreateInfo = static_cast<const VkLayerDeviceCreateInfo*>(layerCreateInfo->pNext);
    }

    if(layerCreateInfo == nullptr)
    {
        // No loader instance create info
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    PFN_vkGetInstanceProcAddr gipa = layerCreateInfo->u.pLayerInfo->pfnNextGetInstanceProcAddr;
    PFN_vkGetDeviceProcAddr gdpa = layerCreateInfo->u.pLayerInfo->pfnNextGetDeviceProcAddr;
    // move chain on for next layer
    const_cast<VkLayerDeviceCreateInfo*>(layerCreateInfo)->u.pLayerInfo = layerCreateInfo->u.pLayerInfo->pNext;

    PFN_vkCreateDevice createFunc = (PFN_vkCreateDevice)gipa(VK_NULL_HANDLE, "vkCreateDevice");

    VkResult ret = createFunc(physicalDevice, pCreateInfo, pAllocator, pDevice);

    if (ret != VK_SUCCESS)
    {
        return ret;
    }
    layer_init_device_dispatch_table(*pDevice, &vk_capture::g_dispatchTable, gdpa);

    return VK_SUCCESS;
}

PFN_vkVoidFunction vkGetDeviceProcAddr(VkDevice device, const char* pName)
{
    PFN_vkVoidFunction ptr = getLayerFunctionPtr(pName);
    if (ptr) {
        return ptr;
    }
    else {
        return g_dispatchTable.GetDeviceProcAddr(device, pName);
    }
}

PFN_vkVoidFunction vkGetInstanceProcAddr(VkInstance instance, const char* pName)
{
    PFN_vkVoidFunction ptr = getLayerFunctionPtr(pName);
    if (ptr) {
        return ptr;
    }
    else {
        return g_instanceDispatchTable.GetInstanceProcAddr(instance, pName);
    }
}

PFN_vkVoidFunction getLayerFunctionPtr(const char* pName)
{
    if (std::strcmp(pName, "vkCreateInstance") == 0) {
        return reinterpret_cast<PFN_vkVoidFunction>(&vk_capture::vkCreateInstance);
    }
    else if (std::strcmp(pName, "vkCreateDevice") == 0) {
        return reinterpret_cast<PFN_vkVoidFunction>(&vk_capture::vkCreateDevice);
    }
    else if (std::strcmp(pName, "vkGetInstanceProcAddr") == 0) {
        return reinterpret_cast<PFN_vkVoidFunction>(&vk_capture::vkGetInstanceProcAddr);
    }
    if (std::strcmp(pName, "vkGetDeviceProcAddr") == 0) {
        return reinterpret_cast<PFN_vkVoidFunction>(&vkGetDeviceProcAddr);
    }
    else if (std::strcmp(pName, "vkAllocateRecordListMWN") == 0) {
        return reinterpret_cast<PFN_vkVoidFunction>(&vkAllocateRecordListMWN);
    }
    else if (std::strcmp(pName, "vkFreeRecordListMWN") == 0) {
        return reinterpret_cast<PFN_vkVoidFunction>(&vkFreeRecordListMWN);
    }
    else if (std::strcmp(pName, "vkBeginRecordingMWN") == 0) {
        return reinterpret_cast<PFN_vkVoidFunction>(&vkBeginRecordingMWN);
    }
    else if (std::strcmp(pName, "vkEndRecordingMWN") == 0) {
        return reinterpret_cast<PFN_vkVoidFunction>(&vkEndRecordingMWN);
    }
    else if (std::strcmp(pName, "vkGetRecordListItemsMWN") == 0) {
        return reinterpret_cast<PFN_vkVoidFunction>(&vkGetRecordListItemsMWN);
    }
    else if (std::strcmp(pName, "vkResetRecordListMWN") == 0) {
        return reinterpret_cast<PFN_vkVoidFunction>(&vkResetRecordListMWN);
    }
    else if (std::strcmp(pName, "vkGetFunctionNameMWN") == 0) {
        return reinterpret_cast<PFN_vkVoidFunction>(&vkGetFunctionNameMWN);
    }
    else {
        return getWrappedFunctionPtr(pName);
    }
}
} // namespace vk_capture

extern "C"
{

VK_LAYER_EXPORT VkResult VKAPI_CALL vk_capture_vkNegotiateLoaderLayerInterfaceVersion(VkNegotiateLayerInterface* pVersionStruct)
{
    if (pVersionStruct->loaderLayerInterfaceVersion < vk_capture::LOADER_LAYER_VERSION) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }
    pVersionStruct->sType = LAYER_NEGOTIATE_INTERFACE_STRUCT;
    pVersionStruct->pNext = nullptr;
    pVersionStruct->loaderLayerInterfaceVersion = vk_capture::LOADER_LAYER_VERSION;
    pVersionStruct->pfnGetInstanceProcAddr = &vk_capture::vkGetInstanceProcAddr;
    pVersionStruct->pfnGetDeviceProcAddr = &vk_capture::vkGetDeviceProcAddr;
    pVersionStruct->pfnGetPhysicalDeviceProcAddr = nullptr;
    return VK_SUCCESS;
}
} // extern "C"