#include <XCEngine/Resources/Material/Material.h>
#include <XCEngine/Resources/Shader/Shader.h>
#include <XCEngine/Core/Asset/ResourceManager.h>

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <vector>

namespace XCEngine {
namespace Resources {

namespace {

constexpr size_t kMaterialConstantSlotSize = 16;

bool IsPackedMaterialPropertyType(MaterialPropertyType type) {
    switch (type) {
    case MaterialPropertyType::Float:
    case MaterialPropertyType::Float2:
    case MaterialPropertyType::Float3:
    case MaterialPropertyType::Float4:
    case MaterialPropertyType::Int:
    case MaterialPropertyType::Int2:
    case MaterialPropertyType::Int3:
    case MaterialPropertyType::Int4:
    case MaterialPropertyType::Bool:
        return true;
    case MaterialPropertyType::Texture:
    case MaterialPropertyType::Cubemap:
    default:
        return false;
    }
}

void RemoveTextureBindingByName(
    Containers::Array<MaterialTextureBinding>& textureBindings,
    const Containers::String& name) {
    for (size_t bindingIndex = 0; bindingIndex < textureBindings.Size(); ++bindingIndex) {
        if (textureBindings[bindingIndex].name == name) {
            if (bindingIndex != textureBindings.Size() - 1) {
                textureBindings[bindingIndex] = std::move(textureBindings.Back());
                textureBindings[bindingIndex].slot = static_cast<Core::uint32>(bindingIndex);
            }
            textureBindings.PopBack();
            break;
        }
    }

    for (size_t bindingIndex = 0; bindingIndex < textureBindings.Size(); ++bindingIndex) {
        textureBindings[bindingIndex].slot = static_cast<Core::uint32>(bindingIndex);
    }
}

void WritePackedMaterialProperty(Core::uint8* destination, const MaterialProperty& property) {
    std::memset(destination, 0, kMaterialConstantSlotSize);

    switch (property.type) {
    case MaterialPropertyType::Float:
        std::memcpy(destination, property.value.floatValue, sizeof(float));
        break;
    case MaterialPropertyType::Float2:
        std::memcpy(destination, property.value.floatValue, sizeof(float) * 2);
        break;
    case MaterialPropertyType::Float3:
        std::memcpy(destination, property.value.floatValue, sizeof(float) * 3);
        break;
    case MaterialPropertyType::Float4:
        std::memcpy(destination, property.value.floatValue, sizeof(float) * 4);
        break;
    case MaterialPropertyType::Int:
        std::memcpy(destination, property.value.intValue, sizeof(Core::int32));
        break;
    case MaterialPropertyType::Int2:
        std::memcpy(destination, property.value.intValue, sizeof(Core::int32) * 2);
        break;
    case MaterialPropertyType::Int3:
        std::memcpy(destination, property.value.intValue, sizeof(Core::int32) * 3);
        break;
    case MaterialPropertyType::Int4:
        std::memcpy(destination, property.value.intValue, sizeof(Core::int32) * 4);
        break;
    case MaterialPropertyType::Bool: {
        const Core::uint32 boolValue = property.value.boolValue ? 1u : 0u;
        std::memcpy(destination, &boolValue, sizeof(boolValue));
        break;
    }
    case MaterialPropertyType::Texture:
    case MaterialPropertyType::Cubemap:
    default:
        break;
    }
}

} // namespace

Material::Material() = default;

Material::~Material() = default;

void Material::Release() {
    m_shader.Reset();
    m_renderQueue = static_cast<Core::int32>(MaterialRenderQueue::Geometry);
    m_renderState = MaterialRenderState();
    m_shaderPass.Clear();
    m_tags.Clear();
    m_properties.Clear();
    m_textureBindings.Clear();
    m_constantBufferData.Clear();
    m_changeVersion = 1;
    m_isValid = false;
    UpdateMemorySize();
}

void Material::SetShader(const ResourceHandle<Shader>& shader) {
    m_shader = shader;
    MarkChanged(false);
}

void Material::SetRenderQueue(Core::int32 renderQueue) {
    m_renderQueue = renderQueue;
    MarkChanged(false);
}

void Material::SetRenderQueue(MaterialRenderQueue renderQueue) {
    SetRenderQueue(static_cast<Core::int32>(renderQueue));
}

void Material::SetRenderState(const MaterialRenderState& renderState) {
    m_renderState = renderState;
    MarkChanged(false);
}

void Material::SetShaderPass(const Containers::String& shaderPass) {
    m_shaderPass = shaderPass;
    MarkChanged(false);
}

void Material::SetTag(const Containers::String& name, const Containers::String& value) {
    for (MaterialTagEntry& tag : m_tags) {
        if (tag.name == name) {
            tag.value = value;
            MarkChanged(false);
            return;
        }
    }

    MaterialTagEntry tag;
    tag.name = name;
    tag.value = value;
    m_tags.PushBack(tag);
    MarkChanged(false);
}

Containers::String Material::GetTag(const Containers::String& name) const {
    for (const MaterialTagEntry& tag : m_tags) {
        if (tag.name == name) {
            return tag.value;
        }
    }

    return Containers::String();
}

bool Material::HasTag(const Containers::String& name) const {
    for (const MaterialTagEntry& tag : m_tags) {
        if (tag.name == name) {
            return true;
        }
    }

    return false;
}

void Material::RemoveTag(const Containers::String& name) {
    for (size_t tagIndex = 0; tagIndex < m_tags.Size(); ++tagIndex) {
        if (m_tags[tagIndex].name == name) {
            if (tagIndex != m_tags.Size() - 1) {
                m_tags[tagIndex] = std::move(m_tags.Back());
            }
            m_tags.PopBack();
            MarkChanged(false);
            return;
        }
    }
}

void Material::ClearTags() {
    m_tags.Clear();
    MarkChanged(false);
}

Containers::String Material::GetTagName(Core::uint32 index) const {
    return index < m_tags.Size() ? m_tags[index].name : Containers::String();
}

Containers::String Material::GetTagValue(Core::uint32 index) const {
    return index < m_tags.Size() ? m_tags[index].value : Containers::String();
}

void Material::SetFloat(const Containers::String& name, float value) {
    RemoveTextureBindingByName(m_textureBindings, name);
    MaterialProperty prop;
    prop.name = name;
    prop.type = MaterialPropertyType::Float;
    prop.value.floatValue[0] = value;
    prop.refCount = 1;
    m_properties.Insert(name, prop);
    MarkChanged(true);
}

void Material::SetFloat2(const Containers::String& name, const Math::Vector2& value) {
    RemoveTextureBindingByName(m_textureBindings, name);
    MaterialProperty prop;
    prop.name = name;
    prop.type = MaterialPropertyType::Float2;
    prop.value.floatValue[0] = value.x;
    prop.value.floatValue[1] = value.y;
    prop.refCount = 1;
    m_properties.Insert(name, prop);
    MarkChanged(true);
}

void Material::SetFloat3(const Containers::String& name, const Math::Vector3& value) {
    RemoveTextureBindingByName(m_textureBindings, name);
    MaterialProperty prop;
    prop.name = name;
    prop.type = MaterialPropertyType::Float3;
    prop.value.floatValue[0] = value.x;
    prop.value.floatValue[1] = value.y;
    prop.value.floatValue[2] = value.z;
    prop.refCount = 1;
    m_properties.Insert(name, prop);
    MarkChanged(true);
}

void Material::SetFloat4(const Containers::String& name, const Math::Vector4& value) {
    RemoveTextureBindingByName(m_textureBindings, name);
    MaterialProperty prop;
    prop.name = name;
    prop.type = MaterialPropertyType::Float4;
    prop.value.floatValue[0] = value.x;
    prop.value.floatValue[1] = value.y;
    prop.value.floatValue[2] = value.z;
    prop.value.floatValue[3] = value.w;
    prop.refCount = 1;
    m_properties.Insert(name, prop);
    MarkChanged(true);
}

void Material::SetInt(const Containers::String& name, Core::int32 value) {
    RemoveTextureBindingByName(m_textureBindings, name);
    MaterialProperty prop;
    prop.name = name;
    prop.type = MaterialPropertyType::Int;
    prop.value.intValue[0] = value;
    prop.refCount = 1;
    m_properties.Insert(name, prop);
    MarkChanged(true);
}

void Material::SetBool(const Containers::String& name, bool value) {
    RemoveTextureBindingByName(m_textureBindings, name);
    MaterialProperty prop;
    prop.name = name;
    prop.type = MaterialPropertyType::Bool;
    prop.value.boolValue = value;
    prop.refCount = 1;
    m_properties.Insert(name, prop);
    MarkChanged(true);
}

void Material::SetTexture(const Containers::String& name, const ResourceHandle<Texture>& texture) {
    MaterialProperty prop;
    prop.name = name;
    prop.type = MaterialPropertyType::Texture;
    prop.refCount = 1;
    m_properties.Insert(name, prop);

    for (auto& binding : m_textureBindings) {
        if (binding.name == name) {
            binding.texture = texture;
            MarkChanged(false);
            return;
        }
    }

    MaterialTextureBinding binding;
    binding.name = name;
    binding.slot = static_cast<Core::uint32>(m_textureBindings.Size());
    binding.texture = texture;
    m_textureBindings.PushBack(binding);
    MarkChanged(false);
}

float Material::GetFloat(const Containers::String& name) const {
    auto* prop = m_properties.Find(name);
    if (prop && prop->type == MaterialPropertyType::Float) {
        return prop->value.floatValue[0];
    }
    return 0.0f;
}

Math::Vector2 Material::GetFloat2(const Containers::String& name) const {
    auto* prop = m_properties.Find(name);
    if (prop && prop->type == MaterialPropertyType::Float2) {
        return Math::Vector2(prop->value.floatValue[0], prop->value.floatValue[1]);
    }
    return Math::Vector2::Zero();
}

Math::Vector3 Material::GetFloat3(const Containers::String& name) const {
    auto* prop = m_properties.Find(name);
    if (prop && prop->type == MaterialPropertyType::Float3) {
        return Math::Vector3(prop->value.floatValue[0], prop->value.floatValue[1], prop->value.floatValue[2]);
    }
    return Math::Vector3::Zero();
}

Math::Vector4 Material::GetFloat4(const Containers::String& name) const {
    auto* prop = m_properties.Find(name);
    if (prop && prop->type == MaterialPropertyType::Float4) {
        return Math::Vector4(prop->value.floatValue[0], prop->value.floatValue[1], 
                            prop->value.floatValue[2], prop->value.floatValue[3]);
    }
    return Math::Vector4::Zero();
}

Core::int32 Material::GetInt(const Containers::String& name) const {
    auto* prop = m_properties.Find(name);
    if (prop && prop->type == MaterialPropertyType::Int) {
        return prop->value.intValue[0];
    }
    return 0;
}

bool Material::GetBool(const Containers::String& name) const {
    auto* prop = m_properties.Find(name);
    if (prop && prop->type == MaterialPropertyType::Bool) {
        return prop->value.boolValue;
    }
    return false;
}

ResourceHandle<Texture> Material::GetTexture(const Containers::String& name) const {
    for (const auto& binding : m_textureBindings) {
        if (binding.name == name) {
            return binding.texture;
        }
    }
    return ResourceHandle<Texture>();
}

Containers::String Material::GetTextureBindingName(Core::uint32 index) const {
    return index < m_textureBindings.Size() ? m_textureBindings[index].name : Containers::String();
}

ResourceHandle<Texture> Material::GetTextureBindingTexture(Core::uint32 index) const {
    return index < m_textureBindings.Size() ? m_textureBindings[index].texture : ResourceHandle<Texture>();
}

std::vector<MaterialProperty> Material::GetProperties() const {
    std::vector<MaterialProperty> properties;
    const auto pairs = m_properties.GetPairs();
    properties.reserve(pairs.Size());
    for (const auto& pair : pairs) {
        properties.push_back(pair.second);
    }
    return properties;
}

void Material::UpdateConstantBuffer() {
    std::vector<const MaterialProperty*> packedProperties;
    const auto pairs = m_properties.GetPairs();
    packedProperties.reserve(pairs.Size());
    for (const auto& pair : pairs) {
        if (IsPackedMaterialPropertyType(pair.second.type)) {
            packedProperties.push_back(&pair.second);
        }
    }

    std::sort(
        packedProperties.begin(),
        packedProperties.end(),
        [](const MaterialProperty* left, const MaterialProperty* right) {
            return std::strcmp(left->name.CStr(), right->name.CStr()) < 0;
        });

    m_constantBufferData.Clear();
    m_constantBufferData.Resize(packedProperties.size() * kMaterialConstantSlotSize);
    if (!packedProperties.empty()) {
        std::memset(m_constantBufferData.Data(), 0, m_constantBufferData.Size());
    }

    for (size_t propertyIndex = 0; propertyIndex < packedProperties.size(); ++propertyIndex) {
        WritePackedMaterialProperty(
            m_constantBufferData.Data() + propertyIndex * kMaterialConstantSlotSize,
            *packedProperties[propertyIndex]);
    }
    UpdateMemorySize();
}

void Material::RecalculateMemorySize() {
    UpdateMemorySize();
}

bool Material::HasProperty(const Containers::String& name) const {
    return m_properties.Contains(name);
}

void Material::RemoveProperty(const Containers::String& name) {
    const MaterialProperty* property = m_properties.Find(name);
    const bool removeTextureBinding =
        property != nullptr &&
        (property->type == MaterialPropertyType::Texture || property->type == MaterialPropertyType::Cubemap);

    m_properties.Erase(name);

    if (removeTextureBinding) {
        RemoveTextureBindingByName(m_textureBindings, name);
    }

    MarkChanged(true);
}

void Material::ClearAllProperties() {
    m_properties.Clear();
    m_textureBindings.Clear();
    m_constantBufferData.Clear();
    MarkChanged(false);
}

void Material::MarkChanged(bool updateConstantBuffer) {
    if (updateConstantBuffer) {
        UpdateConstantBuffer();
    } else {
        UpdateMemorySize();
    }

    ++m_changeVersion;
}

void Material::UpdateMemorySize() {
    m_memorySize = m_constantBufferData.Size() +
                   sizeof(MaterialRenderState) +
                   m_shaderPass.Length() +
                   m_tags.Size() * sizeof(MaterialTagEntry) +
                   m_textureBindings.Size() * sizeof(MaterialTextureBinding) +
                   m_properties.Size() * sizeof(MaterialProperty) +
                   m_name.Length() +
                   m_path.Length();

    for (const MaterialTagEntry& tag : m_tags) {
        m_memorySize += tag.name.Length();
        m_memorySize += tag.value.Length();
    }

    for (const auto& binding : m_textureBindings) {
        m_memorySize += binding.name.Length();
    }
}

} // namespace Resources
} // namespace XCEngine