XCEngine/engine/src/Resources/Mesh/MeshLoader.cpp

#include <XCEngine/Resources/Mesh/MeshLoader.h>
#include <XCEngine/Resources/Mesh/MeshImportSettings.h>
#include <XCEngine/Core/Asset/ResourceManager.h>
#include <XCEngine/Core/Math/Matrix4.h>
#include <assimp/Importer.hpp>
#include <assimp/postprocess.h>
#include <assimp/scene.h>
#include <filesystem>
#include <fstream>
#include <limits>
#include <string>
#include <vector>

namespace XCEngine {
namespace Resources {

namespace {

struct ImportedMeshData {
    std::vector<StaticMeshVertex> vertices;
    std::vector<Core::uint32> indices;
    VertexAttribute attributes = VertexAttribute::Position;
};

Math::Matrix4 ConvertAssimpMatrix(const aiMatrix4x4& matrix) {
    Math::Matrix4 result;
    result.m[0][0] = matrix.a1; result.m[0][1] = matrix.a2; result.m[0][2] = matrix.a3; result.m[0][3] = matrix.a4;
    result.m[1][0] = matrix.b1; result.m[1][1] = matrix.b2; result.m[1][2] = matrix.b3; result.m[1][3] = matrix.b4;
    result.m[2][0] = matrix.c1; result.m[2][1] = matrix.c2; result.m[2][2] = matrix.c3; result.m[2][3] = matrix.c4;
    result.m[3][0] = matrix.d1; result.m[3][1] = matrix.d2; result.m[3][2] = matrix.d3; result.m[3][3] = matrix.d4;
    return result;
}

Core::uint32 BuildPostProcessFlags(const MeshImportSettings& settings) {
    Core::uint32 flags = aiProcess_Triangulate |
                         aiProcess_JoinIdenticalVertices |
                         aiProcess_SortByPType |
                         aiProcess_ValidateDataStructure |
                         aiProcess_FindInvalidData;

    if (settings.GetAxisConversion()) {
        flags |= aiProcess_MakeLeftHanded;
        flags |= aiProcess_FlipWindingOrder;
    }

    if (settings.HasImportFlag(MeshImportFlags::FlipUVs)) {
        flags |= aiProcess_FlipUVs;
    }

    if (settings.HasImportFlag(MeshImportFlags::FlipWindingOrder)) {
        flags |= aiProcess_FlipWindingOrder;
    }

    if (settings.HasImportFlag(MeshImportFlags::GenerateNormals)) {
        flags |= aiProcess_GenSmoothNormals;
    }

    if (settings.HasImportFlag(MeshImportFlags::GenerateTangents)) {
        flags |= aiProcess_CalcTangentSpace;
    }

    if (settings.HasImportFlag(MeshImportFlags::OptimizeMesh)) {
        flags |= aiProcess_ImproveCacheLocality;
        flags |= aiProcess_OptimizeMeshes;
        flags |= aiProcess_OptimizeGraph;
    }

    return flags;
}

ImportedMeshData ImportSingleMesh(const aiMesh& mesh,
                                  const Math::Matrix4& worldTransform,
                                  float globalScale,
                                  const Math::Vector3& offset) {
    ImportedMeshData result;
    result.vertices.reserve(mesh.mNumVertices);
    result.indices.reserve(mesh.mNumFaces * 3);

    VertexAttribute attributes = VertexAttribute::Position;
    if (mesh.HasNormals()) {
        attributes = attributes | VertexAttribute::Normal;
    }
    if (mesh.HasTangentsAndBitangents()) {
        attributes = attributes | VertexAttribute::Tangent | VertexAttribute::Bitangent;
    }
    if (mesh.HasTextureCoords(0)) {
        attributes = attributes | VertexAttribute::UV0;
    }

    const Math::Matrix4 normalTransform = worldTransform.Inverse().Transpose();
    const float appliedScale = globalScale;

    for (Core::uint32 vertexIndex = 0; vertexIndex < mesh.mNumVertices; ++vertexIndex) {
        StaticMeshVertex vertex;

        const aiVector3D& position = mesh.mVertices[vertexIndex];
        const Math::Vector3 transformedPosition = worldTransform.MultiplyPoint(Math::Vector3(position.x, position.y, position.z));
        vertex.position = transformedPosition * appliedScale + offset;

        if (mesh.HasNormals()) {
            const aiVector3D& normal = mesh.mNormals[vertexIndex];
            vertex.normal = normalTransform.MultiplyVector(Math::Vector3(normal.x, normal.y, normal.z)).Normalized();
        }

        if (mesh.HasTangentsAndBitangents()) {
            const aiVector3D& tangent = mesh.mTangents[vertexIndex];
            const aiVector3D& bitangent = mesh.mBitangents[vertexIndex];
            vertex.tangent = normalTransform.MultiplyVector(Math::Vector3(tangent.x, tangent.y, tangent.z)).Normalized();
            vertex.bitangent = normalTransform.MultiplyVector(Math::Vector3(bitangent.x, bitangent.y, bitangent.z)).Normalized();
        }

        if (mesh.HasTextureCoords(0)) {
            const aiVector3D& uv = mesh.mTextureCoords[0][vertexIndex];
            vertex.uv0 = Math::Vector2(uv.x, uv.y);
        }

        result.vertices.push_back(vertex);
    }

    for (Core::uint32 faceIndex = 0; faceIndex < mesh.mNumFaces; ++faceIndex) {
        const aiFace& face = mesh.mFaces[faceIndex];
        if (face.mNumIndices != 3) {
            continue;
        }

        result.indices.push_back(face.mIndices[0]);
        result.indices.push_back(face.mIndices[1]);
        result.indices.push_back(face.mIndices[2]);
    }

    result.attributes = attributes;
    return result;
}

void ProcessNode(const aiNode& node,
                 const aiScene& scene,
                 const Math::Matrix4& parentTransform,
                 const MeshImportSettings& settings,
                 std::vector<StaticMeshVertex>& vertices,
                 std::vector<Core::uint32>& indices,
                 Containers::Array<MeshSection>& sections,
                 VertexAttribute& attributes) {
    const Math::Matrix4 worldTransform = parentTransform * ConvertAssimpMatrix(node.mTransformation);
    const float globalScale = settings.GetScale() * settings.GetImportScale();

    for (Core::uint32 meshIndex = 0; meshIndex < node.mNumMeshes; ++meshIndex) {
        const aiMesh* mesh = scene.mMeshes[node.mMeshes[meshIndex]];
        if (mesh == nullptr || mesh->mNumVertices == 0) {
            continue;
        }

        const ImportedMeshData meshData = ImportSingleMesh(*mesh, worldTransform, globalScale, settings.GetOffset());
        if (meshData.vertices.empty() || meshData.indices.empty()) {
            continue;
        }

        const Core::uint32 baseVertex = static_cast<Core::uint32>(vertices.size());
        const Core::uint32 startIndex = static_cast<Core::uint32>(indices.size());

        vertices.insert(vertices.end(), meshData.vertices.begin(), meshData.vertices.end());
        indices.reserve(indices.size() + meshData.indices.size());
        for (Core::uint32 index : meshData.indices) {
            indices.push_back(baseVertex + index);
        }

        MeshSection section{};
        section.baseVertex = baseVertex;
        section.vertexCount = static_cast<Core::uint32>(meshData.vertices.size());
        section.startIndex = startIndex;
        section.indexCount = static_cast<Core::uint32>(meshData.indices.size());
        section.materialID = mesh->mMaterialIndex;
        sections.PushBack(section);

        attributes = attributes | meshData.attributes;
    }

    for (Core::uint32 childIndex = 0; childIndex < node.mNumChildren; ++childIndex) {
        const aiNode* child = node.mChildren[childIndex];
        if (child != nullptr) {
            ProcessNode(*child, scene, worldTransform, settings, vertices, indices, sections, attributes);
        }
    }
}

} // namespace

MeshLoader::MeshLoader() = default;
MeshLoader::~MeshLoader() = default;

Containers::Array<Containers::String> MeshLoader::GetSupportedExtensions() const {
    Containers::Array<Containers::String> extensions;
    extensions.PushBack(Containers::String("fbx"));
    extensions.PushBack(Containers::String("obj"));
    extensions.PushBack(Containers::String("gltf"));
    extensions.PushBack(Containers::String("glb"));
    extensions.PushBack(Containers::String("dae"));
    extensions.PushBack(Containers::String("stl"));
    return extensions;
}

bool MeshLoader::CanLoad(const Containers::String& path) const {
    Containers::String ext = GetExtension(path).ToLower();
    
    return ext == "fbx" || ext == "obj" || ext == "gltf" || 
           ext == "glb" || ext == "dae" || ext == "stl";
}

LoadResult MeshLoader::Load(const Containers::String& path, const ImportSettings* settings) {
    const Containers::String ext = GetExtension(path).ToLower();
    
    if (!CanLoad(path)) {
        return LoadResult(Containers::String("Unsupported mesh format: ") + ext);
    }

    std::ifstream file(path.CStr(), std::ios::binary);
    if (!file.is_open()) {
        return LoadResult(Containers::String("Failed to read file: ") + path);
    }

    MeshImportSettings defaultSettings;
    const MeshImportSettings* resolvedSettings = dynamic_cast<const MeshImportSettings*>(settings);
    if (resolvedSettings == nullptr) {
        resolvedSettings = &defaultSettings;
    }

    Assimp::Importer importer;
    const aiScene* scene = importer.ReadFile(path.CStr(), BuildPostProcessFlags(*resolvedSettings));
    if (scene == nullptr || scene->mRootNode == nullptr || (scene->mFlags & AI_SCENE_FLAGS_INCOMPLETE) != 0) {
        const char* errorText = importer.GetErrorString();
        return LoadResult(Containers::String("Assimp failed to load mesh: ") +
                          Containers::String(errorText != nullptr ? errorText : "Unknown error"));
    }

    std::vector<StaticMeshVertex> vertices;
    std::vector<Core::uint32> indices;
    Containers::Array<MeshSection> sections;
    VertexAttribute attributes = VertexAttribute::Position;

    ProcessNode(*scene->mRootNode,
                *scene,
                Math::Matrix4::Identity(),
                *resolvedSettings,
                vertices,
                indices,
                sections,
                attributes);

    if (vertices.empty() || indices.empty()) {
        return LoadResult(Containers::String("No triangle mesh data found in: ") + path);
    }

    auto* mesh = new Mesh();

    IResource::ConstructParams params;
    const std::string fileName = std::filesystem::path(path.CStr()).filename().string();
    params.name = Containers::String(fileName.c_str());
    params.path = path;
    params.guid = ResourceGUID::Generate(path);
    params.memorySize = 0;
    
    mesh->Initialize(params);

    mesh->SetVertexData(vertices.data(),
                        vertices.size() * sizeof(StaticMeshVertex),
                        static_cast<Core::uint32>(vertices.size()),
                        sizeof(StaticMeshVertex),
                        attributes);

    const bool use32BitIndex = vertices.size() > std::numeric_limits<Core::uint16>::max();
    if (use32BitIndex) {
        mesh->SetIndexData(indices.data(),
                           indices.size() * sizeof(Core::uint32),
                           static_cast<Core::uint32>(indices.size()),
                           true);
    } else {
        std::vector<Core::uint16> compactIndices;
        compactIndices.reserve(indices.size());
        for (Core::uint32 index : indices) {
            compactIndices.push_back(static_cast<Core::uint16>(index));
        }
        mesh->SetIndexData(compactIndices.data(),
                           compactIndices.size() * sizeof(Core::uint16),
                           static_cast<Core::uint32>(compactIndices.size()),
                           false);
    }

    for (const MeshSection& section : sections) {
        mesh->AddSection(section);
    }
    
    return LoadResult(mesh);
}

ImportSettings* MeshLoader::GetDefaultSettings() const {
    return new MeshImportSettings();
}

REGISTER_RESOURCE_LOADER(MeshLoader);

} // namespace Resources
} // namespace XCEngine