#include <AssetManager/AssetManager.hpp>
#include <Material/Material.hpp>
#include <Utilities/Utilities.hpp>
#include <Log/Log.hpp>
#include <GraphicsContext/GraphicsContext.hpp>
#include <Utilities/Extern.hpp>
#include <Vk/ShaderConst.hpp>
#include <string>

std::unordered_map<std::string, Material*> Material::materialTable;

#define DECL_MATERIAL(name) Material Material::name(#name)
#define DECL_MATERIAL_LM(name, lightValue) Material Material::name(#name, lightValue)
#include <Material/MaterialList.hpp>

Material Material::internalGameIcon;
Material Material::white;

Material::Material(const char *name, unsigned char lightLevel) : name(name)
{
    lightPixel[0] = lightPixel[1] = lightPixel[2] = lightLevel;
    lightPixel[3] = 0xFF;

    if (name != nullptr)
    {
        materialTable[name] = this;
    }
}

void Material::addComponent(std::shared_ptr<Vk::Device> device, Vk::CommandBuffer &cmdbuf, const void *pixels)
{
    addComponent(device, cmdbuf, pixels, widthValue, heightValue);
}

void Material::addComponent(std::shared_ptr<Vk::Device> device, Vk::CommandBuffer &cmdbuf, const void *pixels, int width, int height)
{
    size_t imageDataSize = width * height * 4;

    Component comp;
    
    // Upload the pixels to the staging buffer.
    comp.stagingBuffer = std::make_shared<Vk::Buffer>(
        device,
        imageDataSize,
        VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
    );

    std::memcpy(comp.stagingBuffer->hostdata(), pixels, imageDataSize);

    // Create the texture image.
    comp.image = std::make_shared<Vk::Image>(
        device,
        VK_IMAGE_TYPE_2D,
        width, height,
        VK_FORMAT_R8G8B8A8_SRGB,
        VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
        VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT
    );

    // Create the image view.
    comp.imageView = std::make_shared<Vk::ImageView>(device, comp.image, VK_IMAGE_VIEW_TYPE_2D, VK_IMAGE_ASPECT_COLOR_BIT);

    // Add the component to the list.
    components.push_back(comp);

    // Add commands to the buffer to transfer from the staging buffer to the image.
    cmdbuf.transitionImageLayout(*comp.image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
    cmdbuf.copyBufferToImage(*comp.stagingBuffer, *comp.image, width, height);
    cmdbuf.transitionImageLayout(*comp.image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
}

void Material::buildMetadata(std::shared_ptr<Vk::Device> device, Vk::CommandBuffer &cmdbuf)
{
    metadataStagingBuffer = std::make_shared<Vk::Buffer>(
        device,
        sizeof(Metadata),
        VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
    );

    metadataBuffer = std::make_shared<Vk::Buffer>(
        device,
        sizeof(Metadata),
        VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
        VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT
    );

    Metadata *metadata = (Metadata*) metadataStagingBuffer->hostdata();
    metadata->size = glm::ivec2(widthValue, heightValue);
    metadata->flags = flags;
    
    cmdbuf.transitionBufferToTransferDst(*metadataBuffer);
    cmdbuf.copyBuffer(*metadataStagingBuffer, 0, *metadataBuffer, 0, sizeof(Metadata));
    cmdbuf.transitionBufferToShaderRead(*metadataBuffer);
}

void Material::buildDescriptorSet(std::shared_ptr<Vk::Device> device)
{
    descriptorPool = Vk::DescriptorPool::Builder()
        .addDescriptorType(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, GLSL_IMAGE_COUNT)
        .addDescriptorType(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1)
        .build(device, 1);

    descriptorSetPriv = std::make_shared<Vk::DescriptorSet>(device, descriptorPool, GraphicsContext::materialLayout);

    int i;
    for (i=0; i<(int) components.size(); i++)
    {
        descriptorSetPriv->setImageSampler(i, *components[i].imageView, Vk::Sampler::pixelated());
    }

    descriptorSetPriv->setUniformBuffer(GLSL_IMAGE_METADATA_INDEX, *metadataBuffer, 0, sizeof(Metadata));
}

void Material::finalize()
{
    for (Component &comp : components)
    {
        comp.stagingBuffer = nullptr;
    }

    metadataStagingBuffer = nullptr;
}

void Material::load(std::atomic<unsigned int> &tasksDone)
{
    Log::info(formatString("Loading material '$+' ...", name));

    std::string imageFileName = assets.resolve(formatString("images/$+.png", name));
    load(imageFileName, tasksDone);
}

void Material::load(const std::string &imageFileName, std::atomic<unsigned int> &tasksDone)
{
    std::shared_ptr<Vk::CommandBuffer> cmdbuf = GraphicsContext::instance->createAsyncCommandBuffer();
    cmdbuf->reset();
    cmdbuf->begin();

    int comp;
    pixels = stbi_load(imageFileName.c_str(), &widthValue, &heightValue, &comp, STBI_rgb_alpha);
    
    if (pixels == nullptr)
    {
        crash(formatString("Failed to load image file '$+'!", imageFileName));
    }

    addComponent(GraphicsContext::instance->getDevice(), *cmdbuf, pixels);

    static uint8_t defaultNormal[4] = {0x00, 0x00, 0xFF, 0xFF};
    static uint8_t defaultSpecular[4] = {0x00, 0x00, 0x00, 0xFF};
    static uint8_t defaultBias[4] = {0x00, 0x00, 0x00, 0x00};

    addComponent(GraphicsContext::instance->getDevice(), *cmdbuf, defaultNormal, 1, 1);
    addComponent(GraphicsContext::instance->getDevice(), *cmdbuf, lightPixel, 1, 1);
    addComponent(GraphicsContext::instance->getDevice(), *cmdbuf, defaultSpecular, 1, 1);
    addComponent(GraphicsContext::instance->getDevice(), *cmdbuf, defaultBias, 1, 1);

    buildMetadata(GraphicsContext::instance->getDevice(), *cmdbuf);
    buildDescriptorSet(GraphicsContext::instance->getDevice());

    GraphicsContext::onNextFrame([this, cmdbuf, &tasksDone](GraphicsContext &)
    {
        std::shared_ptr<Vk::Fence> fence = std::make_shared<Vk::Fence>(GraphicsContext::instance->getDevice());
        cmdbuf->end();
        cmdbuf->submit(fence.get());
        fence->wait();

        finalize();

        tasksDone++;
    });
}

int Material::width()
{
    return widthValue;
}

int Material::height()
{
    return heightValue;
}

unsigned char* Material::getPixels()
{
    return pixels;
}

Vk::DescriptorSet& Material::descriptorSet()
{
    return *descriptorSetPriv;
}

Material* Material::byName(const std::string &name)
{
    auto it = materialTable.find(name);
    if (it == materialTable.end())
    {
        return nullptr;
    }

    return it->second;
}