package electrosphere.renderer;

import electrosphere.renderer.actor.ActorAnimationMask;
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.loading.ModelPretransforms.MeshMetadata;
import electrosphere.renderer.anim.AnimNode;
import electrosphere.engine.Globals;
import electrosphere.logger.LoggerInterface;

import java.io.File;
import java.util.ArrayList;
import java.util.List;
import java.util.ListIterator;

import org.joml.Matrix4d;
import org.joml.Matrix4f;
import org.joml.Vector3f;
import org.lwjgl.BufferUtils;
import org.lwjgl.PointerBuffer;
import org.lwjgl.assimp.AIMaterial;
import org.lwjgl.assimp.AIMesh;
import org.lwjgl.assimp.AIScene;
import static org.lwjgl.assimp.Assimp.*;
import static org.lwjgl.opengl.ARBVertexBufferObject.*;
import static org.lwjgl.opengl.ARBVertexProgram.*;
import org.lwjgl.opengl.GL11;
import static org.lwjgl.opengl.GL11.*;
import static org.lwjgl.opengl.GL13.*;
import static org.lwjgl.opengl.GL15.*;
import static org.lwjgl.opengl.GL20.*;
import static org.lwjgl.opengl.GL30.*;
import electrosphere.renderer.texture.TextureMap;
import java.nio.FloatBuffer;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Map;
import java.util.Stack;
import org.joml.Quaternionf;
import org.joml.Vector3d;
import org.lwjgl.assimp.AIAnimation;
import org.lwjgl.assimp.AIBone;
import org.lwjgl.assimp.AINode;

/**
 *
 * @author satellite
 */
public class Model {
    boolean is_Ready_To_Display = false;
    AIScene scene;
    public List<Mesh> meshes;
    public ArrayList<Animation> animations;
    public ArrayList<Bone> bones;
    public HashMap<String,Bone> boneMap;
    HashMap<String,Animation> animMap;
    public HashMap<String,AnimNode> node_map;
    public ArrayList<Material> materials;
    public Matrix4d modelMatrix = new Matrix4d();
    ShaderProgram program;
    Matrix4d globalInverseTransform;
    
    AnimNode root_anim_node;
    ActorMeshMask meshMask;
    Map<String,ActorShaderMask> shaderMask = new HashMap<String,ActorShaderMask>();
    Map<String,ActorTextureMask> textureMap = null;
    
    public static Model createModelFromAiscene(String path, AIScene s){
        Model rVal = new Model();
        //
        //set the scene
        //
        rVal.scene = s;

        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 = Mesh.create_mesh_from_aimesh(aiMesh, meshMetadata);
            rVal.meshes.add(currentMesh);
            currentMesh.parent = rVal;
        }
        //
        //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.node_map = new HashMap<String, AnimNode>();
        while(meshIterator.hasNext()){
            Mesh currentMesh = meshIterator.next();
            Iterator<Bone> boneIterator = currentMesh.bones.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.bones.remove(boneIterator.next());
            }
            boneIterator = to_add_queue.iterator();
            while(boneIterator.hasNext()){
                currentMesh.bones.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());
        }
//        System.out.println("Global inverse transform");
//        System.out.println(globalInverseTransform);
        rVal.root_anim_node = rVal.build_anim_node_map(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;
    }
    
    
    
    public Model(){
        /*
        public boolean is_Ready_To_Display = false;
        public AIScene scene;
        public ArrayList<Mesh> meshes;
        public ArrayList<Animation> animations;
        public ArrayList<Bone> bones;
        public HashMap<String,Bone> boneMap;
        public HashMap<String,AnimNode> node_map = new HashMap();
        public ArrayList<Material> materials;
        public Matrix4f modelMatrix = new Matrix4f().translation(new Vector3f(0,0,0));
        ShaderProgram program;
        public Animation currentAnimation;
        public Matrix4f globalInverseTransform;
        */
        scene = null;
        meshes = null;
        animations = null;
        bones = null;
        boneMap = null;
        node_map = null;
        materials = null;
        modelMatrix = new Matrix4d();
        program = null;
        globalInverseTransform = null;
    }
    
    public void free() {
        aiReleaseImport(scene);
        scene = null;
        meshes = null;
    }
    
    public void draw(RenderPipelineState renderPipelineState){
        // if(node_map != null && !node_map.isEmpty()){
        //     update_node_transform(root_anim_node);
        // }
        Iterator<Mesh> mesh_Iterator = meshes.iterator();
        while(mesh_Iterator.hasNext()){
            Mesh currentMesh = mesh_Iterator.next();
            if(meshMask == null || (meshMask != null && !meshMask.isBlockedMesh(currentMesh.nodeID))){
                //set shader
                ShaderProgram original = currentMesh.shader;
                ShaderProgram shader = getCorrectShader(shaderMask, currentMesh, currentMesh.shader);
                currentMesh.shader = shader;
                //set texture mask
                if(this.textureMap != null && textureMap.containsKey(currentMesh.nodeID)){
                    currentMesh.setTextureMask(textureMap.get(currentMesh.nodeID));
                }
                //draw
                currentMesh.complexDraw(renderPipelineState);
                //reset texture mask
                currentMesh.setTextureMask(null);
                //reset shader
                currentMesh.shader = original;
            }
        }
        if(meshMask != null){
            for(Mesh toDraw : meshMask.getToDrawMeshes()){
                toDraw.bones = bones;
                toDraw.parent = this;
                ShaderProgram original = toDraw.shader;
                ShaderProgram shader = getCorrectShader(shaderMask, toDraw, toDraw.shader);
                toDraw.shader = shader;
                toDraw.complexDraw(renderPipelineState);
                toDraw.shader = original;
            }
        }
    }

    ShaderProgram getCorrectShader(Map<String,ActorShaderMask> shaderMask, Mesh mesh, ShaderProgram oldShader){
        ShaderProgram rVal = oldShader;
        if(shaderMask.containsKey(mesh.nodeID)){
            ActorShaderMask specificMask = shaderMask.get(mesh.nodeID);
            ShaderProgram overwriteShader = null;
            if((overwriteShader = Globals.assetManager.fetchShader(specificMask.getVertexShaderPath(), specificMask.getGeometryShaderPath(), specificMask.getFragmentShaderPath())) != null){
                // ShaderProgram oldProgram = mesh.shader;
                rVal = overwriteShader;
            }
        }
        return rVal;
    }
    
    
//     public void playAnimation(String s){
//         this.currentAnimation = null;
//         Iterator<Animation> animationIterator = animations.iterator();
//         while(animationIterator.hasNext()){
//             Animation current_iterator_item = animationIterator.next();
//             if(current_iterator_item.name.equals(s)){
//                 this.currentAnimation = current_iterator_item;
//             }
//         }
//         if(currentAnimation != null){
// //            if(s.contains("Walk")){
// //                currentAnimation.describeAnimation();
// ////                currentAnimation.fullDescribeAnimation();
// //                System.exit(0);
// //            }
//             currentAnimation.timeCurrent = 0;
//             Iterator<AnimChannel> channelIterator = currentAnimation.channels.iterator();
//             while(channelIterator.hasNext()){
//                 AnimChannel currentChannel = channelIterator.next();
//                 currentChannel.rewind();
//             }
//         }
//     }

    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()));
                }
            }
        }
    }
    
    // public void incrementTime(double time){
    //     if(currentAnimation != null){
    //         boolean isDone = currentAnimation.incrementTime(time);
    //         if(isDone){
    //             currentAnimation.timeCurrent = 0;
    //             Iterator<AnimChannel> channelIterator = currentAnimation.channels.iterator();
    //             while(channelIterator.hasNext()){
    //                 AnimChannel currentChannel = channelIterator.next();
    //                 currentChannel.rewind();
    //             }
    //             Iterator<Bone> boneIterator = bones.iterator();
    //             while(boneIterator.hasNext()){
    //                 Bone currentBone = boneIterator.next();
    //                 currentBone.transform = new Matrix4f();
    //             }
    //             currentAnimation = null;
    //         } else {
    //             //First we push transformFromParent into deform so that later in the pipeline bones without a current animation take on transformFromParent as their transformation to the hierarchy
    //             //I BELIEVE this is so that control bones still apply their offset to the hierarchy even when they're not animated :tm:
    //             //4/5/20
    // //            Iterator<Bone> boneIterator = bones.iterator();
    // //            while(boneIterator.hasNext()){
    // //                Bone currentBone = boneIterator.next();
    // ////                currentBone.deform = currentBone.transformFromParent;
    // //            }
    //             //Once that's done, for every channel we set the corresponding bone's deform to the channels TRS 
    //             Iterator<AnimChannel> channelIterator = currentAnimation.channels.iterator();
    //             while(channelIterator.hasNext()){
    //                 AnimChannel currentChannel = channelIterator.next();
    //                 currentChannel.incrementTime(time);
    //                 Bone currentBone = boneMap.get(currentChannel.getNodeID());
    //                 if(currentBone != null){
    //                     //T * S * R
    //                     currentBone.deform = new Matrix4f();
    //                     currentBone.deform.translate(currentChannel.getCurrentPosition());
    //                     currentBone.deform.rotate(currentChannel.getCurrentRotation());
    //                     currentBone.deform.scale(currentChannel.getCurrentScale());
    //                 }
    //             }
    //         }
    //     }
    // }
    
    public void describeAllAnimations(){
        if(animations.size() > 0){
            System.out.println("=====================");
            System.out.println(animations.size() + " animations available in model!");
            Iterator<Animation> animIterator = animations.iterator();
            while(animIterator.hasNext()){
                Animation currentAnim = animIterator.next();
                currentAnim.describeAnimation();
            }
        }
    }
    
    public final AnimNode build_anim_node_map(AINode node, AnimNode parent){
        AnimNode node_object = new AnimNode(node.mName().dataString(), parent, node);
        node_map.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 = build_anim_node_map(AINode.create(node.mChildren().get(i)),node_object);
//            if(boneMap.containsKey(node_object.id)){
//                Bone parent_bone = boneMap.get(node_object.id);
//                node_object.children.add(temp_child);
//                if(boneMap.containsKey(temp_child.id)){
//                    Bone child_bone = boneMap.get(temp_child.id);
//                    parent_bone.children.add(child_bone);
//                    child_bone.parent = parent_bone;
//                }
//            }
            node_object.children.add(temp_child);
        }
        return node_object;
    }
    
    public void updateNodeTransform(Map<String,ActorBoneRotator> boneRotators, ActorStaticMorph staticMorph){
        if(this.root_anim_node != null){
            update_node_transform(this.root_anim_node,boneRotators,staticMorph);
        }
    }

    void update_node_transform(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 bone = boneList.get(j);
                Node node = rootNode.findByName(bone.getBoneName());
                Matrix4f boneMatrix = Node.getParentTransforms(node, i);
                boneMatrix.mul(bone.getOffsetMatrix());
                boneMatrix = new Matrix4f(rootTransformation).mul(boneMatrix);
                frame.setMatrix(j, boneMatrix);
            */
            //
            //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;
//            if(!toggled){
//                System.out.println(n.id);
//                System.out.println("Transform:");
//                System.out.println(n.parent.transform);
//                System.out.println(target_bone.deform);
//                System.out.println(n.transform);
//                System.out.println("Final transform:");
//                System.out.println(globalInverseTransform);
//                System.out.println(n.transform);
//                System.out.println(target_bone.inverseBindPoseMatrix);
//                System.out.println(target_bone.final_transform);
//                System.out.println("\n\n\n\n");
//            }
        } 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();
            update_node_transform(current_node,boneRotators,staticMorph);
        }
    }
    
    public void drawUI(){
        for(Mesh m : meshes){
            m.complexDraw(Globals.renderingEngine.getRenderPipelineState());
        }
    }
    
    public void pushUniformToMesh(String meshName, String uniformKey, Object uniform){
        for(Mesh m : meshes){
            if(m.nodeID.equals(meshName)){
                m.setUniform(uniformKey, uniform);
            }
        }
    }
    
    public void listAllBoneIDs(){
        for(String id : boneMap.keySet()){
            System.out.println(id);
        }
    }
    
    public Mesh getMesh(String meshName){
        for(Mesh mesh : meshes){
            if(mesh.nodeID.matches(meshName)){
                return mesh;
            }
        }
        return null;
    }

    public Animation getAnimation(String animName){
        return animMap.get(animName);
    }

    public Map<String,ActorShaderMask> getShaderMask(){
        return shaderMask;
    }
    
    public void describeHighLevel(){
        System.out.println("Meshes: ");
        for(Mesh mesh : meshes){
            System.out.println(mesh.nodeID);
        }
        System.out.println("Animations: ");
        for(Animation anim : animations){
            System.out.println(anim.name);
        }
        System.out.println("Bones:");
        for(Bone bone : bones){
            System.out.println(bone.boneID);
        }
    }

    public void setMeshMask(ActorMeshMask meshMask){
        this.meshMask = meshMask;
    }

    public void setTextureMask(Map<String,ActorTextureMask> textureMask){
        this.textureMap = textureMask;
    }

    /**
     * Sets the shader program to use drawing the modal
     * @param program The shader program
     */
    public void setProgram(ShaderProgram program){
        this.program = program;
    }

    /**
     * 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);
        }
    }
}