#include <destrum/Graphics/Caches/MaterialCache.h>

#include <destrum/Graphics/GfxDevice.h>
#include <destrum/Graphics/Util.h>

#include "spdlog/spdlog.h"
#include <algorithm>
#include <string>
#include <stdexcept>
#include <utility>

void MaterialCache::init(
    GfxDevice& gfxDevice,
    MaterialDefaultTextures defaults)
{
    try {
        defaultTextures = defaults;
        device = &gfxDevice;

        materialDataBuffer = gfxDevice.createBuffer(
            MAX_MATERIALS * sizeof(MaterialData),
            VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT |
                VK_BUFFER_USAGE_TRANSFER_DST_BIT,
            VMA_MEMORY_USAGE_CPU_TO_GPU);

        vkutil::addDebugLabel(
            gfxDevice.getDevice(),
            materialDataBuffer.buffer,
            "material data");

        Material placeholderMaterial = Material::SimpleColor(
            glm::vec3{1.0f},
            "PLACEHOLDER_MATERIAL");

        placeholderMaterialId = addMaterial(placeholderMaterial);
        gfxDevice.getMemoryManager().flushAllocation(materialDataBuffer);
    } catch (...) {
        cleanup(gfxDevice);
        throw;
    }
}

void MaterialCache::cleanup(GfxDevice& gfxDevice)
{
    gfxDevice.destroyBuffer(materialDataBuffer);
    materialDataBuffer = {};
    materials.clear();
    materialKeys.clear();
    freeIds = std::make_shared<std::vector<MaterialID>>();
    placeholderMaterialId = NULL_MATERIAL_ID;
    device = nullptr;
}

MaterialID MaterialCache::addMaterial(Material material, std::string key)
{
    MaterialData* data = static_cast<MaterialData*>(materialDataBuffer.info.pMappedData);

    const MaterialID id = getFreeMaterialId();
    if (id >= MAX_MATERIALS) {
        throw std::runtime_error("MaterialCache capacity exhausted");
    }
    if (data == nullptr) {
        throw std::runtime_error("MaterialCache material buffer is not mapped");
    }

    // Safe without waiting: no frame in flight can reference an id that didn't exist yet.
    // Changing an existing material goes through GameRenderer::updateMaterialGPU.
    data[id] = toMaterialData(material);
    device->getMemoryManager().flushAllocation(
        materialDataBuffer,
        static_cast<VkDeviceSize>(id) * sizeof(MaterialData),
        sizeof(MaterialData));

    const MaterialID materialId = id;
    if (key.empty() && !material.assetReference.empty()) {
        key = material.assetReference.cacheKey();
    } else if (key.empty()) {
        key = material.name.empty()
            ? "material:" + std::to_string(materialId)
            : material.name;
    }
    const std::string baseKey = key;
    std::size_t suffix = 1;
    while (std::find(materialKeys.begin(), materialKeys.end(), key) != materialKeys.end()) {
        key = baseKey + "#" + std::to_string(suffix++);
    }

    if (materialId == materials.size()) {
        materials.push_back(std::move(material));
        materialKeys.push_back(std::move(key));
    } else {
        std::erase(*freeIds, materialId);
        materials[materialId] = std::move(material);
        materialKeys[materialId] = std::move(key);
    }

    return materialId;
}

void MaterialCache::removeMaterial(MaterialID id)
{
    if (!isValid(id) || id == placeholderMaterialId) {
        return;
    }

    materials[id] = {};
    materialKeys[id].clear();
    device->deferDestroy([ids = freeIds, id] { ids->push_back(id); });
}

MaterialID MaterialCache::addSimpleTextureMaterial(ImageID textureID)
{
    Material material = Material::SimpleColor(
        glm::vec3{1.0f},
        "simple texture material");
    material.diffuseTexture = textureID;

    return addMaterial(std::move(material));
}

MaterialID MaterialCache::addSimpleColorMaterial(glm::vec3 color, std::string name)
{
    return addMaterial(Material::SimpleColor(std::move(color), std::move(name)));
}

const Material& MaterialCache::getMaterial(MaterialID id) const
{
    return materials.at(id);
}

const std::string& MaterialCache::getMaterialKey(MaterialID id) const
{
    return materialKeys.at(id);
}

std::optional<MaterialID> MaterialCache::findMaterialByKey(std::string_view key) const
{
    if (key.empty()) {
        return std::nullopt;
    }
    const auto it = std::find(materialKeys.begin(), materialKeys.end(), key);
    if (it == materialKeys.end()) {
        return std::nullopt;
    }
    return static_cast<MaterialID>(std::distance(materialKeys.begin(), it));
}

std::optional<MaterialID> MaterialCache::findMaterialByName(
    std::string_view name,
    std::string_view sourceReference) const
{
    std::optional<MaterialID> fallback;
    for (MaterialID id = 0; id < materials.size(); ++id) {
        const auto& material = materials[id];
        if (!isValid(id) || material.name != name) {
            continue;
        }

        if (sourceReference.empty()) {
            return id;
        }

        if (!material.assetReference.empty() &&
            material.assetReference.path == sourceReference) {
            return id;
        }

        // A manually-created material may have the same name but no source
        // reference. Keep it as a fallback for scenes saved before canonical
        // file-backed material keys were introduced.
        if (!fallback.has_value() && material.assetReference.empty()) {
            fallback = id;
        }
    }

    return fallback;
}

MaterialData MaterialCache::toMaterialData(const Material& material) const
{
    const auto getTextureOrElse = [](ImageID imageId, ImageID placeholder)
    {
        return imageId != NULL_IMAGE_ID ? imageId : placeholder;
    };

    return MaterialData{
        .baseColor = glm::vec4(material.baseColor, 1.0f),
        .emissiveColor = glm::vec4(material.emissiveColor, 0.0f),

        .metallicRoughnessNormalOcclusion = glm::vec4{
            material.metallicFactor,
            material.roughnessFactor,
            material.normalScale,
            material.occlusionStrength
        },

        .textureFilteringMode = static_cast<std::uint32_t>(material.textureFilteringMode),
        .diffuseTex = getTextureOrElse(material.diffuseTexture, defaultTextures.white),
        .normalTex = getTextureOrElse(material.normalTexture, defaultTextures.normal),
        .metallicTex = getTextureOrElse(material.metallicTexture, defaultTextures.metallicRoughness),
        .roughnessTex = getTextureOrElse(material.roughnessTexture, defaultTextures.metallicRoughness),
        .occlusionTex = getTextureOrElse(material.occlusionTexture, defaultTextures.white),
        .emissiveTex = getTextureOrElse(material.emissiveTexture, defaultTextures.emissive),
        ._pad0 = 0,
    };
}

MaterialID MaterialCache::getFreeMaterialId() const
{
    return freeIds->empty() ? static_cast<MaterialID>(materials.size()) : freeIds->back();
}

MaterialID MaterialCache::getPlaceholderMaterialId() const
{
    assert(placeholderMaterialId != NULL_MATERIAL_ID && "MaterialCache::init not called");
    return placeholderMaterialId;
}

Material& MaterialCache::getMaterialMutable(MaterialID id)
{
    assert(id < materials.size());
    return materials.at(id);
}

MaterialData MaterialCache::getMaterialData(MaterialID id) const
{
    assert(id < materials.size());
    return toMaterialData(materials[id]);
}