Renderer/src/main/java/electrosphere/engine/assetmanager/AssetManager.java

package electrosphere.engine.assetmanager;

import electrosphere.audio.AudioBuffer;
import electrosphere.collision.CollisionBodyCreation;
import electrosphere.collision.CollisionEngine;
import electrosphere.collision.collidable.Collidable;
import electrosphere.engine.Globals;
import electrosphere.engine.assetmanager.queue.QueuedAsset;
import electrosphere.logger.LoggerInterface;
import electrosphere.renderer.actor.ActorShaderMask;
import electrosphere.renderer.buffer.HomogenousInstancedArray;
import electrosphere.renderer.buffer.HomogenousUniformBuffer;
import electrosphere.renderer.loading.ModelLoader;
import electrosphere.renderer.model.Mesh;
import electrosphere.renderer.model.Model;
import electrosphere.renderer.shader.ComputeShader;
import electrosphere.renderer.shader.VisualShader;
import electrosphere.renderer.texture.Texture;
import electrosphere.renderer.texture.TextureMap;
import electrosphere.server.poseactor.PoseModel;
import electrosphere.util.FileUtils;

import java.io.File;
import java.io.IOException;
import java.util.Collection;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.locks.ReentrantLock;

import org.lwjgl.assimp.AIScene;
import org.ode4j.ode.DBody;

/**
 * Manages all assets loaded into the engine including initially loading and destructing
 */
public class AssetManager {
    
    Map<String,Model> modelsLoadedIntoMemory = new ConcurrentHashMap<String,Model>();
    List<String> modelsInQueue = new CopyOnWriteArrayList<String>();
    List<String> modelsInDeleteQueue = new CopyOnWriteArrayList<String>();
    List<MeshShaderOverride> shaderOverrides = new CopyOnWriteArrayList<MeshShaderOverride>();
    
    Map<String,Texture> texturesLoadedIntoMemory = new ConcurrentHashMap<String,Texture>();
    List<String> texturesInQueue = new CopyOnWriteArrayList<String>();
    List<String> texturesInDeleteQueue = new CopyOnWriteArrayList<String>();
    
    Map<String,AudioBuffer> audioLoadedIntoMemory = new ConcurrentHashMap<String,AudioBuffer>();
    List<String> audioInQueue = new CopyOnWriteArrayList<String>();

    Map<String,VisualShader> shadersLoadedIntoMemory = new ConcurrentHashMap<String,VisualShader>();
    List<ActorShaderMask> shadersInQueue = new CopyOnWriteArrayList<ActorShaderMask>();

    //
    //Compute shader related
    //
    Map<String,ComputeShader> computeShadersLoadedIntoMemory = new ConcurrentHashMap<String,ComputeShader>();
    List<String> computeShadersInQueue = new CopyOnWriteArrayList<String>();

    Map<String,DBody> physicsMeshesLoadedIntoMemory = new ConcurrentHashMap<String,DBody>();
    List<PhysicsMeshQueueItem> physicsMeshesToLoad = new CopyOnWriteArrayList<PhysicsMeshQueueItem>();

    Map<String,PoseModel> poseModelsLoadedIntoMemory = new ConcurrentHashMap<String,PoseModel>();
    List<String> poseModelsInQueue = new CopyOnWriteArrayList<String>();
    List<String> poseModelsInDeleteQueue = new CopyOnWriteArrayList<String>();

    //A queue of homogenous buffers to allocate this render frame
    List<HomogenousUniformBuffer> homogenousBufferAllocationQueue = new CopyOnWriteArrayList<HomogenousUniformBuffer>();

    //A queue of homogenous buffers to allocate this render frame
    List<HomogenousInstancedArray> instanceArrayBufferAllocationQueue = new CopyOnWriteArrayList<HomogenousInstancedArray>();


    //assets queued to be loaded
    ReentrantLock queuedAssetLock = new ReentrantLock();
    List<QueuedAsset<?>> queuedAssets = new LinkedList<QueuedAsset<?>>();
    
    
    
    
    //
    //General asset manager stuff
    //
    
    /**
     * For each class/type of asset, load all assets in queue
     */
    public void loadAssetsInQueue(){
        //models
        LoggerInterface.loggerEngine.DEBUG_LOOP("AssetManager - Load models");
        for(String currentPath : modelsInQueue){
            modelsInQueue.remove(currentPath);
            AIScene aiScene = ModelLoader.loadAIScene(currentPath);
            TextureMap textureMap = null;
            if(this.getLocalTextureMapPath(currentPath) != null){
                textureMap = TextureMap.construct(this.getLocalTextureMapPath(currentPath));
            }
            if(aiScene != null){
                modelsLoadedIntoMemory.put(currentPath, ModelLoader.createModelFromAiScene(aiScene,textureMap,currentPath));
                for(PhysicsMeshQueueItem physicsMeshQueueItem : physicsMeshesToLoad){
                    if(physicsMeshQueueItem.modelPath.contains(currentPath)){
                        //create physics
                        physicsMeshesToLoad.remove(physicsMeshQueueItem);
                        physicsMeshesLoadedIntoMemory.put(
                            this.getCollisionMeshMapKey(physicsMeshQueueItem.collisionEngine,currentPath),
                            CollisionBodyCreation.generateRigidBodyFromAIScene(physicsMeshQueueItem.collisionEngine,aiScene,Collidable.TYPE_STATIC_BIT)
                        );
                    }
                }
            }
        }
        //textures from disk to gpu
        LoggerInterface.loggerEngine.DEBUG_LOOP("AssetManager - Load textures");
        for(String currentPath : texturesInQueue){
            texturesInQueue.remove(currentPath);
            texturesLoadedIntoMemory.put(currentPath, new Texture(Globals.renderingEngine.getOpenGLState(), currentPath));
        }
        //audio from disk
        LoggerInterface.loggerEngine.DEBUG_LOOP("AssetManager - Load audio");
        if(Globals.audioEngine != null && Globals.audioEngine.initialized()){
            for(String currentPath : audioInQueue){
                audioInQueue.remove(currentPath);
                audioLoadedIntoMemory.put(currentPath, new AudioBuffer(currentPath));
            }
        }
        //shaders
        LoggerInterface.loggerEngine.DEBUG_LOOP("AssetManager - Load shaders");
        for(ActorShaderMask currentShader : shadersInQueue){
            shadersInQueue.remove(currentShader);
            String key = getShaderKey(currentShader.getVertexShaderPath(),currentShader.getFragmentShaderPath());
            shadersLoadedIntoMemory.put(
                key,
                VisualShader.loadSpecificShader(currentShader.getVertexShaderPath(),currentShader.getFragmentShaderPath())
            );
        }
        //compute shaders
        LoggerInterface.loggerEngine.DEBUG_LOOP("AssetManager - Load compute shaders");
        for(String computePath : computeShadersInQueue){
            computeShadersInQueue.remove(computePath);
            String key = getComputeShaderKey(computePath);
            try {
                computeShadersLoadedIntoMemory.put(
                    key, 
                    ComputeShader.create(FileUtils.getAssetFileAsString(computePath))
                );
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        //pose models
        LoggerInterface.loggerEngine.DEBUG_LOOP("AssetManager - Load pose models");
        for(String currentPath: poseModelsInQueue){
            poseModelsInQueue.remove(currentPath);
            AIScene scene = ModelLoader.loadAIScene(currentPath);
            poseModelsLoadedIntoMemory.put(currentPath, new PoseModel(currentPath, scene));
        }
        //queued assets
        LoggerInterface.loggerEngine.DEBUG_LOOP("AssetManager - Load queued assets");
        queuedAssetLock.lock();
        for(QueuedAsset<?> queuedAsset : queuedAssets){
            queuedAsset.load();
            if(queuedAsset.get() instanceof Model){
                this.modelsLoadedIntoMemory.put(queuedAsset.getPromisedPath(),(Model)queuedAsset.get());
            } else if(queuedAsset.get() instanceof Texture){
                this.texturesLoadedIntoMemory.put(queuedAsset.getPromisedPath(),(Texture)queuedAsset.get());
            }
        }
        queuedAssets.clear();
        queuedAssetLock.unlock();

        //allocate homogenous buffers
        LoggerInterface.loggerEngine.DEBUG_LOOP("AssetManager - Allocate homogenous buffers");
        this.allocateHomogenousBuffers();

        //allocate instance array buffers
        LoggerInterface.loggerEngine.DEBUG_LOOP("AssetManager - Allocate instance array buffers");
        this.allocateInstanceArrayBuffers();

        //override meshes
        LoggerInterface.loggerEngine.DEBUG_LOOP("AssetManager - Override meshes");
        this.performMeshOverrides();
    }


    //
    //Updating assets
    //
    public void updateAsset(String path){
        LoggerInterface.loggerEngine.DEBUG("AssetManager - updateAsset");
        //models
        if(modelsLoadedIntoMemory.containsKey(path)){
            this.queueModelForDeletion(path);
            this.deleteModelsInDeleteQueue();
            this.addModelPathToQueue(path);
        }
        //textures
        if(texturesLoadedIntoMemory.containsKey(path)){
            this.queueTextureForDeletion(path);
            this.deleteTexturesInDeleteQueue();
            this.addTexturePathtoQueue(path);
        }
        //pose models
        if(poseModelsLoadedIntoMemory.containsKey(path)){
            this.queuePoseModelForDeletion(path);
            this.deletePoseModelsInDeleteQueue();
            this.addPoseModelPathToQueue(path);
        }
        if(audioLoadedIntoMemory.containsKey(path)){
            throw new Error("Unhandled asset type! (Audio)");
        }
        if(shadersLoadedIntoMemory.containsKey(path)){
            throw new Error("Unhandled asset type! (Shader - Visual)");
        }
        if(computeShadersLoadedIntoMemory.containsKey(path)){
            throw new Error("Unhandled asset type! (Shader - Compute)");
        }
    }

    /**
     * Handles the delete queues
     */
    public void handleDeleteQueue(){
        this.deleteModelsInDeleteQueue();
        this.deletePoseModelsInDeleteQueue();
        this.deleteTexturesInDeleteQueue();
    }
    
    
    
    
    
    
    
    
    //
    //Models
    //
    
    public void addModelPathToQueue(String path){
        if(!modelsInQueue.contains(path) && !modelsLoadedIntoMemory.containsKey(path)){
            modelsInQueue.add(path);
        }
    }
    
    public Model fetchModel(String path){
        Model rVal = null;
        if(modelsLoadedIntoMemory.containsKey(path)){
            rVal = modelsLoadedIntoMemory.get(path);
        }
        return rVal;
    }

    /**
     * Queues a model for deletion
     * @param modelPath The path to the model
     */
    public void queueModelForDeletion(String modelPath){
        modelsInDeleteQueue.add(modelPath);
    }
    
    /**
    Registers a (presumably generated in code) model with the asset manager
    @returns a random string that represents the model in the asset manager
    */
    public String registerModel(Model m){
        String rVal;
        UUID newUUID = UUID.randomUUID();
        rVal = newUUID.toString();
        modelsLoadedIntoMemory.put(rVal,m);
        return rVal;
    }

    /**
     * Registers a (presumably generated in code) model to a given path in the asset manager
     * Used particularly if you have a specific path you want to relate to a specific model (eg terrain chunk code)
     * @param m The model to register
     * @param path The path to register the model to
     */
    public void registerModelWithPath(Model m, String path){
        modelsLoadedIntoMemory.put(path, m);
    }

    public void deregisterModelPath(String path){
        modelsLoadedIntoMemory.remove(path);
    }
    
    public void queueOverrideMeshShader(String modelName, String meshName, String vertPath, String fragPath){
        MeshShaderOverride override = new MeshShaderOverride(modelName,meshName,vertPath,fragPath);
        shaderOverrides.add(override);
    }

    public void performMeshOverrides(){
        List<MeshShaderOverride> toRemove = new LinkedList<MeshShaderOverride>();
        for(MeshShaderOverride shaderOverride : shaderOverrides){
            Model model = null;
            if((model = fetchModel(shaderOverride.modelName)) != null){
                for(Mesh mesh : model.getMeshes()){
                    if(mesh.getMeshName().equals(shaderOverride.getMeshName())){
                        mesh.setShader(VisualShader.loadSpecificShader(shaderOverride.vertPath, shaderOverride.fragPath));
                    }
                }
                toRemove.add(shaderOverride);
            }
        }
        for(MeshShaderOverride shaderOverride : toRemove){
            shaderOverrides.remove(shaderOverride);
        }
    }

    /**
     * Nuclear function, reloads all shaders loaded into memory
     */
    public void forceReloadAllModels(){
        for(String modelKey : modelsLoadedIntoMemory.keySet()){
            if(modelKey.contains("Models")){
                modelsInQueue.add(modelKey);
                modelsLoadedIntoMemory.remove(modelKey);
            }
        }
        // modelsLoadedIntoMemory.clear();
    }

    /**
     * Deletes all models in the delete queue
     */
    public void deleteModelsInDeleteQueue(){
        for(String modelPath : modelsInDeleteQueue){
            Model model = this.fetchModel(modelPath);
            if(model != null){
                model.delete();
            }
            this.modelsLoadedIntoMemory.remove(modelPath);
        }
    }

    
    
    
    



    //
    //Pose Models
    //
    /**
     * Adds a pose model to the list of pose models to load
     * @param path The path to load
     */
    public void addPoseModelPathToQueue(String path){
        if(!poseModelsInQueue.contains(path) && !poseModelsLoadedIntoMemory.containsKey(path)){
            poseModelsInQueue.add(path);
        }
    }
    
    /**
     * Fetches a pose model
     * @param path The path to fetch
     * @return The pose model if it exists, null otherwise
     */
    public PoseModel fetchPoseModel(String path){
        PoseModel rVal = null;
        if(poseModelsLoadedIntoMemory.containsKey(path)){
            rVal = poseModelsLoadedIntoMemory.get(path);
        }
        return rVal;
    }

    /**
     * Registers a (presumably generated in code) pose model to a given path in the asset manager
     * Used particularly if you have a specific path you want to relate to a specific pose model (eg basic shapes)
     * @param m The pose model to register
     * @param path The path to register the pose model to
     */
    public void registerPoseModelWithPath(PoseModel m, String path){
        poseModelsLoadedIntoMemory.put(path, m);
    }

    /**
     * Queues a pose model for deletion
     * @param modelPath The path to the pose model
     */
    public void queuePoseModelForDeletion(String modelPath){
        poseModelsInDeleteQueue.add(modelPath);
    }

    /**
     * Deletes all pose models in the delete queue
     */
    public void deletePoseModelsInDeleteQueue(){
        for(String modelPath : poseModelsInDeleteQueue){
            PoseModel poseModel = this.fetchPoseModel(modelPath);
            if(poseModel != null){
                poseModel.delete();
            }
            this.poseModelsLoadedIntoMemory.remove(modelPath);
        }
    }
    
    
    
    
    
    
    
    
    
    //
    // Textures
    //
    
    
    public void addTexturePathtoQueue(String path){
        if(!texturesInQueue.contains(path) && !texturesLoadedIntoMemory.containsKey(path)){
            texturesInQueue.add(path);
        }
    }

    /**
     * Queues a texture for deletion
     * @param texturePath The path to the texture
     */
    public void queueTextureForDeletion(String texturePath){
        texturesInDeleteQueue.add(texturePath);
    }
    
    public Texture fetchTexture(String path){
        Texture rVal = null;
        if(texturesLoadedIntoMemory.containsKey(path)){
            rVal = texturesLoadedIntoMemory.get(path);
        }
        return rVal;
    }
    
    public String registerTexture(Texture t){
        String rVal;
        UUID newUUID = UUID.randomUUID();
        rVal = newUUID.toString();
        texturesLoadedIntoMemory.put(rVal,t);
        return rVal;
    }

    /**
     * Registers a texture to a path
     * @param t The texture
     * @param path The path
     */
    public void registerTextureToPath(Texture t, String path){
        texturesLoadedIntoMemory.put(path,t);
    }

    public boolean hasLoadedTexture(String path){
        return texturesLoadedIntoMemory.containsKey(path);
    }

    /**
     * Gets a local texture map's path from a model's path
     * @param modelPath The model's path
     */
    private String getLocalTextureMapPath(String modelPath){
        File modelFile = FileUtils.getAssetFile(modelPath);
        File containingDirectory = modelFile.getParentFile();
        File[] children = containingDirectory.listFiles();
        if(children != null){
            for(File child : children){
                if(child.getName().equals("texturemap.json")){
                    String rVal = child.getPath();
                    String fixed = rVal.replace(".\\assets", "").replace("./assets", "");
                    return fixed;
                }
            }
        }
        return null;
    }

    /**
     * Deletes all textures in the delete queue
     */
    public void deleteTexturesInDeleteQueue(){
        for(String texturePath : texturesInDeleteQueue){
            Texture texture = this.fetchTexture(texturePath);
            if(texture != null){
                texture.free();
            }
            this.texturesLoadedIntoMemory.remove(texturePath);
        }
    }
    
    
    
    
    
    
    
    
    
    //
    //AUDIO
    //
    
    public void addAudioPathToQueue(String path){
        String sanitizedPath = FileUtils.sanitizeFilePath(path);
        if(!audioInQueue.contains(sanitizedPath) && !audioLoadedIntoMemory.containsKey(sanitizedPath)){
            audioInQueue.add(sanitizedPath);
        }
    }
    
    public AudioBuffer fetchAudio(String path){
        AudioBuffer rVal = null;
        String sanitizedPath = FileUtils.sanitizeFilePath(path);
        if(audioLoadedIntoMemory.containsKey(sanitizedPath)){
            rVal = audioLoadedIntoMemory.get(sanitizedPath);
        } else {
            LoggerInterface.loggerAudio.WARNING("Failed to find audio " + sanitizedPath);
        }
        return rVal;
    }

    /**
     * Gets all audio loaded into the engine
     * @return The collection of all audio buffers
     */
    public Collection<AudioBuffer> getAllAudio(){
        return Collections.unmodifiableCollection(this.audioLoadedIntoMemory.values());
    }
    
    






    //
    //SHADERS
    //
    public void addShaderToQueue(String vertexShader, String fragmentShader){
        shadersInQueue.add(new ActorShaderMask("","",vertexShader,fragmentShader));
    }

    public VisualShader fetchShader(String vertexPath, String fragmentPath){
        String path = getShaderKey(vertexPath,fragmentPath);
        VisualShader rVal = null;
        if(shadersLoadedIntoMemory.containsKey(path)){
            rVal = shadersLoadedIntoMemory.get(path);
        }
        return rVal;
    }

    static String getShaderKey(String vertexPath, String fragmentPath){
        return vertexPath + "-" + fragmentPath;
    }

    /**
     * Nuclear function, reloads all shaders loaded into memory
     */
    public void forceReloadAllShaders(){
        for(String shaderKey : shadersLoadedIntoMemory.keySet()){
            String shaderPaths[] = shaderKey.split("-");
            shadersInQueue.add(new ActorShaderMask("","",shaderPaths[0],shaderPaths[1]));
        }
        shadersLoadedIntoMemory.clear();
    }

    //
    //SHADERS
    //
    /**
     * Adds a compute shader to the queue to be loaded
     * @param computePath The path to the source code for the shader
     */
    public void addComputeShaderToQueue(String computePath){
        computeShadersInQueue.add(getComputeShaderKey(computePath));
    }

    /**
     * Gets a compute shader
     * @param computePath The path to the source code for the compute shader
     * @return The compute shader if it exists, null otherwise
     */
    public ComputeShader fetchComputeShader(String computePath){
        String key = getComputeShaderKey(computePath);
        ComputeShader rVal = null;
        if(computeShadersLoadedIntoMemory.containsKey(key)){
            rVal = computeShadersLoadedIntoMemory.get(key);
        }
        return rVal;
    }

    /**
     * Gets the key for a compute shader
     * @param computePath The path to the shader source
     * @return The key
     */
    static String getComputeShaderKey(String computePath){
        return computePath;
    }
    
    
    
    
    //
    //COLLISION MESH
    //
    public void addCollisionMeshToQueue(PhysicsMeshQueueItem physicsMeshQueueItem){
        if(
            !physicsMeshesToLoad.contains(physicsMeshQueueItem) && 
            !physicsMeshesLoadedIntoMemory.containsKey(getCollisionMeshMapKey(
                physicsMeshQueueItem.collisionEngine,
                physicsMeshQueueItem.modelPath
        ))){
            physicsMeshesToLoad.add(physicsMeshQueueItem);
        }
    }

    public DBody fetchCollisionObject(CollisionEngine collisionEngine, String path){
        return physicsMeshesLoadedIntoMemory.get(getCollisionMeshMapKey(collisionEngine,path));
    }

    /**
     * Gets a key based on collision engine object hash and path of model
     * @param collisionEngine collision engine
     * @param path The path
     * @return The key
     */
    private String getCollisionMeshMapKey(CollisionEngine collisionEngine, String path){
        return collisionEngine + path;
    }


    //
    //HOMOGENOUS UNIFORM BUFFERS
    //
    /**
     * Allocates all uniform buffers in queue
     */
    public void allocateHomogenousBuffers(){
        for(HomogenousUniformBuffer buffer : homogenousBufferAllocationQueue){
            buffer.allocate();
        }
        homogenousBufferAllocationQueue.clear();
    }

    /**
     * Adds a uniform buffer to the queue to be allocated
     * @param buffer The buffer
     */
    public void addHomogenousBufferToQueue(HomogenousUniformBuffer buffer){
        homogenousBufferAllocationQueue.add(buffer);
    }


    //
    //INSTANCE ARRAY BUFFERS
    //
    /**
     * Allocates all instance array buffers in queue
     */
    public void allocateInstanceArrayBuffers(){
        for(HomogenousInstancedArray buffer : instanceArrayBufferAllocationQueue){
            buffer.allocate();
        }
        instanceArrayBufferAllocationQueue.clear();
    }

    /**
     * Adds an instance array buffer to the queue to be allocated
     * @param buffer The buffer
     */
    public void addInstanceArrayBufferToQueue(HomogenousInstancedArray buffer){
        instanceArrayBufferAllocationQueue.add(buffer);
    }


    //
    //Async generic queued assets
    //
    /**
     * Queues an asset to be loaded on the main thread
     * @param asset the asset
     */
    public String queuedAsset(QueuedAsset<?> asset){
        queuedAssetLock.lock();
        this.queuedAssets.add(asset);

        //promise a specific string for this asset
        String promisedPath;
        if(asset.suppliedPath()){
            promisedPath = asset.getPromisedPath();
            if(promisedPath == null || promisedPath == ""){
                String message = "Queued an asset with an empty promised path!" +
                " " + promisedPath +
                " " + asset
                ;
                throw new Error(message);
            }
        } else {
            UUID newUUID = UUID.randomUUID();
            promisedPath = newUUID.toString();
            asset.setPromisedPath(promisedPath);
        }

        queuedAssetLock.unlock();

        return promisedPath;
    }
    
    
    
}