Mesh.cpp

#include <Phoenix/renderer/Mesh.h>
#include <Phoenix/core/log.h>

#include <limits>
#include <assimp/Importer.hpp>
#include <assimp/scene.h>
#include <assimp/postprocess.h>
#include <assimp/config.h>

#include <glm/gtc/matrix_transform.hpp>
#include <chrono>
#include <algorithm>
#include <cctype>
#include <filesystem>

namespace Phoenix{

    Mesh::Mesh(const std::vector<Vertex>& vertices, const std::vector<uint32_t>& indices){
        m_IndexCount = (uint32_t)indices.size();

        // Retain positions + indices on the CPU for physics collider generation.
        m_Positions.reserve(vertices.size());
        for (const auto& v : vertices) { m_Positions.push_back(v.Position); }
        m_Indices = indices;

        m_VertexArray = CreateRef<VertexArray>();
        m_VertexArray->Bind();

        Ref<VertexBuffer> vertexBuffer = CreateRef<VertexBuffer>(
            (float*)vertices.data(), (int)(vertices.size() * sizeof(Vertex)));
        // Bone indices/weights are part of every vertex (defaulted for static meshes)
        // so the stride matches sizeof(Vertex) and the skinning shader can read them.
        BufferLayout layout = {
            { ShaderDataType::Float3, "a_Position" },
            { ShaderDataType::Float3, "a_Normal" },
            { ShaderDataType::Float2, "a_TexCoords" },
            { ShaderDataType::Float4, "a_BoneIDs" },
            { ShaderDataType::Float4, "a_Weights" }
        };
        vertexBuffer->SetLayout(layout);
        m_VertexArray->AddVertexBuffer(vertexBuffer);

        Ref<IndexBuffer> indexBuffer = CreateRef<IndexBuffer>(
            (uint32_t*)indices.data(), (uint32_t)indices.size());
        m_VertexArray->SetIndexBuffer(indexBuffer);
    }

    // ---- Assimp -> glm conversions ----

    static glm::mat4 ToGlm(const aiMatrix4x4& m){
        glm::mat4 r;
        r[0][0] = m.a1; r[1][0] = m.a2; r[2][0] = m.a3; r[3][0] = m.a4;
        r[0][1] = m.b1; r[1][1] = m.b2; r[2][1] = m.b3; r[3][1] = m.b4;
        r[0][2] = m.c1; r[1][2] = m.c2; r[2][2] = m.c3; r[3][2] = m.c4;
        r[0][3] = m.d1; r[1][3] = m.d2; r[2][3] = m.d3; r[3][3] = m.d4;
        return r;
    }
    static glm::vec3 ToGlm(const aiVector3D& v){ return { v.x, v.y, v.z }; }
    static glm::quat ToGlm(const aiQuaternion& q){ return glm::quat(q.w, q.x, q.y, q.z); }

    // ---- Worker-thread parsing (no GL calls) ----

    static void SetVertexBoneData(Vertex& v, int boneID, float weight){
        if (weight <= 0.0f) { return; }
        for (int i = 0; i < 4; i++){
            if (v.BoneIDs[i] < 0.0f){
                v.BoneIDs[i] = (float)boneID;
                v.Weights[i] = weight;
                return;
            }
        }
    }

    // Read this mesh's bones, assigning each a stable id in the shared boneInfoMap,
    // and distribute the per-vertex weights.
    static void ExtractBoneWeights(std::vector<Vertex>& vertices, aiMesh* mesh,
            std::map<std::string, BoneInfo>& boneInfoMap, int& boneCounter){
        for (unsigned int i = 0; i < mesh->mNumBones; i++){
            aiBone* bone = mesh->mBones[i];
            std::string name = bone->mName.C_Str();
            int boneID;
            auto it = boneInfoMap.find(name);
            if (it == boneInfoMap.end()){
                BoneInfo info;
                info.id = boneCounter;
                info.offset = ToGlm(bone->mOffsetMatrix);
                boneInfoMap[name] = info;
                boneID = boneCounter++;
            }
            else{
                boneID = it->second.id;
            }
            for (unsigned int w = 0; w < bone->mNumWeights; w++){
                unsigned int vid = bone->mWeights[w].mVertexId;
                if (vid < vertices.size())
                    SetVertexBoneData(vertices[vid], boneID, bone->mWeights[w].mWeight);
            }
        }
    }

    // "<dir>/bricks_n.jpg" -> "<dir>/bricks_n_height.png" (or .jpg) if such a file exists.
    // Formats like glTF have no height-map slot, so parallax maps sit beside the normal map.
    static std::string FindHeightBeside(const std::string& normalPath){
        std::filesystem::path p(normalPath);
        for (const char* ext : { ".png", ".jpg", ".jpeg" }){
            std::filesystem::path candidate = p.parent_path() / (p.stem().string() + "_height" + ext);
            std::error_code ec;
            if (std::filesystem::exists(candidate, ec)) { return candidate.string(); }
        }
        return {};
    }

    static void ProcessMeshData(aiMesh* mesh, const aiScene* scene, const std::string& directory,
            std::vector<MeshData>& out, std::map<std::string, BoneInfo>& boneInfoMap, int& boneCounter){
        MeshData data;
        data.vertices.reserve(mesh->mNumVertices);

        for (unsigned int i = 0; i < mesh->mNumVertices; i++){
            Vertex vertex;
            vertex.Position = { mesh->mVertices[i].x, mesh->mVertices[i].y, mesh->mVertices[i].z };

            if (mesh->HasNormals())
                vertex.Normal = { mesh->mNormals[i].x, mesh->mNormals[i].y, mesh->mNormals[i].z };
            else
                vertex.Normal = { 0.0f, 0.0f, 0.0f };

            if (mesh->mTextureCoords[0])
                vertex.TexCoords = { mesh->mTextureCoords[0][i].x, mesh->mTextureCoords[0][i].y };
            else
                vertex.TexCoords = { 0.0f, 0.0f };

            data.vertices.push_back(vertex);
        }

        for (unsigned int i = 0; i < mesh->mNumFaces; i++){
            const aiFace& face = mesh->mFaces[i];
            for (unsigned int j = 0; j < face.mNumIndices; j++)
                data.indices.push_back(face.mIndices[j]);
        }

        ExtractBoneWeights(data.vertices, mesh, boneInfoMap, boneCounter);

        if (mesh->mMaterialIndex >= 0){
            aiMaterial* material = scene->mMaterials[mesh->mMaterialIndex];
            if (material->GetTextureCount(aiTextureType_DIFFUSE) > 0){
                aiString name;
                material->GetTexture(aiTextureType_DIFFUSE, 0, &name);
                data.diffusePath = directory + "/" + std::string(name.C_Str());
            }
            // Only explicit normal maps (glTF normalTexture, OBJ "norm"): OBJ map_Bump
            // usually holds a height map, which isn't a tangent-space normal.
            if (material->GetTextureCount(aiTextureType_NORMALS) > 0){
                aiString name;
                material->GetTexture(aiTextureType_NORMALS, 0, &name);
                data.normalPath = directory + "/" + std::string(name.C_Str());
            }
            // Parallax height map: OBJ "disp", else a "<normal map>_height" file beside it.
            if (material->GetTextureCount(aiTextureType_DISPLACEMENT) > 0){
                aiString name;
                material->GetTexture(aiTextureType_DISPLACEMENT, 0, &name);
                data.heightPath = directory + "/" + std::string(name.C_Str());
            }
            else if (!data.normalPath.empty()){
                data.heightPath = FindHeightBeside(data.normalPath);
            }
            // glTF cut-out materials (leaves, chains): alpha-test against the cutoff.
            aiString alphaMode;
            if (material->Get("$mat.gltf.alphaMode", 0, 0, alphaMode) == AI_SUCCESS
                && std::string(alphaMode.C_Str()) == "MASK"){
                float cutoff = 0.5f;
                material->Get("$mat.gltf.alphaCutoff", 0, 0, cutoff);
                data.alphaCutoff = std::max(cutoff, 1e-3f);
            }
        }

        out.push_back(std::move(data));
    }

    static void ProcessNodeData(aiNode* node, const aiScene* scene, const std::string& directory,
            std::vector<MeshData>& out, std::map<std::string, BoneInfo>& boneInfoMap, int& boneCounter){
        for (unsigned int i = 0; i < node->mNumMeshes; i++)
            ProcessMeshData(scene->mMeshes[node->mMeshes[i]], scene, directory, out, boneInfoMap, boneCounter);

        for (unsigned int i = 0; i < node->mNumChildren; i++)
            ProcessNodeData(node->mChildren[i], scene, directory, out, boneInfoMap, boneCounter);
    }

    static void ReadHierarchy(AssimpNodeData& dest, const aiNode* src){
        dest.name = src->mName.C_Str();
        dest.transformation = ToGlm(src->mTransformation);
        dest.children.reserve(src->mNumChildren);
        for (unsigned int i = 0; i < src->mNumChildren; i++){
            AssimpNodeData child;
            ReadHierarchy(child, src->mChildren[i]);
            dest.children.push_back(std::move(child));
        }
    }

    static Ref<Animation> LoadAnimation(const aiAnimation* anim, const aiScene* scene,
            std::map<std::string, BoneInfo>& boneInfoMap, int& boneCounter, const glm::mat4& globalInverse){
        auto out = CreateRef<Animation>();
        out->m_Name = anim->mName.C_Str();
        out->m_Duration = (float)anim->mDuration;
        out->m_TicksPerSecond = anim->mTicksPerSecond != 0.0 ? (float)anim->mTicksPerSecond : 25.0f;
        out->m_GlobalInverseTransform = globalInverse;
        ReadHierarchy(out->m_RootNode, scene->mRootNode);

        for (unsigned int i = 0; i < anim->mNumChannels; i++){
            aiNodeAnim* channel = anim->mChannels[i];
            std::string boneName = channel->mNodeName.C_Str();

            // A channel may target a node that wasn't a skinning bone in any mesh; give
            // it an id so the hierarchy walk still finds its animated local transform.
            if (boneInfoMap.find(boneName) == boneInfoMap.end())
                boneInfoMap[boneName].id = boneCounter++;
            int id = boneInfoMap[boneName].id;

            std::vector<KeyPosition> positions;
            positions.reserve(channel->mNumPositionKeys);
            for (unsigned int k = 0; k < channel->mNumPositionKeys; k++)
                positions.push_back({ ToGlm(channel->mPositionKeys[k].mValue), (float)channel->mPositionKeys[k].mTime });

            std::vector<KeyRotation> rotations;
            rotations.reserve(channel->mNumRotationKeys);
            for (unsigned int k = 0; k < channel->mNumRotationKeys; k++)
                rotations.push_back({ ToGlm(channel->mRotationKeys[k].mValue), (float)channel->mRotationKeys[k].mTime });

            std::vector<KeyScale> scales;
            scales.reserve(channel->mNumScalingKeys);
            for (unsigned int k = 0; k < channel->mNumScalingKeys; k++)
                scales.push_back({ ToGlm(channel->mScalingKeys[k].mValue), (float)channel->mScalingKeys[k].mTime });

            out->m_Bones.emplace_back(boneName, id, std::move(positions), std::move(rotations), std::move(scales));
        }

        out->m_BoneInfoMap = boneInfoMap; // snapshot (meshes' bones + channels' bones)
        out->BuildNodeCache();
        return out;
    }

    static ModelData ParseModel(std::string path){
        ModelData result;

        Assimp::Importer importer;
        // Collapse FBX pivot helpers ("$AssimpFbx$" nodes) so each bone is a single node
        // with one full translation/rotation/scale channel -- which the keyframe sampler
        // expects. Without this, FBX (e.g. Mixamo) animations are scattered across
        // synthetic single-purpose nodes and reconstruct incorrectly (splayed limbs).
        importer.SetPropertyBool(AI_CONFIG_IMPORT_FBX_PRESERVE_PIVOTS, false);
        // Assimp's glTF importer already flips V to the bottom-left origin that matches
        // our flipped-on-load textures; flipping again would turn glTF textures upside down.
        std::string ext = path.substr(path.find_last_of('.') + 1);
        std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c){ return (char)std::tolower(c); });
        const bool gltf = ext == "gltf" || ext == "glb";
        const aiScene* scene = importer.ReadFile(path,
            aiProcess_Triangulate
            | aiProcess_GenSmoothNormals
            | (gltf ? 0 : aiProcess_FlipUVs)
            | aiProcess_JoinIdenticalVertices
            | aiProcess_LimitBoneWeights);

        if (!scene || (scene->mFlags & AI_SCENE_FLAGS_INCOMPLETE) || !scene->mRootNode){
            PHX_CORE_ERROR("Assimp failed to load '{0}': {1}", path, importer.GetErrorString());
            return result;
        }

        std::string directory = path.substr(0, path.find_last_of('/'));
        int boneCounter = 0;
        ProcessNodeData(scene->mRootNode, scene, directory, result.meshes, result.boneInfoMap, boneCounter);

        glm::mat4 globalInverse = glm::inverse(ToGlm(scene->mRootNode->mTransformation));
        for (unsigned int i = 0; i < scene->mNumAnimations; i++)
            result.animations.push_back(LoadAnimation(scene->mAnimations[i], scene, result.boneInfoMap, boneCounter, globalInverse));

        return result;
    }

    // ---- Main-thread model handle ----

    Model::Model(const std::string& path) : m_Path(path){
        m_Future = std::async(std::launch::async, ParseModel, path);
    }

    Model::~Model() = default;

    void Model::Update(){
        if (m_Uploaded || !m_Future.valid())
            return;
        if (m_Future.wait_for(std::chrono::seconds(0)) != std::future_status::ready)
            return;

        ModelData data = m_Future.get();
        m_BoneInfoMap = std::move(data.boneInfoMap);
        m_Animations = std::move(data.animations);

        for (auto& md : data.meshes){
            Ref<Mesh> mesh = CreateRef<Mesh>(md.vertices, md.indices);
            if (!md.diffusePath.empty())
                mesh->SetDiffuseMap(Texture2D::Create(md.diffusePath));
            if (!md.normalPath.empty())
                mesh->SetNormalMap(Texture2D::Create(md.normalPath));
            if (!md.heightPath.empty())
                mesh->SetHeightMap(Texture2D::Create(md.heightPath));
            mesh->SetAlphaCutoff(md.alphaCutoff);
            m_Meshes.push_back(mesh);
        }

        m_Uploaded = true;
        PHX_CORE_INFO("Loaded model '{0}' ({1} meshes, {2} bones, {3} animations)",
            m_Path, m_Meshes.size(), m_BoneInfoMap.size(), m_Animations.size());
        if (m_BoneInfoMap.size() > (size_t)MAX_BONES)
            PHX_CORE_ERROR("Model '{0}' has {1} bones, exceeding MAX_BONES={2}; raise it and u_BoneMatrices[] in the shaders",
                m_Path, m_BoneInfoMap.size(), MAX_BONES);
    }

    const std::string& Model::GetAnimationName(size_t index) const{
        static const std::string empty;
        return index < m_Animations.size() ? m_Animations[index]->GetName() : empty;
    }

    bool Model::GetLocalBounds(glm::vec3& outMin, glm::vec3& outMax) const{
        if (!m_Uploaded) { return false; }
        if (!m_BoundsValid){
            glm::vec3 lo(std::numeric_limits<float>::max()), hi(-std::numeric_limits<float>::max());
            for (const auto& mesh : m_Meshes)
                for (const glm::vec3& p : mesh->GetPositions()){ lo = glm::min(lo, p); hi = glm::max(hi, p); }
            if (lo.x > hi.x) { return false; } // no vertices
            m_BoundsMin = lo;
            m_BoundsMax = hi;
            m_BoundsValid = true;
        }
        outMin = m_BoundsMin;
        outMax = m_BoundsMax;
        return true;
    }

    void Model::RemoveRootMotion(){
        for (std::size_t i = 0; i < m_Animations.size(); i++){
            glm::vec3 d = m_Animations[i]->RemoveRootDrift();
            if (glm::length(d) > 1e-3f)
                PHX_CORE_INFO("In place: removed root travel ({0:.1f}, {1:.1f}, {2:.1f}) from clip '{3}'",
                              d.x, d.y, d.z, GetAnimationName(i));
        }
    }

    bool Model::FindBone(const std::string& name, BoneInfo& out) const{
        if (!m_Uploaded || name.empty()) { return false; }
        auto it = m_BoneInfoMap.find(name);
        if (it == m_BoneInfoMap.end()){
            const std::string suffix = ":" + name;
            for (auto i = m_BoneInfoMap.begin(); i != m_BoneInfoMap.end(); ++i){
                const std::string& n = i->first;
                if (n.size() > suffix.size() && n.compare(n.size() - suffix.size(), suffix.size(), suffix) == 0){
                    it = i;
                    break;
                }
            }
        }
        if (it == m_BoneInfoMap.end()) { return false; }
        out = it->second;
        return true;
    }

    int Model::GetAnimationIndex(const std::string& name) const{
        for (size_t i = 0; i < m_Animations.size(); i++)
            if (m_Animations[i]->GetName() == name) { return (int)i; }
        return -1;
    }

    void Model::AddAnimationsFromFile(const std::string& path){
        // Undo/redo and Stop re-apply the scene, which asks for the same files again;
        // re-importing would be slow and append duplicate clips.
        if (m_MergedAnimationFiles.count(path)) { return; }

        Assimp::Importer importer;
        importer.SetPropertyBool(AI_CONFIG_IMPORT_FBX_PRESERVE_PIVOTS, false);
        const aiScene* scene = importer.ReadFile(path,
            aiProcess_Triangulate | aiProcess_LimitBoneWeights);

        if (!scene || !scene->mRootNode || scene->mNumAnimations == 0){
            // Not fatal: the file may simply be absent (the user hasn't added it yet).
            PHX_CORE_WARN("No animations merged from '{0}': {1}", path, importer.GetErrorString());
            return;
        }

        // Clip name = the file's base name (e.g. "idle.fbx" -> "idle"), because Mixamo
        // names every clip "mixamo.com" so the embedded name can't distinguish them.
        size_t slash = path.find_last_of("/\\");
        size_t dot = path.find_last_of('.');
        std::string base = path.substr(slash == std::string::npos ? 0 : slash + 1);
        if (dot != std::string::npos && dot > slash) { base = base.substr(0, base.find_last_of('.')); }

        // Start any channel-only nodes after the model's existing bone ids so skinning
        // bone ids/offsets (seeded from m_BoneInfoMap) are preserved for the merged clips.
        int baseMax = 0;
        for (const auto& kv : m_BoneInfoMap) { baseMax = std::max(baseMax, kv.second.id + 1); }

        glm::mat4 globalInverse = glm::inverse(ToGlm(scene->mRootNode->mTransformation));
        size_t before = m_Animations.size();
        for (unsigned int i = 0; i < scene->mNumAnimations; i++){
            std::map<std::string, BoneInfo> seed = m_BoneInfoMap; // base ids + offsets
            int boneCounter = baseMax;
            Ref<Animation> a = LoadAnimation(scene->mAnimations[i], scene, seed, boneCounter, globalInverse);
            a->m_Name = (scene->mNumAnimations == 1) ? base : base + "_" + std::to_string(i);
            m_Animations.push_back(a);
        }
        m_MergedAnimationFiles.insert(path);
        PHX_CORE_INFO("Merged {0} animation(s) named '{1}' from '{2}' (model now has {3} clips)",
            m_Animations.size() - before, base, path, m_Animations.size());
    }
}