#include "MiscData.h"

#include <array>
#include <SDL3/SDL_log.h>

#include "MacroUtils.h"
#include "RenderDevice.h"

#include <glm/gtc/matrix_transform.hpp>

void MiscData::init( RenderDevice const& renderDevice )
{
  VkDevice const device = renderDevice.device;

  previousCounter = 0;

  // Pipeline Creation
  {
    size_t dataSize;
    void*  rawData = SDL_LoadFile( "Mesh.spv", &dataSize );
    ASSERT( dataSize % 4 == 0 );

    if ( !rawData )
    {
      SDL_LogError( SDL_LOG_CATEGORY_SYSTEM, "%s", SDL_GetError() );
      abort();
    }

    auto data = static_cast<uint32_t const*>(rawData);

    VkShaderModuleCreateInfo const shaderModuleCreateInfo = {
      .sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .codeSize = dataSize,
      .pCode = data,
    };

    VkShaderModule shaderModule;
    VK_CHECK( vkCreateShaderModule(device, &shaderModuleCreateInfo, nullptr, &shaderModule) );

    VkDescriptorSetLayoutBinding constexpr descriptorSetLayoutBinding = {
      .binding = 0,
      .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
      .descriptorCount = 1,
      .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
      .pImmutableSamplers = nullptr,
    };

    VkDescriptorSetLayoutCreateInfo const descriptorSetLayoutCreateInfo = {
      .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .bindingCount = 1,
      .pBindings = &descriptorSetLayoutBinding,
    };
    VK_CHECK( vkCreateDescriptorSetLayout(device, &descriptorSetLayoutCreateInfo, nullptr, &descriptorSetLayout) );

    VkPipelineLayoutCreateInfo const pipelineLayoutCreateInfo = {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .setLayoutCount = 1,
      .pSetLayouts = &descriptorSetLayout,
      .pushConstantRangeCount = 0,
      .pPushConstantRanges = nullptr,
    };
    VK_CHECK( vkCreatePipelineLayout(device, &pipelineLayoutCreateInfo, nullptr, &pipelineLayout) );

    std::array stages = {
      VkPipelineShaderStageCreateInfo{
        .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
        .pNext = nullptr,
        .flags = 0,
        .stage = VK_SHADER_STAGE_VERTEX_BIT,
        .module = shaderModule,
        .pName = "VertexMain",
        .pSpecializationInfo = nullptr,
      },
      VkPipelineShaderStageCreateInfo{
        .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
        .pNext = nullptr,
        .flags = 0,
        .stage = VK_SHADER_STAGE_FRAGMENT_BIT,
        .module = shaderModule,
        .pName = "FragmentMain",
        .pSpecializationInfo = nullptr,
      }
    };

    VkVertexInputBindingDescription constexpr bindingDescription = {
      .binding = 0,
      .stride = sizeof( Vertex ),
      .inputRate = VK_VERTEX_INPUT_RATE_VERTEX,
    };

    std::array attributeDescriptions = {
      VkVertexInputAttributeDescription{
        .location = 0,
        .binding = 0,
        .format = VK_FORMAT_R32G32B32A32_SFLOAT,
        .offset = offsetof( Vertex, position ),
      },
      VkVertexInputAttributeDescription{
        .location = 1,
        .binding = 0,
        .format = VK_FORMAT_R32G32B32A32_SFLOAT,
        .offset = offsetof( Vertex, color ),
      },
    };

    VkPipelineVertexInputStateCreateInfo const vertexInputState = {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .vertexBindingDescriptionCount = 1,
      .pVertexBindingDescriptions = &bindingDescription,
      .vertexAttributeDescriptionCount = static_cast<uint32_t>(attributeDescriptions.size()),
      .pVertexAttributeDescriptions = attributeDescriptions.data(),
    };

    VkPipelineInputAssemblyStateCreateInfo constexpr inputAssembly = {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_STRIP,
      .primitiveRestartEnable = VK_FALSE,
    };

    VkPipelineTessellationStateCreateInfo constexpr tessellationState = {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_TESSELLATION_STATE_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .patchControlPoints = 0,
    };

    VkPipelineViewportStateCreateInfo constexpr viewportState = {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .viewportCount = 1,
      .pViewports = nullptr,
      .scissorCount = 1,
      .pScissors = nullptr,
    };

    VkPipelineRasterizationStateCreateInfo constexpr rasterizationState = {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .depthClampEnable = VK_TRUE,
      .rasterizerDiscardEnable = VK_FALSE,
      .polygonMode = VK_POLYGON_MODE_FILL,
      .cullMode = VK_CULL_MODE_NONE,
      .frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE,
      .depthBiasEnable = VK_FALSE,
      .depthBiasConstantFactor = 0.0f,
      .depthBiasClamp = 0.0f,
      .depthBiasSlopeFactor = 0.0f,
      .lineWidth = 1.0f,
    };

    VkPipelineMultisampleStateCreateInfo constexpr multisampleState = {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .rasterizationSamples = VK_SAMPLE_COUNT_1_BIT,
      .sampleShadingEnable = VK_FALSE,
      .minSampleShading = 0.0f,
      .pSampleMask = nullptr,
      .alphaToCoverageEnable = VK_FALSE,
      .alphaToOneEnable = VK_FALSE,
    };

    VkPipelineDepthStencilStateCreateInfo constexpr depthStencilState = {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .depthTestEnable = VK_FALSE,
      .depthWriteEnable = VK_FALSE,
      .depthCompareOp = VK_COMPARE_OP_ALWAYS,
      .depthBoundsTestEnable = VK_FALSE,
      .stencilTestEnable = VK_FALSE,
      .front = {},
      .back = {},
      .minDepthBounds = 0.0f,
      .maxDepthBounds = 1.0f,
    };

    VkPipelineColorBlendAttachmentState constexpr colorBlendAttachmentState = {
      .blendEnable = VK_FALSE,
      .srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA,
      .dstColorBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA,
      .colorBlendOp = VK_BLEND_OP_ADD,
      .srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE,
      .dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO,
      .alphaBlendOp = VK_BLEND_OP_ADD,
      .colorWriteMask = VK_COLOR_COMPONENT_R_BIT
                        | VK_COLOR_COMPONENT_G_BIT
                        | VK_COLOR_COMPONENT_B_BIT
                        | VK_COLOR_COMPONENT_A_BIT,
    };

    VkPipelineColorBlendStateCreateInfo const colorBlendState = {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .logicOpEnable = VK_FALSE,
      .logicOp = VK_LOGIC_OP_COPY,
      .attachmentCount = 1,
      .pAttachments = &colorBlendAttachmentState,
      .blendConstants = { 0.0f, 0.0f, 0.0f, 0.0f },
    };

    std::array constexpr dynamicStates = {
      VK_DYNAMIC_STATE_VIEWPORT,
      VK_DYNAMIC_STATE_SCISSOR
    };

    VkPipelineDynamicStateCreateInfo const dynamicStateCreateInfo = {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .dynamicStateCount = static_cast<uint32_t>(dynamicStates.size()),
      .pDynamicStates = dynamicStates.data()
    };

    VkPipelineRenderingCreateInfoKHR const renderingCreateInfo = {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_RENDERING_CREATE_INFO_KHR,
      .colorAttachmentCount = 1,
      .pColorAttachmentFormats = &renderDevice.swapchainFormat,
    };

    VkGraphicsPipelineCreateInfo const graphicsPipelineCreateInfo = {
      .sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO,
      .pNext = &renderingCreateInfo,
      .flags = 0,
      .stageCount = static_cast<uint32_t>(stages.size()),
      .pStages = stages.data(),
      .pVertexInputState = &vertexInputState,
      .pInputAssemblyState = &inputAssembly,
      .pTessellationState = &tessellationState,
      .pViewportState = &viewportState,
      .pRasterizationState = &rasterizationState,
      .pMultisampleState = &multisampleState,
      .pDepthStencilState = &depthStencilState,
      .pColorBlendState = &colorBlendState,
      .pDynamicState = &dynamicStateCreateInfo,
      .layout = pipelineLayout,
      .renderPass = nullptr,
      .subpass = 0,
      .basePipelineHandle = nullptr,
      .basePipelineIndex = 0,
    };

    VK_CHECK( vkCreateGraphicsPipelines(device, nullptr, 1, &graphicsPipelineCreateInfo, nullptr, &meshPipeline) );

    vkDestroyShaderModule( device, shaderModule, nullptr );

    SDL_free( rawData );
  }

  // Vertex Buffer Creation
  {
    vertexBufferSize = sizeof vertices[0] * vertices.size();

    //  TL----TR
    //  |  \   |
    //  |   \  |
    //  |    \ |
    //  BL----BR
    //
    // BL -> BR -> TL
    // TL -> BR -> TR

    vertices = std::array{
      // Bottom Left
      Vertex{
        .position = { -1.0f, -1.0f, 0.0f, 1.0f },
        .color = { 0.0f, 0.0f, 1.0f, 1.0f },
      },
      // Bottom Right
      Vertex{
        .position = { 1.0f, -1.0f, 0.0f, 1.0f },
        .color = { 1.0f, 0.0f, 0.0f, 1.0f },
      },
      // Top Left
      Vertex{
        .position = { -1.0f, 1.0f, 0.0f, 1.0f },
        .color = { 0.0f, 1.0f, 0.0f, 1.0f },
      },
      // Top Right
      Vertex{
        .position = { 1.0f, 1.0f, 0.0f, 1.0f },
        .color = { 1.0f, 1.0f, 0.0f, 1.0f },
      }
    };

    VkBufferCreateInfo const bufferCreateInfo = {
      .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .size = vertexBufferSize,
      .usage = VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,
      .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
      .queueFamilyIndexCount = 0,
      .pQueueFamilyIndices = nullptr,
    };

    VmaAllocationCreateInfo constexpr allocationCreateInfo = {
      .flags = VMA_ALLOCATION_CREATE_MAPPED_BIT |
               VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT,
      .usage = VMA_MEMORY_USAGE_AUTO,
      .requiredFlags = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
      .preferredFlags = 0,
      .memoryTypeBits = 0,
      .pool = nullptr,
      .pUserData = nullptr,
      .priority = 1.0f,
    };

    VmaAllocationInfo allocationInfo;

    VK_CHECK(
      vmaCreateBuffer(
        renderDevice.gpuAllocator,
        &bufferCreateInfo,
        &allocationCreateInfo,
        &vertexBuffer,
        &vertexBufferAllocation,
        &allocationInfo) );

    if ( allocationInfo.pMappedData )
    {
      memcpy( allocationInfo.pMappedData, vertices.data(), vertices.size() * sizeof vertices[0] );
    }
  }

  // Camera
  {
    cameraPosition              = glm::vec3{ 0.0f, 0.0f, -5.0f };
    cameraTarget                = glm::vec3{ 0.0f, 0.0f, 0.0f };
    cameraData.modelMatrix      = glm::mat4{ 1.0f };
    cameraData.viewMatrix       = glm::lookAt( cameraPosition, cameraTarget, glm::vec3{ 0.0f, 1.0f, 0.0f } );
    cameraData.projectionMatrix = glm::perspective( glm::radians( 70.0f ), 16.0f / 9.0f, 0.1f, 1000.0f );

    cameraUniformBufferSize = sizeof( CameraData );

    VkBufferCreateInfo const bufferCreateInfo = {
      .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .size = cameraUniformBufferSize,
      .usage = VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
      .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
      .queueFamilyIndexCount = 0,
      .pQueueFamilyIndices = nullptr,
    };

    VmaAllocationCreateInfo constexpr allocationCreateInfo = {
      .flags = VMA_ALLOCATION_CREATE_MAPPED_BIT |
               VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT,
      .usage = VMA_MEMORY_USAGE_AUTO,
      .requiredFlags = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
      .preferredFlags = 0,
      .memoryTypeBits = 0,
      .pool = nullptr,
      .pUserData = nullptr,
      .priority = 1.0f,
    };

    VmaAllocationInfo allocationInfo;

    VK_CHECK(
      vmaCreateBuffer(
        renderDevice.gpuAllocator,
        &bufferCreateInfo,
        &allocationCreateInfo,
        &cameraUniformBuffer,
        &cameraUniformBufferAllocation,
        &allocationInfo) );

    if ( allocationInfo.pMappedData )
    {
      memcpy( allocationInfo.pMappedData, &cameraData, sizeof cameraData );
      cameraUniformBufferPtr = static_cast<uint8_t*>(allocationInfo.pMappedData);
    }
  }

  // Descriptors
  {
    VkDescriptorPoolSize const poolSize = {
      .type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
      .descriptorCount = renderDevice.getNumFrames(),
    };
    VkDescriptorPoolCreateInfo const descriptorPoolCreateInfo = {
      .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .maxSets = renderDevice.getNumFrames(),
      .poolSizeCount = 1,
      .pPoolSizes = &poolSize,
    };

    VK_CHECK( vkCreateDescriptorPool(device, &descriptorPoolCreateInfo, nullptr, &descriptorPool) );

    VkDescriptorSetAllocateInfo const descriptorSetAllocateInfo = {
      .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
      .pNext = nullptr,
      .descriptorPool = descriptorPool,
      .descriptorSetCount = 1,
      .pSetLayouts = &descriptorSetLayout,
    };

    VK_CHECK( vkAllocateDescriptorSets(device, &descriptorSetAllocateInfo, &descriptorSet) );

    VkDescriptorBufferInfo const descriptorBufferInfo = {
      .buffer = cameraUniformBuffer,
      .offset = 0,
      .range = sizeof CameraData,
    };

    VkWriteDescriptorSet writeDescriptorSet = {
      .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
      .pNext = nullptr,
      .dstSet = descriptorSet,
      .dstBinding = 0,
      .dstArrayElement = 0,
      .descriptorCount = 1,
      .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
      .pImageInfo = nullptr,
      .pBufferInfo = &descriptorBufferInfo,
      .pTexelBufferView = nullptr,
    };

    vkUpdateDescriptorSets( device, 1, &writeDescriptorSet, 0, nullptr );
  }

  // Barrier Creation
  {
    acquireToRenderBarrier = {
      .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER_2,
      .pNext = nullptr,
      .srcStageMask = VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT,
      .srcAccessMask = VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT,
      .dstStageMask = VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT,
      .dstAccessMask = VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT,
      .oldLayout = VK_IMAGE_LAYOUT_UNDEFINED,
      .newLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
      .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
      .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
      .subresourceRange = {
        .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
        .baseMipLevel = 0,
        .levelCount = 1,
        .baseArrayLayer = 0,
        .layerCount = 1,
      }
    };
    acquireToRenderDependency = {
      .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO,
      .pNext = nullptr,
      .dependencyFlags = 0,
      .memoryBarrierCount = 0,
      .pMemoryBarriers = nullptr,
      .bufferMemoryBarrierCount = 0,
      .pBufferMemoryBarriers = nullptr,
      .imageMemoryBarrierCount = 1,
      .pImageMemoryBarriers = &acquireToRenderBarrier,
    };

    renderToPresentBarrier = {
      .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER_2,
      .pNext = nullptr,
      .srcStageMask = VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT,
      .srcAccessMask = VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT,
      .dstStageMask = VK_PIPELINE_STAGE_2_BOTTOM_OF_PIPE_BIT,
      .dstAccessMask = VK_ACCESS_2_NONE,
      .oldLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
      .newLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
      .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
      .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
      .subresourceRange = {
        .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
        .baseMipLevel = 0,
        .levelCount = 1,
        .baseArrayLayer = 0,
        .layerCount = 1,
      }
    };
    renderToPresentDependency = {
      .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO,
      .pNext = nullptr,
      .dependencyFlags = 0,
      .memoryBarrierCount = 0,
      .pMemoryBarriers = nullptr,
      .bufferMemoryBarrierCount = 0,
      .pBufferMemoryBarriers = nullptr,
      .imageMemoryBarrierCount = 1,
      .pImageMemoryBarriers = &renderToPresentBarrier,
    };
  }
}

void MiscData::destroy( RenderDevice const& renderDevice )
{
  VkDevice const device = renderDevice.device;

  vkDestroyDescriptorPool( device, Take( descriptorPool ), nullptr );
  vmaDestroyBuffer( renderDevice.gpuAllocator, Take( cameraUniformBuffer ), Take( cameraUniformBufferAllocation ) );
  vmaDestroyBuffer( renderDevice.gpuAllocator, Take( vertexBuffer ), Take( vertexBufferAllocation ) );

  vkDestroyPipeline( device, Take( meshPipeline ), nullptr );
  vkDestroyPipelineLayout( device, Take( pipelineLayout ), nullptr );
  vkDestroyDescriptorSetLayout( device, Take( descriptorSetLayout ), nullptr );
}