/*
Index of this file:
// [SECTION] quick notes
// [SECTION] includes
// [SECTION] structs
// [SECTION] apis
// [SECTION] pl_app_load
// [SECTION] pl_app_shutdown
// [SECTION] pl_app_resize
// [SECTION] pl_app_update
*/

//-----------------------------------------------------------------------------
// [SECTION] includes
//-----------------------------------------------------------------------------

#include <stdlib.h> // malloc, free
#include <string.h> // memset
#include "pl.h"

#define PL_MATH_INCLUDE_FUNCTIONS // required to expose some of the color helpers
#include "pl_math.h"

// extensions
#define PL_GRAPHICS_EXPOSE_VULKAN
#include "pl_graphics_ext.h"
#include "pl_shader_ext.h"
#include "pl_starter_ext.h"
#include "pl_image_ext.h"
#include "pl_vfs_ext.h"
#include "pl_screen_log_ext.h"

#include "vulkan/vulkan.h"

#define VULKAN_CHECK(x) assert(x == VK_SUCCESS)

//-----------------------------------------------------------------------------
// [SECTION] structs
//-----------------------------------------------------------------------------

typedef struct _plAppData
{
    // window
    plWindow* ptWindow;
    plSurface* ptSurface;
    plDevice* ptDevice;

    // shaders
    plShaderHandle tShader;

} plAppData;

//-----------------------------------------------------------------------------
// [SECTION] apis
//-----------------------------------------------------------------------------

const plIOI*        gptIO        = NULL;
const plWindowI*    gptWindows   = NULL;
const plGraphicsI*  gptGfx       = NULL;
const plShaderI*    gptShader    = NULL;
const plStarterI*   gptStarter   = NULL;
const plImageI*     gptImage     = NULL;
const plVfsI*       gptVfs       = NULL;
const plScreenLogI* gptScreenLog = NULL;

//-----------------------------------------------------------------------------
// [SECTION] pl_app_load
//-----------------------------------------------------------------------------

PL_EXPORT void*
pl_app_load(plApiRegistryI* ptApiRegistry, plAppData* ptAppData)
{

    // this path is taken only during first load, so we
    // allocate app memory here
    ptAppData = malloc(sizeof(plAppData));
    memset(ptAppData, 0, sizeof(plAppData));

    // retrieve extension registry
    const plExtensionRegistryI* ptExtensionRegistry = pl_get_api_latest(ptApiRegistry, plExtensionRegistryI);

    // load extensions
    ptExtensionRegistry->add_path("../../pl-template/out");
    ptExtensionRegistry->load("pl_unity_ext", NULL, NULL, true);
    ptExtensionRegistry->load("pl_platform_ext", NULL, NULL, false);
    
    // load required apis
    gptIO      = pl_get_api_latest(ptApiRegistry, plIOI);
    gptWindows = pl_get_api_latest(ptApiRegistry, plWindowI);

    // load required apis (these are provided though extensions)
    gptGfx       = pl_get_api_latest(ptApiRegistry, plGraphicsI);
    gptShader    = pl_get_api_latest(ptApiRegistry, plShaderI);
    gptStarter   = pl_get_api_latest(ptApiRegistry, plStarterI);
    gptImage     = pl_get_api_latest(ptApiRegistry, plImageI);
    gptVfs       = pl_get_api_latest(ptApiRegistry, plVfsI);
    gptScreenLog = pl_get_api_latest(ptApiRegistry, plScreenLogI);

    // use window API to create a window
    plWindowDesc tWindowDesc = {
        .pcTitle = "Vulkan Example",
        .iXPos   = 200,
        .iYPos   = 200,
        .uWidth  = 600,
        .uHeight = 600,
    };
    gptWindows->create(tWindowDesc, &ptAppData->ptWindow);
    gptWindows->show(ptAppData->ptWindow);

    plGraphicsInit tGraphicsDesc = {
        .tFlags = PL_GRAPHICS_INIT_FLAGS_SWAPCHAIN_ENABLED | PL_GRAPHICS_INIT_FLAGS_VALIDATION_ENABLED
    };
    gptGfx->initialize(&tGraphicsDesc);

    ptAppData->ptSurface = gptGfx->create_surface(ptAppData->ptWindow);

    ptAppData->ptDevice = gptStarter->create_device(ptAppData->ptSurface);
    
    // initialize shader extension (we are doing this ourselves so we can add additional shader directories)
    static const plShaderOptions tDefaultShaderOptions = {
        .apcIncludeDirectories = {
            "../../pl-template/shaders/"
        },
        .apcDirectories = {
            "../shaders/",
            "../../pl-template/shaders/"
        },
        .tFlags = PL_SHADER_FLAGS_AUTO_OUTPUT | PL_SHADER_FLAGS_ALWAYS_COMPILE
    };
    gptShader->initialize(&tDefaultShaderOptions);
    return ptAppData;
}

//-----------------------------------------------------------------------------
// [SECTION] pl_app_shutdown
//-----------------------------------------------------------------------------

PL_EXPORT void
pl_app_shutdown(plAppData* ptAppData)
{
    gptGfx->flush_device(ptAppData->ptDevice);
    gptShader->cleanup();
    gptWindows->destroy(ptAppData->ptWindow);
    gptScreenLog->cleanup();
    gptGfx->cleanup_surface(ptAppData->ptSurface);
    gptGfx->cleanup_device(ptAppData->ptDevice);
    gptGfx->cleanup();
    free(ptAppData);
}

//-----------------------------------------------------------------------------
// [SECTION] pl_app_resize
//-----------------------------------------------------------------------------

PL_EXPORT void
pl_app_resize(plAppData* ptAppData)
{

}

//-----------------------------------------------------------------------------
// [SECTION] pl_app_update
//-----------------------------------------------------------------------------

PL_EXPORT void
pl_app_update(plAppData* ptAppData)
{
    gptIO->new_frame(); // must be called once at the beginning of a frame

    plIO* ptIO = gptIO->get_io();
    ptIO->bRunning = false;

    // core (basically every vulkan application will have these)

    // global resources (for simplicity, we are treating these as global resources)
    VkDescriptorPool                 gtDescriptorPool;       // pool of descriptors
    VkCommandPool                    gtComputeCommandPool;   // pool of command buffers
    VkFence                          gtComputeWorkFence;     // fence to signal compute work is done (not needed for this example but required by API)

    // example specific resources
    VkBuffer                         gtStorageBuffer;       // storage buffer (reading & writing for this example)
    VkDeviceMemory                   gtStorageBufferMemory; // storage buffer actual device memory
    VkDescriptorSet                  gtDescriptorSet;       // descriptor set which will contain a descriptor to the storage buffer
    float*                           gpfMemoryMap;          // persistant mapping to the storage buffer

    // example specific compute pipeline stuff
    VkDescriptorSetLayout            gtDescriptorSetLayout; // describes the compute shader descriptor set layout (single set in this example)
    VkPipelineLayout                 gtPipelineLayout;      // describes the compute shader pipeline layout (descriptor set layout + other stuff in general) 
    VkPipeline                       gtPipeline;            // complete pipeline state object for compute shader (this really is the "shader")
    VkShaderModule                   gtShaderModule;        // SPIR-V code

    VkDevice gtLogicalDevice = gptGfx->get_vulkan_device(ptAppData->ptDevice);

    plShaderModule tShader = gptShader->load_glsl("shader.comp", "main", NULL, NULL);

    //-----------------------------------------------------------------------------
    // load an image
    //-----------------------------------------------------------------------------

    size_t szImageFileSize = gptVfs->get_file_size_str("../../pl-template/data/test.hdr");
    unsigned char* pucBuffer = malloc(szImageFileSize);
    plVfsFileHandle tEnvMapHandle = gptVfs->open_file("../../pl-template/data/test.hdr", PL_VFS_FILE_MODE_READ);
    gptVfs->read_file(tEnvMapHandle, pucBuffer, &szImageFileSize);
    gptVfs->close_file(tEnvMapHandle);

    int iImageWidth = 0;
    int iImageHeight = 0;
    int iImageChannels = 0;
    float* pfRawImageBytes = gptImage->load_hdr(pucBuffer, (int)szImageFileSize, &iImageWidth, &iImageHeight, &iImageChannels, 4);
    assert(pfRawImageBytes);
    free(pucBuffer);

    //-----------------------------------------------------------------------------
    // create storage buffer
    //   - we are using a storage buffer for this example to read & write from
    //   - we are also storing it in host visible memory for simplicity
    //-----------------------------------------------------------------------------

    // create storage buffer
    VkBufferCreateInfo tBufferInfo = {
        .sType       = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
        .size        = iImageWidth * iImageHeight * 4 * sizeof(float),
        .usage       = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
        .sharingMode = VK_SHARING_MODE_EXCLUSIVE
    };
    VULKAN_CHECK(vkCreateBuffer(gtLogicalDevice, &tBufferInfo, NULL, &gtStorageBuffer));

    // get buffer memory requirements
    VkMemoryRequirements tMemRequirements = {0};
    vkGetBufferMemoryRequirements(gtLogicalDevice, gtStorageBuffer, &tMemRequirements);

    // allocate actual device memory
    VkMemoryAllocateInfo tStorageAllocInfo = {
        .sType           = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
        .allocationSize  = tMemRequirements.size,
    };

    // find suitable memory type
    VkPhysicalDeviceMemoryProperties gtMemProps = gptGfx->get_vulkan_memory_properties(ptAppData->ptDevice);
    for (uint32_t i = 0; i < gtMemProps.memoryTypeCount; i++) 
    {
        if ((tMemRequirements.memoryTypeBits & (1 << i)) && (gtMemProps.memoryTypes[i].propertyFlags & (VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)) == (VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)) 
        {
            tStorageAllocInfo.memoryTypeIndex = i;
            break;
        }
    }
    VULKAN_CHECK(vkAllocateMemory(gtLogicalDevice, &tStorageAllocInfo, NULL, &gtStorageBufferMemory));

    // bind memory to resource
    VULKAN_CHECK(vkBindBufferMemory(gtLogicalDevice, gtStorageBuffer, gtStorageBufferMemory, 0));

    //-----------------------------------------------------------------------------
    // map to the memory persistantly
    //   - this is really just to make the example simpler
    //   - more generally, you would need to transfer the data & syncronize access
    //-----------------------------------------------------------------------------

    VULKAN_CHECK(vkMapMemory(gtLogicalDevice, gtStorageBufferMemory, 0, iImageWidth * iImageHeight * 4 * sizeof(float), 0, (void**)&gpfMemoryMap));
    memcpy(gpfMemoryMap, pfRawImageBytes, iImageWidth * iImageHeight * 4 * sizeof(float));
    gptImage->free(pfRawImageBytes);

    //-----------------------------------------------------------------------------
    // create descriptor pool, layout, & set
    //-----------------------------------------------------------------------------

    // create descriptor pool
    VkDescriptorPoolSize tPoolSize = {VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 100 };
    const VkDescriptorPoolCreateInfo tDescriptorPoolInfo = {
        .sType         = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
        .flags         = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT,
        .maxSets       = 100,
        .poolSizeCount = 1,
        .pPoolSizes    = &tPoolSize,
    };
    VULKAN_CHECK(vkCreateDescriptorPool(gtLogicalDevice, &tDescriptorPoolInfo, NULL, &gtDescriptorPool));

    // create descriptor set layout
    const VkDescriptorSetLayoutBinding tBindings = {
        .binding         = 0u,
        .descriptorType  = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        .descriptorCount = 1,
        .stageFlags      = VK_SHADER_STAGE_COMPUTE_BIT
    };

    const VkDescriptorSetLayoutCreateInfo tLayoutInfo = {
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
        .bindingCount = 1,
        .pBindings = &tBindings
    };
    VULKAN_CHECK(vkCreateDescriptorSetLayout(gtLogicalDevice, &tLayoutInfo, NULL, &gtDescriptorSetLayout));

    // allocate descriptor set from descriptor pool
    const VkDescriptorSetAllocateInfo tAllocInfo = {
        .sType              = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
        .descriptorPool     = gtDescriptorPool,
        .descriptorSetCount = 1,
        .pSetLayouts        = &gtDescriptorSetLayout
    };
    VULKAN_CHECK(vkAllocateDescriptorSets(gtLogicalDevice, &tAllocInfo, &gtDescriptorSet));

    //-----------------------------------------------------------------------------
    // write data to descriptors
    //   - you have already created the storage buffer, created the descriptor set,
    //     but you haven't "connected" them, which is what you are doing here
    //-----------------------------------------------------------------------------

    // descriptor info
    VkDescriptorBufferInfo tBufferDescriptorInfo = {
        .buffer = gtStorageBuffer,
        .range  = VK_WHOLE_SIZE
    };

    VkWriteDescriptorSet tDescriptorWrites = {
        .sType              = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .dstBinding         = 0u,
        .dstArrayElement    = 0,
        .descriptorType     = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        .descriptorCount    = 1,
        .dstSet             = gtDescriptorSet,
        .pBufferInfo        = &tBufferDescriptorInfo
    };
    vkUpdateDescriptorSets(gtLogicalDevice, 1, &tDescriptorWrites, 0, NULL);

    //-----------------------------------------------------------------------------
    // create shader & pipeline & pipeline layout
    //-----------------------------------------------------------------------------

    // create pipeline layout (resource bindings)
    VkPipelineLayoutCreateInfo tPipelineLayoutInfo = {
        .sType          = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
        .setLayoutCount = 1,
        .pSetLayouts    = &gtDescriptorSetLayout
    };
    VULKAN_CHECK(vkCreatePipelineLayout(gtLogicalDevice, &tPipelineLayoutInfo, NULL, &gtPipelineLayout));

    VkShaderModuleCreateInfo tShaderCreateInfo = {
        .sType    = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
        .codeSize = tShader.szCodeSize,
        .pCode    = (const uint32_t*)(tShader.puCode)
    };
    VULKAN_CHECK(vkCreateShaderModule(gtLogicalDevice, &tShaderCreateInfo, NULL, &gtShaderModule));

    VkPipelineShaderStageCreateInfo tShaderStage = {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
        .stage = VK_SHADER_STAGE_COMPUTE_BIT,
        .module = gtShaderModule,
        .pName = "main"
    };

    // create compute pipeline (FYI, this is when actual shader compilation is performed by the driver)
    VkComputePipelineCreateInfo tPipelineCreateInfo = {
        .sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO,
        .layout = gtPipelineLayout,
        .stage  = tShaderStage
    };
    VULKAN_CHECK(vkCreateComputePipelines(gtLogicalDevice, VK_NULL_HANDLE, 1, &tPipelineCreateInfo, NULL, &gtPipeline));

    //-----------------------------------------------------------------------------
    //-----------------------------------------------------------------------------
    //-----------------------------------------------------------------------------
    //-----------------------------------------------------------------------------
    //                                  WORK
    //-----------------------------------------------------------------------------
    //-----------------------------------------------------------------------------
    //-----------------------------------------------------------------------------
    //-----------------------------------------------------------------------------

    //-----------------------------------------------------------------------------
    // create command pool
    //-----------------------------------------------------------------------------

    // create command pool from compute queue
    const VkCommandPoolCreateInfo tCommandPoolInfo = {
        .sType            = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
        .queueFamilyIndex = gptGfx->get_vulkan_queue_family(ptAppData->ptDevice),
        .flags            = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT
    };
    VULKAN_CHECK(vkCreateCommandPool(gtLogicalDevice, &tCommandPoolInfo, NULL, &gtComputeCommandPool));

    VkCommandBufferAllocateInfo tCommandAllocInfo = {
        .sType              = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
        .level              = VK_COMMAND_BUFFER_LEVEL_PRIMARY,
        .commandPool        = gtComputeCommandPool,
        .commandBufferCount = 1,
    };

    //-----------------------------------------------------------------------------
    // allocate a command buffer from the command pool
    //-----------------------------------------------------------------------------

    VkCommandBuffer tCommandBuffer = {0};
    vkAllocateCommandBuffers(gtLogicalDevice, &tCommandAllocInfo, &tCommandBuffer);

    //-----------------------------------------------------------------------------
    // begin recording work into the command buffer
    //-----------------------------------------------------------------------------
    
    VkCommandBufferBeginInfo tCommandBufferBeginInfo = {
        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
        .flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT,
    };
    VULKAN_CHECK(vkBeginCommandBuffer(tCommandBuffer, &tCommandBufferBeginInfo));

    // bind compute pipeline
    vkCmdBindPipeline(tCommandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, gtPipeline);

    // bind resources (remember this is done in sets)
    vkCmdBindDescriptorSets(tCommandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, gtPipelineLayout, 0, 1, &gtDescriptorSet, 0, 0);

    // dispatch work
    // vkCmdDispatch(tCommandBuffer, iImageWidth, iImageHeight, 1);
    vkCmdDispatch(tCommandBuffer, iImageWidth, iImageHeight, 1);

    // finish recording
    VULKAN_CHECK(vkEndCommandBuffer(tCommandBuffer));

    // submit command buffer to compute queue
    const VkPipelineStageFlags tWaitStageMask = VK_PIPELINE_STAGE_TRANSFER_BIT;
    const VkSubmitInfo tComputeSubmitInfo = {
        .sType              = VK_STRUCTURE_TYPE_SUBMIT_INFO,
        .pWaitDstStageMask  = &tWaitStageMask,
        .commandBufferCount = 1,
        .pCommandBuffers    = &tCommandBuffer
    };

    //-----------------------------------------------------------------------------
    // create syncronization primitives
    //   - not really needed for this example but is required when submitting work
    //     I believe.
    //-----------------------------------------------------------------------------

    // fence for compute command buffer sync
    const VkFenceCreateInfo tFenceCreateInfo = {
        .sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO,
        .flags = VK_FENCE_CREATE_SIGNALED_BIT
    };
    VULKAN_CHECK(vkCreateFence(gtLogicalDevice, &tFenceCreateInfo, NULL, &gtComputeWorkFence));

    vkResetFences(gtLogicalDevice, 1, &gtComputeWorkFence);
    VULKAN_CHECK(vkQueueSubmit(gptGfx->get_vulkan_queue(ptAppData->ptDevice), 1, &tComputeSubmitInfo, gtComputeWorkFence));

    // wait for fence to be signaled
    VULKAN_CHECK(vkWaitForFences(gtLogicalDevice, 1, &gtComputeWorkFence, VK_TRUE, UINT64_MAX));

    //-----------------------------------------------------------------------------
    //-----------------------------------------------------------------------------
    //-----------------------------------------------------------------------------
    //-----------------------------------------------------------------------------
    //                                  RESULTS
    //-----------------------------------------------------------------------------
    //-----------------------------------------------------------------------------
    //-----------------------------------------------------------------------------
    //-----------------------------------------------------------------------------

    // can only output non-normalized unfortunately
    unsigned char* pucResults = malloc(iImageHeight * iImageWidth * 4);
    for(uint32_t i = 0; i < (uint32_t)(iImageHeight * iImageWidth * 4); i++)
        pucResults[i] = (unsigned char)gpfMemoryMap[i];

    plImageWriteInfo tWriteInfo = {
        .iWidth = iImageWidth,
        .iHeight = iImageHeight,
        .iComponents = 4,
        .iByteStride = iImageWidth * 4
    };
    gptImage->write("../../pl-template/out/test_image_results.png", pucResults, &tWriteInfo);

    free(pucResults);

    // unmap memory & cleanup storage buffer
    vkUnmapMemory(gtLogicalDevice, gtStorageBufferMemory);
    vkDestroyBuffer(gtLogicalDevice, gtStorageBuffer, NULL);
    vkFreeMemory(gtLogicalDevice, gtStorageBufferMemory, NULL);

    // cleanup pipeline & shader
    vkDestroyShaderModule(gtLogicalDevice, gtShaderModule, NULL);
    vkDestroyPipelineLayout(gtLogicalDevice, gtPipelineLayout, NULL);
    vkDestroyPipeline(gtLogicalDevice, gtPipeline, NULL);

    // cleanup some global resources
    vkDestroyCommandPool(gtLogicalDevice, gtComputeCommandPool, NULL);
    vkDestroyDescriptorSetLayout(gtLogicalDevice, gtDescriptorSetLayout, NULL);
    vkDestroyDescriptorPool(gtLogicalDevice, gtDescriptorPool, NULL);
    vkDestroyFence(gtLogicalDevice, gtComputeWorkFence, NULL);
}
