#ifndef COMPUTEANDGRAPHICS_H
#define COMPUTEANDGRAPHICS_H

#pragma once

#ifdef NDEBUG
    const bool enableValidationLayers = false;
#else
    const bool enableValidationLayers = true;
#endif

#include "Types.h"

#include <chrono>


#include <optional>
#include <cstdint> // Necessary for uint32_t
#include <limits> // Necessary for std::numeric_limits
#include <algorithm> // Necessary for std::clamp
#include <fstream>
#include <array>

#ifdef Q_OS_ANDROID
    const QString TARGET_PLATFORM = "Android";
#else
    const QString TARGET_PLATFORM = "Linux";
#endif

const int MAX_FRAMES_IN_FLIGHT = 2;

const uint32_t PARTICLE_COUNT = 8192;

struct Particle {
    alignas(16) glm::vec2 position;
    alignas(16) glm::vec2 velocity;
    alignas(16) glm::vec4 color;

    static VkVertexInputBindingDescription getBindingDescription() {
        VkVertexInputBindingDescription bindingDescription{};
        bindingDescription.binding = 0;
        bindingDescription.stride = sizeof(Particle);
        bindingDescription.inputRate = VK_VERTEX_INPUT_RATE_VERTEX;

        return bindingDescription;
    }

    static std::array<VkVertexInputAttributeDescription, 2> getAttributeDescriptions() {
        std::array<VkVertexInputAttributeDescription, 2> attributeDescriptions{};

        attributeDescriptions[0].binding = 0;
        attributeDescriptions[0].location = 0;
        attributeDescriptions[0].format = VK_FORMAT_R32G32_SFLOAT;
        attributeDescriptions[0].offset = offsetof(Particle, position);

        attributeDescriptions[1].binding = 0;
        attributeDescriptions[1].location = 1;
        attributeDescriptions[1].format = VK_FORMAT_R32G32B32A32_SFLOAT;
        attributeDescriptions[1].offset = offsetof(Particle, color);

        return attributeDescriptions;
    }
};

struct QueueFamilyIndices {
    std::optional<uint32_t> graphicsFamily;
    std::optional<uint32_t> presentFamily;

    bool isComplete() {
        return graphicsFamily.has_value() && presentFamily.has_value();
    }
};

struct Vertex {
    alignas(16) glm::vec3 pos;
    alignas(16) glm::vec3 color;
    alignas(16) glm::vec2 texCoord;

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

        return bindingDescription;
    }

    static std::array<VkVertexInputAttributeDescription, 3> getAttributeDescriptions() {
        std::array<VkVertexInputAttributeDescription, 3> attributeDescriptions{};

        attributeDescriptions[0].binding = 0;
        attributeDescriptions[0].location = 0;
        attributeDescriptions[0].format = VK_FORMAT_R32G32B32_SFLOAT;
        attributeDescriptions[0].offset = offsetof(Vertex, pos);

        attributeDescriptions[1].binding = 0;
        attributeDescriptions[1].location = 1;
        attributeDescriptions[1].format = VK_FORMAT_R32G32B32_SFLOAT;
        attributeDescriptions[1].offset = offsetof(Vertex, color);

        attributeDescriptions[2].binding = 0;
        attributeDescriptions[2].location = 2;
        attributeDescriptions[2].format = VK_FORMAT_R32G32_SFLOAT;
        attributeDescriptions[2].offset = offsetof(Vertex, texCoord);

        return attributeDescriptions;
    }

    bool operator==(const Vertex& other) const {
        return pos == other.pos && color == other.color && texCoord == other.texCoord;
    }
};

struct UniformBufferObject {
    alignas(16) glm::mat4 model;
    alignas(16) glm::mat4 view;
    alignas(16) glm::mat4 proj;
};

class ComputeAndGraphics
{
public:
    ComputeAndGraphics();
    void run();

    // Message Callback
    static VKAPI_ATTR VkBool32 VKAPI_CALL debugCallback(VkDebugUtilsMessageSeverityFlagBitsEXT /*messageSeverity*/,
        VkDebugUtilsMessageTypeFlagsEXT /*messageType*/,const VkDebugUtilsMessengerCallbackDataEXT* pCallbackData,void* /*pUserData*/) {

        std::cerr << "validation layer: " << pCallbackData->pMessage << std::endl;

        return VK_FALSE;
    }
    static std::vector<char> readFile(const std::string& filename);

    static void framebufferResizeCallback(GLFWwindow* window, int width, int height);
//    static void framebufferResizeCallback(QQuickWindow* window, int width, int height);

private:
    void initWindow();
    void initVulkan();
    void mainLoop();
    void cleanup();
    void createInstance();

    // Enable Validation Layers
    bool checkValidationLayerSupport();

    // Message Callback
    std::vector<const char*> getRequiredExtensions();
    void setupDebugMessenger();

    void pickPhysicalDevice();
    void chooseDevice(uint32_t deviceCount);
    bool isDeviceSuitable(VkPhysicalDevice device);

    void createLogicalDevice();
    QueueFamilyIndices findQueueFamilies(VkPhysicalDevice device);

    void createSurface();

    bool checkDeviceExtensionSupport(VkPhysicalDevice device);
    SwapChainSupportDetails querySwapChainSupport(VkPhysicalDevice device);
    VkSurfaceFormatKHR chooseSwapSurfaceFormat(const std::vector<VkSurfaceFormatKHR>& availableFormats);
    VkPresentModeKHR chooseSwapPresentMode(const std::vector<VkPresentModeKHR>& availablePresentModes);
    VkExtent2D chooseSwapExtent(const VkSurfaceCapabilitiesKHR& capabilities);
    void createSwapChain();

    void createImageViews();

    void createGraphicsPipeline();
    VkShaderModule createShaderModule(const std::vector<char>& code);
    void createRenderPass();

    void createFramebuffers();
    void createCommandPool();
    void createCommandBuffer();
    void recordCommandBuffer(VkCommandBuffer commandBuffer, uint32_t imageIndex);

    void drawFrame();
    void createSyncObjects();

    void recreateSwapChain();
    void cleanupSwapChain();

    void createVertexBuffer();
    uint32_t findMemoryType(uint32_t typeFilter, VkMemoryPropertyFlags properties);
    void createBuffer(VkDeviceSize size, VkBufferUsageFlags usage, VkMemoryPropertyFlags properties, VkBuffer& buffer, VkDeviceMemory& bufferMemory);
    void copyBuffer(VkBuffer srcBuffer, VkBuffer dstBuffer, VkDeviceSize size);
    void createIndexBuffer();

    void createDescriptorSetLayout();
    void createUniformBuffers();
    void updateUniformBuffer(uint32_t currentImage);
    void createDescriptorPool();
    void createDescriptorSets();

    void createTextureImage();
    void createImage(uint32_t width, uint32_t height, uint32_t mipLevels, VkSampleCountFlagBits numSamples, VkFormat format, VkImageTiling tiling, VkImageUsageFlags usage
                     , VkMemoryPropertyFlags properties, VkImage& image, VkDeviceMemory& imageMemory);
    VkCommandBuffer beginSingleTimeCommands();
    void endSingleTimeCommands(VkCommandBuffer commandBuffer);
    void transitionImageLayout(VkImage image, VkFormat format, VkImageLayout oldLayout, VkImageLayout newLayout, uint32_t mipLevels);
    void copyBufferToImage(VkBuffer buffer, VkImage image, uint32_t width, uint32_t height);
    void createTextureImageView();
    VkImageView createImageView(VkImage image, VkFormat format, VkImageAspectFlags aspectFlags, uint32_t mipLevels);
    void createTextureSampler();

    void createDepthResources();
    VkFormat findSupportedFormat(const std::vector<VkFormat>& candidates, VkImageTiling tiling, VkFormatFeatureFlags features);
    VkFormat findDepthFormat();
    bool hasStencilComponent(VkFormat format);

    void createShaderStorageBuffers();
    void createComputePipeline();
    void createComputeCommandBuffers();
    void recordComputeCommandBuffer(VkCommandBuffer commandBuffer);
    void createComputeDescriptorSets();
    void createComputeDescriptorSetLayout();
    void createComputeDescriptorPool();

    VkApplicationInfo createInstanceAppInfo();
    VkDebugUtilsMessengerCreateInfoEXT createInstanceDebugMessengerInfo();
    VkInstanceCreateInfo createInstanceInfo(VkApplicationInfo* appInfo, std::vector<const char*>& extensions
                                            , VkDebugUtilsMessengerCreateInfoEXT& debugCreateInfo);

    VkDeviceQueueCreateInfo createQueueInfo(uint32_t queueFamily, uint32_t queueCount, float* queuePriority);
    VkPhysicalDeviceFeatures createDeviceFeatures();
    VkDeviceCreateInfo createLogicalDeviceInfo(std::vector<VkDeviceQueueCreateInfo>& queueCreateInfos, VkPhysicalDeviceFeatures* deviceFeatures);

    VkSwapchainCreateInfoKHR createSwapChainInfo(SwapChainSupportDetails swapChainSupport, VkSurfaceKHR _surface, uint32_t imageCount
                                                 , VkSurfaceFormatKHR surfaceFormat, VkExtent2D extent, VkPresentModeKHR presentMode);

    VkAttachmentDescription createColorAttachmentInfo(VkFormat format, VkSampleCountFlagBits msaaSamples);
    VkAttachmentDescription createDepthAttachmentInfo(VkFormat format, VkSampleCountFlagBits msaaSamples);
    VkAttachmentReference createAttachmentRefInfo(uint32_t offset, VkImageLayout layout);
    VkSubpassDescription createSubpassInfo(VkAttachmentReference* colorAttachRef, VkAttachmentReference* depthAttachRef);
    VkRenderPassCreateInfo createRenderPassInfo(std::vector<VkAttachmentDescription>& attachments, VkSubpassDescription* subpass, VkSubpassDependency* dependency);
    VkSubpassDependency createSubpassDependencyInfo();

    VkDescriptorSetLayoutBinding createLayoutBindingInfo(uint32_t offset, VkDescriptorType descType, VkShaderStageFlagBits StageFlags);

    VkPipelineShaderStageCreateInfo createPipelineShaderInfo(VkShaderStageFlagBits stage, VkShaderModule module);
    VkPipelineVertexInputStateCreateInfo createPipelineVertexInputInfo();
    VkPipelineInputAssemblyStateCreateInfo createPipelineInputAssemblyStateInfo();
    VkPipelineDynamicStateCreateInfo createPipelineDynamicStateInfo(std::vector<VkDynamicState>& dynamicStates);
    VkPipelineViewportStateCreateInfo createPipelineViewportStateInfo();
    VkPipelineRasterizationStateCreateInfo createPipelineRasterizationStateInfo();
    VkPipelineMultisampleStateCreateInfo createPipelineMultisampleStateInfo();
    VkPipelineColorBlendAttachmentState colorPipelineBlendAttachmentStateInfo();
    VkPipelineColorBlendStateCreateInfo colorPipelineBlendStateInfo(VkPipelineColorBlendAttachmentState* colorBlendAttachment);
    VkPipelineLayoutCreateInfo createPipelineLayoutInfo();
    VkPipelineDepthStencilStateCreateInfo createPipelineDepthStencilStateInfo();
    VkPipelineLayoutCreateInfo createComputePipelineLayoutInfo();

    VkFramebufferCreateInfo createFramebufferInfo(std::vector<VkImageView>& attachments);

    void copyImageToStagingBuffer(VkBuffer& stagingBuffer, VkDeviceMemory& stagingBufferMemory, stbi_uc* pixels, VkDeviceSize imageSize);

    VkImageMemoryBarrier initTransitionLayoutBarrierInfo(VkImage image, VkFormat format, VkImageLayout oldLayout, VkImageLayout newLayout, uint32_t mipLevels);
    void setSrcAndDst(VkImageMemoryBarrier& barrier, VkPipelineStageFlags& sourceStage, VkPipelineStageFlags& destinationStage, VkImageLayout oldLayout
                      , VkImageLayout newLayout);

    void allocateAndBindImageMemory(VkImage& image, VkDeviceMemory& imageMemory, VkMemoryPropertyFlags properties);

    VkBufferImageCopy createBufferImageCopyInfo(uint32_t width, uint32_t height);

    void copyVerticesToStagingBuffer(VkBuffer& stagingBuffer, VkDeviceMemory& stagingBufferMemory);
    void copyIndicesToStagingBuffer(VkBuffer& stagingBuffer, VkDeviceMemory& stagingBufferMemory);

    void allocateAndBindBufferMemory(VkBuffer& buffer, VkDeviceMemory& bufferMemory, VkMemoryPropertyFlags properties);
    void copyParticlesToStagingBuffer(VkBuffer& stagingBuffer, VkDeviceMemory& stagingBufferMemory);

    std::vector<Particle> initParticles();
    void initShaderStorageBuffers(VkBuffer& stagingBuffer);

    void allocateDescriptorSets();
    VkDescriptorBufferInfo createDescriptorBufferInfo(VkBuffer buffer, VkDeviceSize range);
    VkDescriptorImageInfo createDescriptorImageInfo();
    void addUniformBufferWriteDescriptor(VkWriteDescriptorSet& descriptorWrite, VkDescriptorBufferInfo* bufferInfo, size_t index, uint32_t dstBinding);
    void addImageWriteDescriptor(VkWriteDescriptorSet& descriptorWrite, VkDescriptorImageInfo* imageInfo, size_t index, uint32_t dstBinding);

    void allocateComputeDescriptorSets();
    void addStorageBufferWriteDescriptor(VkWriteDescriptorSet& descriptorWrite, VkDescriptorBufferInfo* bufferInfo, size_t index, uint32_t dstBinding);

    void computeSubmission();
    VkSubmitInfo createComputeSubmitInfo();
    VkResult graphicsSubmission();
    void drawSubmission(uint32_t imageIndex, VkSemaphore* waitSemaphores, VkPipelineStageFlags* waitStages, VkSemaphore* signalSemaphores);
    VkSubmitInfo createDrawSubmitInfo(VkSemaphore* waitSemaphores, VkPipelineStageFlags* waitStages, VkSemaphore* signalSemaphores);
    VkResult presentSubmission(uint32_t imageIndex, VkSemaphore* signalSemaphores);

    VkViewport createViewportInfo();
    VkRect2D createScissorInfo();
    void beginPipelineCommands(VkCommandBuffer commandBuffer);
    void beginRenderPass(VkCommandBuffer commandBuffer, uint32_t imageIndex);

    void cleanPipeline();
    void cleanBuffers();
    void cleanSyncObjects();
    void cleanBase();

    const int32_t WIDTH = 800;
    const int32_t HEIGHT = 600;

    const std::vector<const char*> validationLayers = {
        "VK_LAYER_KHRONOS_validation"
    };
    const std::vector<const char*> deviceExtensions = {
        VK_KHR_SWAPCHAIN_EXTENSION_NAME
    };
    const std::vector<Vertex> _vertices = {
        {{-0.5f, -0.5f, 0.0f}, {1.0f, 0.0f, 0.0f}, {1.0f, 0.0f}},
        {{0.5f, -0.5f, 0.0f}, {0.0f, 1.0f, 0.0f}, {0.0f, 0.0f}},
        {{0.5f, 0.5f, 0.0f}, {0.0f, 0.0f, 1.0f}, {0.0f, 1.0f}},
        {{-0.5f, 0.5f, 0.0f}, {1.0f, 1.0f, 1.0f}, {1.0f, 1.0f}},

//        {{-0.5f, -0.5f, -0.5f}, {1.0f, 0.0f, 0.0f}, {0.0f, 0.0f}},
//        {{0.5f, -0.5f, -0.5f}, {0.0f, 1.0f, 0.0f}, {1.0f, 0.0f}},
//        {{0.5f, 0.5f, -0.5f}, {0.0f, 0.0f, 1.0f}, {1.0f, 1.0f}},
//        {{-0.5f, 0.5f, -0.5f}, {1.0f, 1.0f, 1.0f}, {0.0f, 1.0f}}
    };
    const std::vector<uint32_t> _indices = {
        0, 1, 2, 2, 3, 0,
//        4, 5, 6, 6, 7, 4
    };

    uint32_t currentFrame = 0;
    bool framebufferResized = false;

    GLFWwindow *_window;
//    QQuickWindow* _window;
    VkInstance _instance;
    VkDebugUtilsMessengerEXT _debugMessenger;
    VkPhysicalDevice _physicalDevice;
    VkDevice _device;
    VkQueue _graphicsQueue;
    VkSurfaceKHR _surface;
    VkQueue _presentQueue;
    VkSwapchainKHR _swapChain;
    std::vector<VkImage> _swapChainImages;
    VkFormat _swapChainImageFormat;
    VkExtent2D _swapChainExtent;
    std::vector<VkImageView> _swapChainImageViews;
    VkRenderPass _renderPass;
    VkDescriptorSetLayout _descriptorSetLayout;
    VkPipelineLayout _pipelineLayout;
    VkPipeline _graphicsPipeline;
    std::vector<VkFramebuffer> _swapChainFramebuffers;
    VkCommandPool _commandPool;
    std::vector<VkCommandBuffer> _commandBuffers;
    std::vector<VkSemaphore> _imageAvailableSemaphores;
    std::vector<VkSemaphore> _renderFinishedSemaphores;
    std::vector<VkFence> _inFlightFences;
    VkBuffer _vertexBuffer;
    VkDeviceMemory _vertexBufferMemory;
    VkBuffer _indexBuffer;
    VkDeviceMemory _indexBufferMemory;

    std::vector<VkBuffer> _uniformBuffers;
    std::vector<VkDeviceMemory> _uniformBuffersMemory;
    std::vector<void*> _uniformBuffersMapped;
    VkDescriptorPool _descriptorPool;
    std::vector<VkDescriptorSet> _descriptorSets;

    VkImage _textureImage;
    VkDeviceMemory _textureImageMemory;
    VkImageView _textureImageView;
    VkSampler _textureSampler;

    VkImage _depthImage;
    VkDeviceMemory _depthImageMemory;
    VkImageView _depthImageView;

    std::vector<VkBuffer> _shaderStorageBuffers;
    std::vector<VkDeviceMemory> _shaderStorageBuffersMemory;
    VkPipelineLayout _computePipelineLayout;
    VkPipeline _computePipeline;
    VkDescriptorSetLayout _computeDescriptorSetLayout;
    VkDescriptorPool _computeDescriptorPool;
    std::vector<VkDescriptorSet> _computeDescriptorSets;
    std::vector<VkCommandBuffer> _computeCommandBuffers;
    std::vector<VkFence> _computeInFlightFences;
    std::vector<VkSemaphore> _computeFinishedSemaphores;

};

#endif // COMPUTEANDGRAPHICS_H