// Copyright (c) 2018-2026 Jean-Philippe Bruyère <jp_bruyere@hotmail.com>
//
// This code is licensed under the MIT license (MIT) (http://opensource.org/licenses/MIT)

#define GetInstProcAddress(inst, func) (PFN_##func) vkGetInstanceProcAddr(inst, #func);

#define GetVkProcAddress(dev, inst, func)                                                                              \
    (vkGetDeviceProcAddr(dev, #func) == NULL) ? (PFN_##func)vkGetInstanceProcAddr(inst, #func)                         \
                                              : (PFN_##func)vkGetDeviceProcAddr(dev, #func)

#include "vkvg_device_internal.h"
#include "vkvg_context_internal.h"

uint32_t vkvg_log_level = VKVG_LOG_DEBUG;
#ifdef VKVG_WIRED_DEBUG
vkvg_wired_debug_mode vkvg_wired_debug = vkvg_wired_debug_mode_normal;
#endif

PFN_vkCmdBindPipeline       CmdBindPipeline;
PFN_vkCmdBindDescriptorSets CmdBindDescriptorSets;
PFN_vkCmdBindIndexBuffer    CmdBindIndexBuffer;
PFN_vkCmdBindVertexBuffers  CmdBindVertexBuffers;

PFN_vkCmdDrawIndexed        CmdDrawIndexed;
PFN_vkCmdDraw               CmdDraw;

PFN_vkCmdSetStencilCompareMask CmdSetStencilCompareMask;
PFN_vkCmdSetStencilReference   CmdSetStencilReference;
PFN_vkCmdSetStencilWriteMask   CmdSetStencilWriteMask;
PFN_vkCmdBeginRenderPass       CmdBeginRenderPass;
PFN_vkCmdEndRenderPass         CmdEndRenderPass;
PFN_vkCmdSetViewport           CmdSetViewport;
PFN_vkCmdSetScissor            CmdSetScissor;

PFN_vkCmdPushConstants   CmdPushConstants;

PFN_vkWaitForFences      WaitForFences;
PFN_vkResetFences        ResetFences;
PFN_vkResetCommandBuffer ResetCommandBuffer;

bool _device_try_get_phyinfo(VkhPhyInfo* phys, uint32_t phyCount, VkPhysicalDeviceType gpuType, VkhPhyInfo* phy) {
    for (uint32_t i = 0; i < phyCount; i++) {
        if (vkh_phyinfo_get_properties(phys[i]).deviceType == gpuType) {
            *phy = phys[i];
            return true;
        }
    }
    return false;
}
// TODO:save/reload cache in user temp directory
void _device_create_pipeline_cache(VkvgDevice dev) {

    VkPipelineCacheCreateInfo pipelineCacheCreateInfo = {.sType = VK_STRUCTURE_TYPE_PIPELINE_CACHE_CREATE_INFO};
    VK_CHECK_RESULT(vkCreatePipelineCache(dev->vkDev, &pipelineCacheCreateInfo, NULL, &dev->pipelineCache));
}

void _device_createDescriptorSetLayout(VkvgDevice dev) {
    VkSamplerCreateInfo samplerInfo = {.sType         = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
                                       .maxAnisotropy = 1.0,
                                       .addressModeU  = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
                                       .addressModeV  = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
                                       .addressModeW  = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
                                       .magFilter     = VK_FILTER_NEAREST,
                                       .minFilter     = VK_FILTER_NEAREST,
                                       .mipmapMode    = VK_SAMPLER_MIPMAP_MODE_NEAREST,
                                       .unnormalizedCoordinates = VK_TRUE
    };
    vkCreateSampler(dev->vkDev, &samplerInfo, NULL, &dev->smpFont);

    VkDescriptorBindingFlags bindingFlags[] = {
        VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT,  // Allows empty/unbound slots in this binding
        VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT
    };
    VkDescriptorSetLayoutBindingFlagsCreateInfo bindingFlagsInfo = {
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_BINDING_FLAGS_CREATE_INFO,
        .bindingCount = 2,
        .pBindingFlags = bindingFlags
    };
    VkDescriptorSetLayoutBinding    dsLayoutBinding[]    = {
        {0, VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, 1, VK_SHADER_STAGE_FRAGMENT_BIT, NULL},
        {1, VK_DESCRIPTOR_TYPE_SAMPLER, 1, VK_SHADER_STAGE_FRAGMENT_BIT, NULL}
    };
    VkDescriptorSetLayoutCreateInfo dsLayoutCreateInfo = {
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
        .pNext = &bindingFlagsInfo,
        .bindingCount = 2,
        .pBindings = dsLayoutBinding
    };
    VK_CHECK_RESULT(vkCreateDescriptorSetLayout(dev->vkDev, &dsLayoutCreateInfo, NULL, &dev->dslSrc));

    bindingFlagsInfo.bindingCount = 1;
    dsLayoutCreateInfo.bindingCount = 1;
    dsLayoutBinding[0].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    dsLayoutBinding[0].pImmutableSamplers = &dev->smpFont;
    VK_CHECK_RESULT(vkCreateDescriptorSetLayout(dev->vkDev, &dsLayoutCreateInfo, NULL, &dev->dslFont));

    dsLayoutBinding[0].pImmutableSamplers = NULL;
    dsLayoutBinding[0].descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC;
    VK_CHECK_RESULT(vkCreateDescriptorSetLayout(dev->vkDev, &dsLayoutCreateInfo, NULL, &dev->dslGrad));

    VkPushConstantRange pushConstantRange[] = {
        {VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(push_constants)},
        //{VK_SHADER_STAGE_FRAGMENT_BIT,0,sizeof(push_constants)}
    };
    VkDescriptorSetLayout dsls[] = {dev->dslFont, dev->dslSrc, dev->dslGrad};
    VkPipelineLayoutCreateInfo pipelineLayoutCreateInfo = {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
        .pushConstantRangeCount = 1,
        .pPushConstantRanges = (VkPushConstantRange*)&pushConstantRange,
        .setLayoutCount = 3,
        .pSetLayouts    = dsls};
    VK_CHECK_RESULT(vkCreatePipelineLayout(dev->vkDev, &pipelineLayoutCreateInfo, NULL, &dev->pipelineLayout));
}
void _device_allocalte_dummy_descriptors (VkvgDevice dev) {
    const VkDescriptorPoolSize descriptorPoolSize[] = {
                {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1},
                {VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1}
    };
    VkDescriptorPoolCreateInfo descriptorPoolCreateInfo = {.sType   = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
                                                           .maxSets = 2,
                                                           .flags   = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT,
                                                           .poolSizeCount = 2,
                                                           .pPoolSizes    = descriptorPoolSize};
    VK_CHECK_RESULT(vkCreateDescriptorPool(dev->vkDev, &descriptorPoolCreateInfo, NULL, &dev->descriptorPool));

    VkDescriptorSetAllocateInfo descriptorSetAllocateInfo = {.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
                                                             .descriptorPool     = dev->descriptorPool,
                                                             .descriptorSetCount = 1,
                                                             .pSetLayouts        = &dev->dslSrc};
    VK_CHECK_RESULT(vkAllocateDescriptorSets(dev->vkDev, &descriptorSetAllocateInfo, &dev->dummySrcSet));
    descriptorSetAllocateInfo.pSetLayouts = &dev->dslGrad;
    VK_CHECK_RESULT(vkAllocateDescriptorSets(dev->vkDev, &descriptorSetAllocateInfo, &dev->dummyGradSet));
}
void _device_wait_idle(VkvgDevice dev) { vkDeviceWaitIdle(dev->vkDev); }
void _device_wait_and_reset_device_fence(VkvgDevice dev) {
    vkWaitForFences(dev->vkDev, 1, &dev->fence, VK_TRUE, UINT64_MAX);
    ResetFences(dev->vkDev, 1, &dev->fence);
}

bool _device_try_get_cached_context(VkvgDevice dev, VkvgContext* pCtx) {
    LOCK_DEVICE

    if (dev->cachedContextCount) {
        thrd_t       curThread = thrd_current();
        _cached_ctx* prev      = NULL;
        _cached_ctx* cur       = dev->cachedContextLast;
        while (cur) {
            if (thrd_equal(cur->thread, curThread)) {
                if (prev)
                    prev->pNext = cur->pNext;
                else
                    dev->cachedContextLast = cur->pNext;

                dev->cachedContextCount--;

                LOG(VKVG_LOG_THREAD, "get cached context: %p, thd:%lu cached ctx: %d\n", cur->ctx, cur->thread,
                    dev->cachedContextCount);

                *pCtx = cur->ctx;
                free(cur);
                UNLOCK_DEVICE
                return true;
            }
            prev = cur;
            cur  = cur->pNext;
        }
    }
    *pCtx = NULL;
    UNLOCK_DEVICE
    return false;
}
void _device_store_context(VkvgContext ctx) {
    VkvgDevice dev = ctx->dev;

    LOCK_DEVICE

    _cached_ctx* cur = (_cached_ctx*)calloc(1, sizeof(_cached_ctx));
    cur->ctx         = ctx;
    cur->thread      = thrd_current();
    cur->pNext       = dev->cachedContextLast;

    dev->cachedContextLast = cur;
    dev->cachedContextCount++;

    LOG(VKVG_LOG_THREAD, "store context: %p, thd:%lu cached ctx: %d\n", cur->ctx, cur->thread, dev->cachedContextCount);

    ctx->references++;

    UNLOCK_DEVICE
}
void _device_submit_cmd(VkvgDevice dev, VkCommandBuffer* cmd, VkFence fence) {
    LOCK_DEVICE
    vkh_cmd_submit(dev->gQueue, cmd, fence);
    UNLOCK_DEVICE
}

bool _device_init_function_pointers(VkvgDevice dev) {
#if defined(DEBUG) && defined(VKVG_DBG_UTILS)
    if (vkGetInstanceProcAddr(dev->instance, "vkSetDebugUtilsObjectNameEXT") == VK_NULL_HANDLE) {
        LOGE("vkvg create device failed: 'VK_EXT_debug_utils' has to be loaded for Debug build\n");
        return false;
    }
    vkh_device_init_debug_utils(dev->vkhDev);
#endif
    CmdBindPipeline          = GetVkProcAddress(dev->vkDev, dev->instance, vkCmdBindPipeline);
    CmdBindDescriptorSets    = GetVkProcAddress(dev->vkDev, dev->instance, vkCmdBindDescriptorSets);
    CmdBindIndexBuffer       = GetVkProcAddress(dev->vkDev, dev->instance, vkCmdBindIndexBuffer);
    CmdBindVertexBuffers     = GetVkProcAddress(dev->vkDev, dev->instance, vkCmdBindVertexBuffers);
    CmdDrawIndexed           = GetVkProcAddress(dev->vkDev, dev->instance, vkCmdDrawIndexed);
    CmdDraw                  = GetVkProcAddress(dev->vkDev, dev->instance, vkCmdDraw);
    CmdSetStencilReference   = GetVkProcAddress(dev->vkDev, dev->instance, vkCmdSetStencilReference);
    CmdBeginRenderPass       = GetVkProcAddress(dev->vkDev, dev->instance, vkCmdBeginRenderPass);
    CmdEndRenderPass         = GetVkProcAddress(dev->vkDev, dev->instance, vkCmdEndRenderPass);
    CmdSetViewport           = GetVkProcAddress(dev->vkDev, dev->instance, vkCmdSetViewport);
    CmdSetScissor            = GetVkProcAddress(dev->vkDev, dev->instance, vkCmdSetScissor);
    CmdPushConstants         = GetVkProcAddress(dev->vkDev, dev->instance, vkCmdPushConstants);
    WaitForFences            = GetVkProcAddress(dev->vkDev, dev->instance, vkWaitForFences);
    ResetFences              = GetVkProcAddress(dev->vkDev, dev->instance, vkResetFences);
    ResetCommandBuffer       = GetVkProcAddress(dev->vkDev, dev->instance, vkResetCommandBuffer);
    return true;
}

void _device_check_best_image_tiling(VkvgDevice dev, VkFormat format) {
    VkFlags            stencilFormats[] = {VK_FORMAT_S8_UINT, VK_FORMAT_D16_UNORM_S8_UINT, VK_FORMAT_D24_UNORM_S8_UINT,
                                           VK_FORMAT_D32_SFLOAT_S8_UINT};
    VkFormatProperties phyStencilProps = {0}, phyImgProps = {0};

    // check png blit format
    VkFlags pngBlitFormats[] = {VK_FORMAT_R8G8B8A8_UNORM, VK_FORMAT_R8G8B8A8_SRGB};
    dev->pngStagFormat       = VK_FORMAT_UNDEFINED;
    for (int i = 0; i < 2; i++) {
        vkGetPhysicalDeviceFormatProperties(dev->phy, pngBlitFormats[i], &phyImgProps);
        if ((phyImgProps.linearTilingFeatures & VKVG_PNG_WRITE_IMG_REQUIREMENTS) == VKVG_PNG_WRITE_IMG_REQUIREMENTS) {
            dev->pngStagFormat = pngBlitFormats[i];
            dev->pngStagTiling = VK_IMAGE_TILING_LINEAR;
            break;
        } else if ((phyImgProps.optimalTilingFeatures & VKVG_PNG_WRITE_IMG_REQUIREMENTS) ==
                   VKVG_PNG_WRITE_IMG_REQUIREMENTS) {
            dev->pngStagFormat = pngBlitFormats[i];
            dev->pngStagTiling = VK_IMAGE_TILING_OPTIMAL;
            break;
        }
    }

    if (dev->pngStagFormat == VK_FORMAT_UNDEFINED)
        LOGW("vkvg create device warning: no suitable image format for png write\n");

    dev->stencilFormat     = VK_FORMAT_UNDEFINED;
    dev->stencilAspectFlag = VK_IMAGE_ASPECT_STENCIL_BIT;
    dev->supportedTiling   = 0xff;

    vkGetPhysicalDeviceFormatProperties(dev->phy, format, &phyImgProps);

    if ((phyImgProps.optimalTilingFeatures & VKVG_SURFACE_IMGS_REQUIREMENTS) == VKVG_SURFACE_IMGS_REQUIREMENTS) {
        for (int i = 0; i < 4; i++) {
            vkGetPhysicalDeviceFormatProperties(dev->phy, stencilFormats[i], &phyStencilProps);
            if (phyStencilProps.optimalTilingFeatures & VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT) {
                dev->stencilFormat = stencilFormats[i];
                if (i > 0)
                    dev->stencilAspectFlag |= VK_IMAGE_ASPECT_DEPTH_BIT;
                dev->supportedTiling = VK_IMAGE_TILING_OPTIMAL;
                return;
            }
        }
    }
    if ((phyImgProps.linearTilingFeatures & VKVG_SURFACE_IMGS_REQUIREMENTS) == VKVG_SURFACE_IMGS_REQUIREMENTS) {
        for (int i = 0; i < 4; i++) {
            vkGetPhysicalDeviceFormatProperties(dev->phy, stencilFormats[i], &phyStencilProps);
            if (phyStencilProps.linearTilingFeatures & VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT) {
                dev->stencilFormat = stencilFormats[i];
                if (i > 0)
                    dev->stencilAspectFlag |= VK_IMAGE_ASPECT_DEPTH_BIT;
                dev->supportedTiling = VK_IMAGE_TILING_LINEAR;
                return;
            }
        }
    }
    dev->status = VKVG_STATUS_INVALID_FORMAT;
    LOGE("vkvg create device failed: image format not supported: %d\n", format);
}

void _dump_image_format_properties(VkvgDevice dev, VkFormat format) {
    /*VkImageFormatProperties imgProps;
    VK_CHECK_RESULT(vkGetPhysicalDeviceImageFormatProperties(dev->phy,
                                                             format, VK_IMAGE_TYPE_2D, VKVG_TILING,
                                                             VK_IMAGE_USAGE_SAMPLED_BIT|VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT|VK_IMAGE_USAGE_TRANSFER_SRC_BIT|VK_IMAGE_USAGE_TRANSFER_DST_BIT,
                                                             0, &imgProps));
    printf ("tiling			  = %d\n", VKVG_TILING);
    printf ("max extend		  = (%d, %d, %d)\n", imgProps.maxExtent.width, imgProps.maxExtent.height,
    imgProps.maxExtent.depth); printf ("max mip levels	  = %d\n", imgProps.maxMipLevels); printf ("max array layers =
    %d\n", imgProps.maxArrayLayers); printf ("sample counts	  = "); if (imgProps.sampleCounts & VK_SAMPLE_COUNT_1_BIT)
        printf ("1,");
    if (imgProps.sampleCounts & VK_SAMPLE_COUNT_2_BIT)
        printf ("2,");
    if (imgProps.sampleCounts & VK_SAMPLE_COUNT_4_BIT)
        printf ("4,");
    if (imgProps.sampleCounts & VK_SAMPLE_COUNT_8_BIT)
        printf ("8,");
    if (imgProps.sampleCounts & VK_SAMPLE_COUNT_16_BIT)
        printf ("16,");
    if (imgProps.sampleCounts & VK_SAMPLE_COUNT_32_BIT)
        printf ("32,");
    printf ("\n");
    printf ("max resource size= %lu\n", imgProps.maxResourceSize);
*/
}