//
// Created by Matthew on 2025/3/22.
//

#ifndef MPPREVIEW_VULKAN_HDR_GENERATOR_H
#define MPPREVIEW_VULKAN_HDR_GENERATOR_H

#include <vector>
#include <string>
#include <vulkan/vulkan.h>
#include "BmpLoader.h"

struct HdrMergeParams {
    uint32_t imageCount;
    uint32_t width;
    uint32_t height;
    float exposureValues[16]; // Support up to 16 images
};

class VulkanHdrGenerator {
private:
    VkInstance instance;
    VkPhysicalDevice physicalDevice;
    VkDevice device;
    VkQueue computeQueue;
    VkCommandPool commandPool;
    VkDescriptorPool descriptorPool;
    VkCommandBuffer commandBuffer;

    VkBuffer stagingBuffer;
    VkDeviceMemory stagingBufferMemory;

    VkBuffer inputBuffer = nullptr;
    VkDeviceMemory inputBufferMemory;
    VkBuffer outputBuffer;
    VkDeviceMemory outputBufferMemory;
    VkBuffer paramsBuffer;
    VkDeviceMemory paramsBufferMemory;

    VkShaderModule computeShaderModule;
    VkPipeline computePipeline;
    VkPipelineLayout pipelineLayout;
    VkDescriptorSetLayout descriptorSetLayout;
    VkDescriptorSet descriptorSet;

    uint32_t computeQueueFamilyIndex;

    std::string compFilePath;

    void setupVulkan();
    void createComputeResources();
    void createBuffers(VkDeviceSize inputSize, VkDeviceSize outputSize);
    void createDescriptorSet();
    uint32_t findComputeQueueFamily(VkPhysicalDevice device);
    uint32_t findMemoryType(uint32_t typeFilter, VkMemoryPropertyFlags properties);
    void processImageBatch(const std::vector<std::vector<float>>& images,
                           std::vector<float>& outputData,
                           int32_t width, int32_t height,
                           const std::vector<float>& exposures);

    VkShaderModule createShaderModule(const std::vector<uint8_t>& code);
    std::vector<uint8_t> readFile(const std::string& filename);
    void cleanup();

public:
    VulkanHdrGenerator(const std::string& compFile);
    ~VulkanHdrGenerator();

    // Generate HDR from multiple BMP files with tile-based processing
    bool generateHdr(const std::vector<std::string>& bmpFiles,
                     const std::string& outputFile,
                     const std::vector<float>& exposureValues,
                     int32_t tileWidth = 256, int32_t tileHeight = 256);
};

#endif //MPPREVIEW_VULKAN_HDR_GENERATOR_H