package electrosphere.renderer.model;

import electrosphere.renderer.RenderPipelineState;
import electrosphere.renderer.actor.ActorBoneRotator;
import electrosphere.renderer.actor.ActorMeshMask;
import electrosphere.renderer.actor.ActorShaderMask;
import electrosphere.renderer.actor.ActorStaticMorph;
import electrosphere.renderer.actor.ActorTextureMask;
import electrosphere.renderer.anim.AnimChannel;
import electrosphere.renderer.anim.Animation;
import electrosphere.renderer.loading.ModelPretransforms;
import electrosphere.renderer.meshgen.MeshLoader;
import electrosphere.renderer.shader.ShaderProgram;
import electrosphere.renderer.anim.AnimNode;
import electrosphere.engine.Globals;
import electrosphere.logger.LoggerInterface;

import java.util.ArrayList;
import java.util.List;

import org.joml.Matrix4d;
import org.lwjgl.PointerBuffer;
import org.lwjgl.assimp.AIMaterial;
import org.lwjgl.assimp.AIMesh;
import org.lwjgl.assimp.AIScene;
import org.lwjgl.assimp.Assimp;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import org.joml.Sphered;
import org.joml.Vector3d;
import org.lwjgl.assimp.AIAnimation;
import org.lwjgl.assimp.AIBone;
import org.lwjgl.assimp.AINode;

/**
 * A model
 * Contains a list of meshes to draw
 * Contains animations to apply to those meshes
 */
public class Model {

    //the model matrix for this model
    private Matrix4d modelMatrix = new Matrix4d();

    //an optional global transform applied to the parent bone. Typically found in models loaded from files
    private Matrix4d globalInverseTransform = new Matrix4d();

    //the meshes in the model
    private List<Mesh> meshes = new ArrayList<Mesh>();

    //the materials in the model
    private List<Material> materials = new ArrayList<Material>();

    //bone data for the model
    private List<Bone> bones = new ArrayList<Bone>();
    private Map<String,Bone> boneMap = new HashMap<String,Bone>();
    
    //these structures are used to parse the tree of bones of the model
    private AnimNode rootAnimNode;
    private Map<String,AnimNode> nodeMap = new HashMap<String,AnimNode>();

    //animation data for the model
    private List<Animation> animations = new ArrayList<Animation>();
    private Map<String,Animation> animMap = new HashMap<String,Animation>();


    //these are masks that can be applied to the model to overwrite meshes, shaders, and textures of it
    private ActorMeshMask meshMask;
    private Map<String,ActorShaderMask> shaderMask = new HashMap<String,ActorShaderMask>();
    private Map<String,ActorTextureMask> textureMap = new HashMap<String,ActorTextureMask>();

    //The bounding sphere for this particular model
    private Sphered boundingSphere = new Sphered();
    

    
    
    /**
     * Constructor
     */
    public Model(){
    }

    /**
     * Loads a model from an ai scene object
     * @param path The path of the model
     * @param s the ai scene
     * @return The model object
     */
    public static Model createModelFromAiscene(String path, AIScene s){
        Model rVal = new Model();

        ModelPretransforms.ModelMetadata modelMetadata = Globals.modelPretransforms.getModel(path);
        ModelPretransforms.GlobalTransform globalTransform = null;
        if(modelMetadata != null){
            globalTransform = modelMetadata.getGlobalTransform();
        }
        
        //
        //load meshes
        //
        int meshCount = s.mNumMeshes();
        PointerBuffer meshesBuffer = s.mMeshes();
        rVal.meshes = new ArrayList<Mesh>();
        while(meshesBuffer.hasRemaining()){
            AIMesh aiMesh = AIMesh.create(meshesBuffer.get());
            ModelPretransforms.MeshMetadata meshMetadata = null;
            if(modelMetadata != null){
                meshMetadata = modelMetadata.getMesh(aiMesh.mName().dataString());
            }
            Mesh currentMesh = MeshLoader.createMeshFromAIScene(aiMesh, meshMetadata);
            rVal.meshes.add(currentMesh);
            currentMesh.setParent(rVal);
            //update model bounding sphere
            if(currentMesh.getBoundingSphere().r > rVal.boundingSphere.r){
                rVal.boundingSphere.r = currentMesh.getBoundingSphere().r;
            }
        }
        //
        //register bones
        //
        rVal.bones = new ArrayList<Bone>();
        rVal.boneMap = new HashMap<String, Bone>();
        meshesBuffer.rewind();
        while(meshesBuffer.hasRemaining()){
            AIMesh currentMeshData = AIMesh.createSafe(meshesBuffer.get());
            LoggerInterface.loggerRenderer.DEBUG("mesh name:" + currentMeshData.mName().dataString());
            PointerBuffer boneBuffer = currentMeshData.mBones();
            for(int j = 0; j < currentMeshData.mNumBones(); j++){
                AIBone currentBone = AIBone.createSafe(boneBuffer.get());
                String currentBoneName = currentBone.mName().dataString();
                if(!rVal.boneMap.containsKey(currentBoneName)){
                    Bone boneObject = new Bone(currentBone);
                    rVal.boneMap.put(currentBoneName, boneObject);
                    rVal.bones.add(boneObject);
                }
            }
        }
        Iterator<Mesh> meshIterator = rVal.meshes.iterator();
        rVal.bones = new ArrayList<Bone>();
        rVal.boneMap = new HashMap<String,Bone>();
        rVal.nodeMap = new HashMap<String, AnimNode>();
        while(meshIterator.hasNext()){
            Mesh currentMesh = meshIterator.next();
            Iterator<Bone> boneIterator = currentMesh.getBones().iterator();
            ArrayList<Bone> to_remove_queue = new ArrayList<Bone>();
            ArrayList<Bone> to_add_queue = new ArrayList<Bone>();
            while(boneIterator.hasNext()){
                Bone currentBone = boneIterator.next();
                if(!rVal.boneMap.containsKey(currentBone.boneID)){
                    rVal.bones.add(currentBone);
                    rVal.boneMap.put(currentBone.boneID, currentBone);
                }
            }
            boneIterator = to_remove_queue.iterator();
            while(boneIterator.hasNext()){
                currentMesh.getBones().remove(boneIterator.next());
            }
            boneIterator = to_add_queue.iterator();
            while(boneIterator.hasNext()){
                currentMesh.getBones().add(boneIterator.next());
            }
        }
        //
        //parse animation nodes and form hierarchy
        //
        AINode rootNode = s.mRootNode();
        rVal.globalInverseTransform = electrosphere.util.Utilities.convertAIMatrixd(rootNode.mTransformation());
        if(globalTransform != null){
            rVal.globalInverseTransform.scale(globalTransform.getScale());
        }
        LoggerInterface.loggerRenderer.DEBUG("Global Inverse Transform");
        LoggerInterface.loggerRenderer.DEBUG(rVal.globalInverseTransform + "");
        rVal.rootAnimNode = rVal.buildAnimNodeMap(s.mRootNode(),null);
        //
        //load animations
        //
        int animCount = s.mNumAnimations();
        PointerBuffer animBuffer = s.mAnimations();
        rVal.animations = new ArrayList<Animation>();
        rVal.animMap = new HashMap<String,Animation>();
        for(int i = 0; i < animCount; i++){
            Animation newAnim = new Animation(AIAnimation.create(animBuffer.get(i)), i);
            rVal.animations.add(newAnim);
            rVal.animMap.put(newAnim.name,newAnim);
        }
        //
        //Load materials
        //
        if(s.mNumMaterials() > 0){
            rVal.materials = new ArrayList<Material>();
            PointerBuffer material_buffer = s.mMaterials();
            while(material_buffer.hasRemaining()){
                rVal.materials.add(Material.load_material_from_aimaterial(AIMaterial.create(material_buffer.get())));
            }
        }
        //garbage collect
        System.gc();
        return rVal;
    }
    
    /**
     * Draws the model
     * @param renderPipelineState the render pipeline state
     */
    public void draw(RenderPipelineState renderPipelineState){
        Iterator<Mesh> mesh_Iterator = meshes.iterator();
        while(mesh_Iterator.hasNext()){
            Mesh currentMesh = mesh_Iterator.next();
            if(meshMask == null || (meshMask != null && !meshMask.isBlockedMesh(currentMesh.getMeshName()))){
                //set shader
                ShaderProgram original = currentMesh.getShader();
                ShaderProgram shader = getCorrectShader(shaderMask, currentMesh, currentMesh.getShader());
                currentMesh.setShader(shader);
                //set texture mask
                if(this.textureMap != null && textureMap.containsKey(currentMesh.getMeshName())){
                    currentMesh.setTextureMask(textureMap.get(currentMesh.getMeshName()));
                }
                //draw
                currentMesh.complexDraw(renderPipelineState);
                //reset texture mask
                currentMesh.setTextureMask(null);
                //reset shader
                currentMesh.setShader(original);
            }
        }
        if(meshMask != null){
            for(Mesh toDraw : meshMask.getToDrawMeshes()){
                toDraw.setBones(bones);
                toDraw.setParent(this);
                ShaderProgram original = toDraw.getShader();
                ShaderProgram shader = getCorrectShader(shaderMask, toDraw, toDraw.getShader());
                toDraw.setShader(shader);
                toDraw.complexDraw(renderPipelineState);
                toDraw.setShader(original);
            }
        }
    }

    /**
     * Determines the correct shader to use for a given mesh
     * @param shaderMask The shader mask
     * @param mesh The mesh
     * @param oldShader The original shader on the mesh
     * @return The correct shader program to use for this mesh
     */
    ShaderProgram getCorrectShader(Map<String,ActorShaderMask> shaderMask, Mesh mesh, ShaderProgram oldShader){
        ShaderProgram rVal = oldShader;
        if(shaderMask.containsKey(mesh.getMeshName())){
            ActorShaderMask specificMask = shaderMask.get(mesh.getMeshName());
            ShaderProgram overwriteShader = null;
            if((overwriteShader = Globals.assetManager.fetchShader(specificMask.getVertexShaderPath(), specificMask.getGeometryShaderPath(), specificMask.getFragmentShaderPath())) != null){
                rVal = overwriteShader;
            }
        }
        return rVal;
    }
    

    /**
     * Applies an animation mask to the model
     * @param animationName The name of the animation
     * @param time The temporal location in the animation to mask
     * @param mask The mask
     */
    public void applyAnimationMask(String animationName, double time, List<String> mask){
        Animation animationCurrent = animMap.get(animationName);
        if(animationCurrent != null){
            for(String boneName : mask){
                AnimChannel currentChannel = animationCurrent.getChannel(boneName);
                Bone currentBone = boneMap.get(currentChannel.getNodeID());
                currentChannel.setTime(time);
                // System.out.println(currentChannel + " " + currentBone);
                if(currentBone != null){
                    // System.out.println("Applying to bone");
                    //T * S * R
                    currentBone.deform = new Matrix4d();
                    currentBone.deform.translate(currentChannel.getCurrentPosition());
                    currentBone.deform.rotate(currentChannel.getCurrentRotation());
                    currentBone.deform.scale(new Vector3d(currentChannel.getCurrentScale()));
                }
            }
        }
    }
    
    /**
     * Logs all animations for a given model
     */
    public void describeAllAnimations(){
        if(animations.size() > 0){
            LoggerInterface.loggerRenderer.DEBUG("=====================");
            LoggerInterface.loggerRenderer.DEBUG(animations.size() + " animations available in model!");
            Iterator<Animation> animIterator = animations.iterator();
            while(animIterator.hasNext()){
                Animation currentAnim = animIterator.next();
                currentAnim.describeAnimation();
            }
        }
    }
    
    /**
     * Recursively builds the bone tree
     * @param node The current assimp bone to operate on
     * @param parent The parent bone
     * @return The current bone
     */
    public final AnimNode buildAnimNodeMap(AINode node, AnimNode parent){
        AnimNode node_object = new AnimNode(node.mName().dataString(), parent, node);
        nodeMap.put(node_object.id, node_object);
        if(boneMap.containsKey(node_object.id)){
            node_object.is_bone = true;
        }
        int num_children = node.mNumChildren();
        for(int i = 0; i < num_children; i++){
            AnimNode temp_child = buildAnimNodeMap(AINode.create(node.mChildren().get(i)),node_object);
            node_object.children.add(temp_child);
        }
        return node_object;
    }
    
    /**
     * Updates the position of all bones to their correct locations given bone rotators, static morph, and current animation
     * @param boneRotators The bone rotators to apply
     * @param staticMorph The static morph to apply
     */
    public void updateNodeTransform(Map<String,ActorBoneRotator> boneRotators, ActorStaticMorph staticMorph){
        if(this.rootAnimNode != null){
            updateNodeTransform(this.rootAnimNode,boneRotators,staticMorph);
        }
    }

    /**
     * Recursively updates the position of all bones to their correct locations given bone rotators, static morph, and current animation
     * @param n the current node to operate on
     * @param boneRotators The bone rotators to apply
     * @param staticMorph The static morph to apply
     */
    void updateNodeTransform(AnimNode n, Map<String,ActorBoneRotator> boneRotators, ActorStaticMorph staticMorph){
        if(n.parent != null){
            n.transform = new Matrix4d(n.parent.transform);
        } else {
            n.transform = new Matrix4d();
        }
        if(n.is_bone){
            //
            //bone rotators (turrets, hair, etc)
            Bone target_bone = boneMap.get(n.id);
            n.transform = n.transform.mul(target_bone.deform);
            if(boneRotators.containsKey(target_bone.boneID)){
                n.transform.rotate(boneRotators.get(target_bone.boneID).getRotation());
            }
            //
            //static morph (changing nose size, eye distance, etc)
            Matrix4d bone_matrix = new Matrix4d(n.transform);
            if(staticMorph != null && staticMorph.getBoneTransforms(n.id) != null){
                bone_matrix.mul(staticMorph.getBoneTransforms(n.id).getTransform());
                n.transform.mul(staticMorph.getBoneTransforms(n.id).getTransform());
            }
            //
            //Calculate final offset from initial bone
            bone_matrix.mul(target_bone.inverseBindPoseMatrix);
            bone_matrix = new Matrix4d(globalInverseTransform).mul(bone_matrix);
            target_bone.final_transform = bone_matrix;
        } else {
            n.transform = n.transform.mul(electrosphere.util.Utilities.convertAIMatrix(n.raw_data.mTransformation()));
        }
        Iterator<AnimNode> node_iterator = n.children.iterator();
        while(node_iterator.hasNext()){
            AnimNode current_node = node_iterator.next();
            updateNodeTransform(current_node,boneRotators,staticMorph);
        }
    }

    /**
     * Draws a ui model
     */
    public void drawUI(){
        for(Mesh m : meshes){
            m.complexDraw(Globals.renderingEngine.getRenderPipelineState());
        }
    }
    

    /**
     * Pushes a uniform to a given mesh
     * @param meshName The name of the mesh
     * @param uniformKey The uniform key
     * @param uniform The value of the uniform
     */
    public void pushUniformToMesh(String meshName, String uniformKey, Object uniform){
        for(Mesh m : meshes){
            if(m.getMeshName().equals(meshName)){
                m.setUniform(uniformKey, uniform);
            }
        }
    }
    
    /**
     * Logs all bone IDs (names)
     */
    public void listAllBoneIDs(){
        for(String id : boneMap.keySet()){
            LoggerInterface.loggerRenderer.DEBUG(id);
        }
    }
    
    /**
     * Gets a given mesh by its name
     * @param meshName the name of the mesh
     * @return The mesh if it exists, or null
     */
    public Mesh getMesh(String meshName){
        for(Mesh mesh : meshes){
            if(mesh.getMeshName().matches(meshName)){
                return mesh;
            }
        }
        return null;
    }

    /**
     * Gets an animation by its name
     * @param animName The name of the animation
     * @return the animation if it exists, or null
     */
    public Animation getAnimation(String animName){
        return animMap.get(animName);
    }
    
    /**
     * Describes the model at a high level
     */
    public void describeHighLevel(){
        LoggerInterface.loggerRenderer.DEBUG("Meshes: ");
        for(Mesh mesh : meshes){
            LoggerInterface.loggerRenderer.DEBUG(mesh.getMeshName());
        }
        LoggerInterface.loggerRenderer.DEBUG("Animations: ");
        for(Animation anim : animations){
            LoggerInterface.loggerRenderer.DEBUG(anim.name);
        }
        LoggerInterface.loggerRenderer.DEBUG("Bones:");
        for(Bone bone : bones){
            LoggerInterface.loggerRenderer.DEBUG(bone.boneID);
        }
    }



    //
    // Pure getters and setters
    //
    

    /**
     * Gets the meshes in this model
     * @return the list of meshes
     */
    public List<Mesh> getMeshes(){
        return meshes;
    }

    /**
     * Sets the model matrix for this model
     * @param modelMatrix the model matrix
     */
    public void setModelMatrix(Matrix4d modelMatrix){
        this.modelMatrix = modelMatrix;
    }

    /**
     * Gets the model matrix for this model
     * @return the model matrix
     */
    public Matrix4d getModelMatrix(){
        return modelMatrix;
    }

    /**
     * Gets the list of bones
     * @return the list of bones
     */
    public List<Bone> getBones(){
        return bones;
    }

    /**
     * Gets the map of bone name -> bone
     * @return the map
     */
    public Map<String,Bone> getBoneMap(){
        return boneMap;
    }

    /**
     * Gets the list of materials in this model
     * @return The list of materials
     */
    public List<Material> getMaterials(){
        return materials;
    }

    /**
     * Sets the mesh mask
     * @param meshMask the mesh mask to set
     */
    public void setMeshMask(ActorMeshMask meshMask){
        this.meshMask = meshMask;
    }

    /**
     * Sets the texture mask
     * @param textureMask the texture mask
     */
    public void setTextureMask(Map<String,ActorTextureMask> textureMask){
        this.textureMap = textureMask;
    }

    /**
     * Gets the shader mask map
     * @return The shader mask map
     */
    public Map<String,ActorShaderMask> getShaderMask(){
        return shaderMask;
    }


    /**
     * Utility method to attempt overwriting the model's meshes with a new material
     * @param material The material
     */
    public void tryOverwriteMaterial(Material material){
        for(Mesh mesh : meshes){
            mesh.setMaterial(material);
        }
    }

    /**
     * Gets the bounding sphere for this model
     * @return The bounding sphere
     */
    public Sphered getBoundingSphere(){
        return boundingSphere;
    }

    /**
     * Sets the bounding sphere
     * @param boundingSphere The bounding sphere to be stored
     */
    public void setBoundingSphere(Sphered boundingSphere){
        this.boundingSphere = boundingSphere;
    }
}