#include "texture_manager.hpp"
#include "exceptions.hpp"
#include "vulkan_instance.hpp"

#include <glm/glm.hpp>
#include <vulkan/vulkan_core.h>

namespace gz::vk {
    TextureManager::TextureManager(VulkanInstance& instance)
        : vk(instance)
    {
        LogCreateInfo logCI{};
        logCI.logfile = "textureManager.log";
        logCI.storeLog = false;
        logCI.prefix = "Texture";
        logCI.prefixColor = Color::GREEN;
        logCI.timeColor = VULKAN_MESSAGE_TIME_COLOR;
        tLog = Log(std::move(logCI));

        vk.registerCleanupCallback(std::bind(&TextureManager::cleanup, this));
        atlases.insert({ 0, TextureAtlas(vk, 24, 24, 24, 24) });
        createDescriptorSetLayout();
        createDescriptorPool();
        createDescriptorSet();
        VulkanInstance::registerObjectUsingVulkan(ObjectUsingVulkan("TextureManager",
                    { &descriptorSetLayout, &descriptorPool, &descriptorSet },
                    {}));
    }


    void TextureManager::cleanup() {
        tLog.log0("cleanup: cleaning up");
        /* vkFreeDescriptorSets(vk.getDevice(), descriptorPool, 1, &descriptorSet); */
        vkDestroyDescriptorSetLayout(vk.getDevice(), descriptorSetLayout, NO_ALLOC);
        vkDestroyDescriptorPool(vk.getDevice(), descriptorPool, NO_ALLOC);
        for (auto& [i, atlas] : atlases) {
            atlas.cleanup();

        }
    }


	void TextureManager::getTexCoords(const std::string& textureName, glm::vec2& texCoords) {
        if (!textureInfos.contains(textureName)) {
            loadTextureFromFile(textureName);
        }
        TextureInfo& texI = textureInfos[textureName];
        tLog.log0("getTexCoords", texCoords, "textureInfo: topleft", texI.texCoordTopLeft, "botright", texI.texCoordBottomRight);
        /* texCoords.x = texI.texCoordTopLeft.x + texCoords.x * (texI.texCoordBottomRight.x - texI.texCoordTopLeft.x); */
        /* texCoords.y = texI.texCoordTopLeft.y + texCoords.y * (texI.texCoordBottomRight.y - texI.texCoordTopLeft.y); */
        texCoords = texI.texCoordTopLeft + texCoords * (texI.texCoordBottomRight - texI.texCoordTopLeft);
	}


	void TextureManager::loadTextureFromFile(const std::string& textureName) {
        int textureWidth, textureHeight, textureChannels;
        stbi_uc* pixels = stbi_load(("textures/" + textureName).c_str(), &textureWidth, &textureHeight, &textureChannels, STBI_rgb_alpha);
        if (!pixels) {
            throw Exception("Failed to load texture: textures/" + textureName, "TextureManager::loadTexture");
        }
        TextureInfo& texI = textureInfos[textureName];
        texI.atlas = 0;
        auto texCoords = atlases.at(0).addTexture(pixels, textureWidth, textureHeight);
        stbi_image_free(pixels);
        tLog.log0("TextureManager::loadTextureFromFile: TexCoords of new image:", texCoords.first, texCoords.second);
        texI.texCoordTopLeft = std::move(texCoords.first);
        texI.texCoordBottomRight = std::move(texCoords.second);
        tLog.log0("TextureManager::loadTextureFromFile: After loading:", atlases.at(0));
	}


//
// DESCRIPTOR SET
//
void TextureManager::createDescriptorSetLayout() {
    // combined image sampler
    VkDescriptorSetLayoutBinding samplerLayoutBinding{};
    samplerLayoutBinding.binding = 0;
    samplerLayoutBinding.descriptorCount = 1;
    samplerLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    samplerLayoutBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;
    /* samplerLayoutBinding.pImmutableSamplers = nullptr; */

    VkDescriptorSetLayoutCreateInfo descriptorSetLayoutCI{};
    descriptorSetLayoutCI.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    descriptorSetLayoutCI.bindingCount = 1;
    descriptorSetLayoutCI.pBindings = &samplerLayoutBinding;

    VkResult result = vkCreateDescriptorSetLayout(vk.getDevice(), &descriptorSetLayoutCI, nullptr, &descriptorSetLayout);
    if (result != VK_SUCCESS) {
        throw getVkException(result, "Could not create descriptorSetLayout", "TextureManager::createDescriptorSetLayout");
    }
    /* vLog("createDescriptorSetLayout: Created descriptor set layout."); */
}


void TextureManager::createDescriptorPool() {
    std::array<VkDescriptorPoolSize, 1> poolSizes;
    /* poolSizes[0].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER; */
    /* poolSizes[0].descriptorCount = static_cast<uint32_t>(vk.getMaxFramesInFlight()); */
    poolSizes[0].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    poolSizes[0].descriptorCount = static_cast<uint32_t>(vk.getMaxFramesInFlight());

    VkDescriptorPoolCreateInfo poolCI{};
    poolCI.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
    poolCI.poolSizeCount = static_cast<uint32_t>(poolSizes.size());
    poolCI.pPoolSizes = poolSizes.data();
    poolCI.maxSets = 1;
    VkResult result = vkCreateDescriptorPool(vk.getDevice(), &poolCI, nullptr, &descriptorPool);
   if (result != VK_SUCCESS) {
       throw getVkException(result, "Failed to create descriptor pool", "TextureManager::createDescriptorPool");
   }
   /* vLog("createDescriptorPool: Created descriptor pool."); */
}


void TextureManager::createDescriptorSet() {
    /* std::vector<VkDescriptorSetLayout> layouts(vk.getMaxFramesInFlight(), descriptorSetLayout); */
    VkDescriptorSetAllocateInfo setAI{};
    setAI.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    setAI.descriptorPool = descriptorPool;
    /* setAI.descriptorSetCount = static_cast<uint32_t>(vk.getMaxFramesInFlight()); */
    /* setAI.pSetLayouts = layouts.data(); */
    setAI.descriptorSetCount = 1;
    setAI.pSetLayouts = &descriptorSetLayout;

    /* descriptorSets.resize(vk.getMaxFramesInFlight()); */
    /* VkResult result = vkAllocateDescriptorSets(vk.getDevice(), &setAI, descriptorSets.data()); */
    VkResult result = vkAllocateDescriptorSets(vk.getDevice(), &setAI, &descriptorSet);
    if (result != VK_SUCCESS) {
       throw getVkException(result, "Failed to create descriptor set", "TextureManager::createDescriptorPool");
    }

    // configure sets
    std::vector<VkDescriptorImageInfo> imageInfos;
    for (auto it = atlases.cbegin(); it != atlases.cend(); it++) {
        imageInfos.emplace_back();
        imageInfos.back().imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        imageInfos.back().imageView = it->second.getTextureImageView();
        imageInfos.back().sampler = it->second.getTextureSampler();
    }

    std::array<VkWriteDescriptorSet, 1> descriptorW{};
    descriptorW[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    descriptorW[0].dstSet = descriptorSet;
    descriptorW[0].dstBinding = 0;
    descriptorW[0].dstArrayElement = 0;
    descriptorW[0].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    descriptorW[0].descriptorCount = static_cast<uint32_t>(imageInfos.size());
    descriptorW[0].pImageInfo = imageInfos.data();
    
    uint32_t descriptorWriteCount = static_cast<uint32_t>(descriptorW.size());
    uint32_t descriptorCopyCount = 0;
    vkUpdateDescriptorSets(vk.getDevice(), descriptorWriteCount, descriptorW.data(), descriptorCopyCount, nullptr);
    /* } */
    /* vLog("createDescriptorSets: Created descriptor sets."); */
}
}  // namespace gz::vk