/*
** Copyright (c) 2015-2018, 2023 The Khronos Group Inc.
**
** Licensed under the Apache License, Version 2.0 (the "License");
** you may not use this file except in compliance with the License.
** You may obtain a copy of the License at
**
**     http://www.apache.org/licenses/LICENSE-2.0
**
** Unless required by applicable law or agreed to in writing, software
** distributed under the License is distributed on an "AS IS" BASIS,
** WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
** See the License for the specific language governing permissions and
** limitations under the License.
*/

#pragma once

#include <stdlib.h>
#include <cstring>

#include <algorithm>
#include <array>
#include <mutex>
#include <unordered_set>
#include <unordered_map>
#include <string>
#include <vector>

#include "vk_video/vulkan_video_codecs_common.h"
#include "vk_video/vulkan_video_codec_h264std.h"
#include "vk_video/vulkan_video_codec_h264std_decode.h"
#include "vk_video/vulkan_video_codec_h264std_encode.h"
#include "vk_video/vulkan_video_codec_h265std.h"
#include "vk_video/vulkan_video_codec_h265std_decode.h"
#include "vk_video/vulkan_video_codec_h265std_encode.h"
#include "vulkan/vulkan.h"

#include "vulkan/vk_icd.h"
#include "vk_typemap_helper.h"
#include "mock_icd_test_config.h"

namespace vkmock
{

using mutex_t       = std::mutex;
using lock_guard_t  = std::lock_guard<mutex_t>;
using unique_lock_t = std::unique_lock<mutex_t>;

static mutex_t        global_lock;
static uint64_t       global_unique_handle                   = 1;
static const uint32_t SUPPORTED_LOADER_ICD_INTERFACE_VERSION = 5;
static uint32_t       loader_interface_version               = 0;
static bool           negotiate_loader_icd_interface_called  = false;
static void*          CreateDispObjHandle()
{
    auto handle = new VK_LOADER_DATA;
    set_loader_magic_value(handle);
    return handle;
}
static void DestroyDispObjHandle(void* handle)
{
    delete reinterpret_cast<VK_LOADER_DATA*>(handle);
}

static constexpr uint32_t icd_physical_device_count = 1;
static std::unordered_map<VkInstance, std::array<VkPhysicalDevice, icd_physical_device_count>> physical_device_map;
static std::unordered_map<VkPhysicalDevice, std::unordered_set<VkDisplayKHR>>                  display_map;

// Map device memory handle to any mapped allocations that we'll need to free on unmap
static std::unordered_map<VkDeviceMemory, std::vector<void*>> mapped_memory_map;

// Map device memory allocation handle to the size
static std::unordered_map<VkDeviceMemory, VkDeviceSize> allocated_memory_size_map;

static std::unordered_map<VkDevice, std::unordered_map<uint32_t, std::unordered_map<uint32_t, VkQueue>>> queue_map;
static VkDeviceAddress current_available_address = 0x10000000;
struct BufferState
{
    VkDeviceSize    size;
    VkDeviceAddress address;
};
static std::unordered_map<VkDevice, std::unordered_map<VkBuffer, BufferState>> buffer_map;
static std::unordered_map<VkDevice, std::unordered_map<VkImage, VkDeviceSize>> image_memory_size_map;
static std::unordered_map<VkDevice, std::unordered_set<VkCommandPool>>         command_pool_map;
static std::unordered_map<VkCommandPool, std::vector<VkCommandBuffer>>         command_pool_buffer_map;

static constexpr uint32_t                                                     icd_swapchain_image_count = 1;
static std::unordered_map<VkSwapchainKHR, VkImage[icd_swapchain_image_count]> swapchain_image_map;

// TODO: Would like to codegen this but limits aren't in XML
static VkPhysicalDeviceLimits SetLimits(VkPhysicalDeviceLimits* limits)
{
    limits->maxImageDimension1D                             = 4096;
    limits->maxImageDimension2D                             = 4096;
    limits->maxImageDimension3D                             = 256;
    limits->maxImageDimensionCube                           = 4096;
    limits->maxImageArrayLayers                             = 256;
    limits->maxTexelBufferElements                          = 65536;
    limits->maxUniformBufferRange                           = 16384;
    limits->maxStorageBufferRange                           = 134217728;
    limits->maxPushConstantsSize                            = 128;
    limits->maxMemoryAllocationCount                        = 4096;
    limits->maxSamplerAllocationCount                       = 4000;
    limits->bufferImageGranularity                          = 1;
    limits->sparseAddressSpaceSize                          = 2147483648;
    limits->maxBoundDescriptorSets                          = 4;
    limits->maxPerStageDescriptorSamplers                   = 16;
    limits->maxPerStageDescriptorUniformBuffers             = 12;
    limits->maxPerStageDescriptorStorageBuffers             = 4;
    limits->maxPerStageDescriptorSampledImages              = 16;
    limits->maxPerStageDescriptorStorageImages              = 4;
    limits->maxPerStageDescriptorInputAttachments           = 4;
    limits->maxPerStageResources                            = 128;
    limits->maxDescriptorSetSamplers                        = 96;
    limits->maxDescriptorSetUniformBuffers                  = 72;
    limits->maxDescriptorSetUniformBuffersDynamic           = 8;
    limits->maxDescriptorSetStorageBuffers                  = 24;
    limits->maxDescriptorSetStorageBuffersDynamic           = 4;
    limits->maxDescriptorSetSampledImages                   = 96;
    limits->maxDescriptorSetStorageImages                   = 24;
    limits->maxDescriptorSetInputAttachments                = 4;
    limits->maxVertexInputAttributes                        = 16;
    limits->maxVertexInputBindings                          = 16;
    limits->maxVertexInputAttributeOffset                   = 2047;
    limits->maxVertexInputBindingStride                     = 2048;
    limits->maxVertexOutputComponents                       = 64;
    limits->maxTessellationGenerationLevel                  = 64;
    limits->maxTessellationPatchSize                        = 32;
    limits->maxTessellationControlPerVertexInputComponents  = 64;
    limits->maxTessellationControlPerVertexOutputComponents = 64;
    limits->maxTessellationControlPerPatchOutputComponents  = 120;
    limits->maxTessellationControlTotalOutputComponents     = 2048;
    limits->maxTessellationEvaluationInputComponents        = 64;
    limits->maxTessellationEvaluationOutputComponents       = 64;
    limits->maxGeometryShaderInvocations                    = 32;
    limits->maxGeometryInputComponents                      = 64;
    limits->maxGeometryOutputComponents                     = 64;
    limits->maxGeometryOutputVertices                       = 256;
    limits->maxGeometryTotalOutputComponents                = 1024;
    limits->maxFragmentInputComponents                      = 64;
    limits->maxFragmentOutputAttachments                    = 4;
    limits->maxFragmentDualSrcAttachments                   = 1;
    limits->maxFragmentCombinedOutputResources              = 4;
    limits->maxComputeSharedMemorySize                      = 16384;
    limits->maxComputeWorkGroupCount[0]                     = 65535;
    limits->maxComputeWorkGroupCount[1]                     = 65535;
    limits->maxComputeWorkGroupCount[2]                     = 65535;
    limits->maxComputeWorkGroupInvocations                  = 128;
    limits->maxComputeWorkGroupSize[0]                      = 128;
    limits->maxComputeWorkGroupSize[1]                      = 128;
    limits->maxComputeWorkGroupSize[2]                      = 64;
    limits->subPixelPrecisionBits                           = 4;
    limits->subTexelPrecisionBits                           = 4;
    limits->mipmapPrecisionBits                             = 4;
    limits->maxDrawIndexedIndexValue                        = UINT32_MAX;
    limits->maxDrawIndirectCount                            = UINT16_MAX;
    limits->maxSamplerLodBias                               = 2.0f;
    limits->maxSamplerAnisotropy                            = 16;
    limits->maxViewports                                    = 16;
    limits->maxViewportDimensions[0]                        = 4096;
    limits->maxViewportDimensions[1]                        = 4096;
    limits->viewportBoundsRange[0]                          = -8192;
    limits->viewportBoundsRange[1]                          = 8191;
    limits->viewportSubPixelBits                            = 0;
    limits->minMemoryMapAlignment                           = 64;
    limits->minTexelBufferOffsetAlignment                   = 16;
    limits->minUniformBufferOffsetAlignment                 = 16;
    limits->minStorageBufferOffsetAlignment                 = 16;
    limits->minTexelOffset                                  = -8;
    limits->maxTexelOffset                                  = 7;
    limits->minTexelGatherOffset                            = -8;
    limits->maxTexelGatherOffset                            = 7;
    limits->minInterpolationOffset                          = 0.0f;
    limits->maxInterpolationOffset                          = 0.5f;
    limits->subPixelInterpolationOffsetBits                 = 4;
    limits->maxFramebufferWidth                             = 4096;
    limits->maxFramebufferHeight                            = 4096;
    limits->maxFramebufferLayers                            = 256;
    limits->framebufferColorSampleCounts                    = 0x7F;
    limits->framebufferDepthSampleCounts                    = 0x7F;
    limits->framebufferStencilSampleCounts                  = 0x7F;
    limits->framebufferNoAttachmentsSampleCounts            = 0x7F;
    limits->maxColorAttachments                             = 4;
    limits->sampledImageColorSampleCounts                   = 0x7F;
    limits->sampledImageIntegerSampleCounts                 = 0x7F;
    limits->sampledImageDepthSampleCounts                   = 0x7F;
    limits->sampledImageStencilSampleCounts                 = 0x7F;
    limits->storageImageSampleCounts                        = 0x7F;
    limits->maxSampleMaskWords                              = 1;
    limits->timestampComputeAndGraphics                     = VK_TRUE;
    limits->timestampPeriod                                 = 1;
    limits->maxClipDistances                                = 8;
    limits->maxCullDistances                                = 8;
    limits->maxCombinedClipAndCullDistances                 = 8;
    limits->discreteQueuePriorities                         = 2;
    limits->pointSizeRange[0]                               = 1.0f;
    limits->pointSizeRange[1]                               = 64.0f;
    limits->lineWidthRange[0]                               = 1.0f;
    limits->lineWidthRange[1]                               = 8.0f;
    limits->pointSizeGranularity                            = 1.0f;
    limits->lineWidthGranularity                            = 1.0f;
    limits->strictLines                                     = VK_TRUE;
    limits->standardSampleLocations                         = VK_TRUE;
    limits->optimalBufferCopyOffsetAlignment                = 1;
    limits->optimalBufferCopyRowPitchAlignment              = 1;
    limits->nonCoherentAtomSize                             = 256;

    return *limits;
}

void SetBoolArrayTrue(VkBool32* bool_array, uint32_t num_bools)
{
    for (uint32_t i = 0; i < num_bools; ++i)
    {
        bool_array[i] = VK_TRUE;
    }
}

VkDeviceSize GetImageSizeFromCreateInfo(const VkImageCreateInfo* pCreateInfo)
{
    VkDeviceSize size = pCreateInfo->extent.width;
    size *= pCreateInfo->extent.height;
    size *= pCreateInfo->extent.depth;
    // TODO: A pixel size is 32 bytes. This accounts for the largest possible pixel size of any format. It could be
    // changed to more accurate size if need be.
    size *= 32;
    size *= pCreateInfo->arrayLayers;
    size *= (pCreateInfo->mipLevels > 1 ? 2 : 1);

    switch (pCreateInfo->format)
    {
        case VK_FORMAT_G8_B8_R8_3PLANE_420_UNORM:
        case VK_FORMAT_G8_B8_R8_3PLANE_422_UNORM:
        case VK_FORMAT_G8_B8_R8_3PLANE_444_UNORM:
        case VK_FORMAT_G10X6_B10X6_R10X6_3PLANE_420_UNORM_3PACK16:
        case VK_FORMAT_G10X6_B10X6_R10X6_3PLANE_422_UNORM_3PACK16:
        case VK_FORMAT_G10X6_B10X6_R10X6_3PLANE_444_UNORM_3PACK16:
        case VK_FORMAT_G12X4_B12X4_R12X4_3PLANE_420_UNORM_3PACK16:
        case VK_FORMAT_G12X4_B12X4_R12X4_3PLANE_422_UNORM_3PACK16:
        case VK_FORMAT_G12X4_B12X4_R12X4_3PLANE_444_UNORM_3PACK16:
        case VK_FORMAT_G16_B16_R16_3PLANE_420_UNORM:
        case VK_FORMAT_G16_B16_R16_3PLANE_422_UNORM:
        case VK_FORMAT_G16_B16_R16_3PLANE_444_UNORM:
            size *= 3;
            break;
        case VK_FORMAT_G8_B8R8_2PLANE_420_UNORM:
        case VK_FORMAT_G8_B8R8_2PLANE_422_UNORM:
        case VK_FORMAT_G10X6_B10X6R10X6_2PLANE_420_UNORM_3PACK16:
        case VK_FORMAT_G10X6_B10X6R10X6_2PLANE_422_UNORM_3PACK16:
        case VK_FORMAT_G12X4_B12X4R12X4_2PLANE_420_UNORM_3PACK16:
        case VK_FORMAT_G12X4_B12X4R12X4_2PLANE_422_UNORM_3PACK16:
        case VK_FORMAT_G16_B16R16_2PLANE_420_UNORM:
        case VK_FORMAT_G16_B16R16_2PLANE_422_UNORM:
            size *= 2;
            break;
        default:
            break;
    }

    return size;
}

typedef void(VKAPI_PTR* PFN_vkCmdExecuteCommands)(VkCommandBuffer        commandBuffer,
                                                  uint32_t               commandBufferCount,
                                                  const VkCommandBuffer* pCommandBuffers);

static TestConfig test_config;

} // namespace vkmock
