package electrosphere.renderer;

import electrosphere.renderer.anim.AnimChannel;
import electrosphere.renderer.anim.Animation;
import electrosphere.renderer.anim.AnimNode;
import electrosphere.main.Globals;
import java.io.File;
import java.util.ArrayList;
import java.util.List;
import java.util.ListIterator;
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.lwjgl.assimp.AIAnimation;
import org.lwjgl.assimp.AIBone;
import org.lwjgl.assimp.AINode;

/**
 *
 * @author satellite
 */
public class 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;
    public ArrayList<Material> materials;
    public Matrix4f modelMatrix = new Matrix4f();
    ShaderProgram program;
    public Animation currentAnimation;
    public Matrix4f globalInverseTransform;
    
    AnimNode root_anim_node;
    
    public static Model createModelFromAiscene(AIScene s){
        Model rVal = new Model();
        //
        //set the scene
        //
        rVal.scene = s;
        
        //
        //load meshes
        //
        int meshCount = s.mNumMeshes();
        PointerBuffer meshesBuffer = s.mMeshes();
        rVal.meshes = new ArrayList();
        while(meshesBuffer.hasRemaining()){
            Mesh currentMesh = Mesh.create_mesh_from_aimesh(AIMesh.create(meshesBuffer.get()));
            rVal.meshes.add(currentMesh);
            currentMesh.parent = rVal;
        }
        //
        //register bones
        //
        rVal.bones = new ArrayList();
        rVal.boneMap = new HashMap();
        meshesBuffer.rewind();
        while(meshesBuffer.hasRemaining()){
            AIMesh currentMeshData = AIMesh.createSafe(meshesBuffer.get());
            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();
        rVal.boneMap = new HashMap();
        rVal.node_map = new HashMap();
        while(meshIterator.hasNext()){
            Mesh currentMesh = meshIterator.next();
            Iterator<Bone> boneIterator = currentMesh.bones.iterator();
            ArrayList<Bone> to_remove_queue = new ArrayList();
            ArrayList<Bone> to_add_queue = new ArrayList();
            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.convertAIMatrix(rootNode.mTransformation());
//        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();
        for(int i = 0; i < animCount; i++){
            rVal.animations.add(new Animation(AIAnimation.create(animBuffer.get(i)), i));
        }
        //
        //Load materials
        //
        if(s.mNumMaterials() > 0){
            rVal.materials = new ArrayList();
            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 Matrix4f();
        program = null;
        currentAnimation = null;
        globalInverseTransform = null;
    }
    
    public void free() {
        aiReleaseImport(scene);
        scene = null;
        meshes = null;
    }
    
    public void draw(boolean setShader, boolean bufferStandardUniforms, boolean bufferNonStandardUniforms, boolean useMaterial, boolean useShadowMap, boolean setBones, boolean useLight){
        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();
            currentMesh.complexDraw(setShader, bufferStandardUniforms, bufferNonStandardUniforms, useMaterial, useShadowMap, setBones, useLight);
        }
    }
    
    
    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){
            currentAnimation.timeCurrent = 0;
            Iterator<AnimChannel> channelIterator = currentAnimation.channels.iterator();
            while(channelIterator.hasNext()){
                AnimChannel currentChannel = channelIterator.next();
                currentChannel.rewind();
            }
        }
    }
    
    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.nodeID);
                    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 update_node_transform(AnimNode n){
        if(n.parent != null){
            n.transform = new Matrix4f(n.parent.transform);
        } else {
            n.transform = new Matrix4f();
        }
        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 target_bone = boneMap.get(n.id);
            n.transform = n.transform.mul(target_bone.deform);
            Matrix4f bone_matrix = new Matrix4f(n.transform).mul(target_bone.inverseBindPoseMatrix);
            bone_matrix = new Matrix4f(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);
        }
    }
    
    
    public void drawForDepthBuffer(){
        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();
            currentMesh.setUniform("model", modelMatrix);
            if(currentMesh.hasBones){
                currentMesh.complexDraw(false, false, true, false, false, true, false);
            } else {
                currentMesh.setUniform("hasBones", 0);
                currentMesh.complexDraw(false, false, true, false, false, false, false);
            }
        }
    }
    
    public void drawUI(){
        for(Mesh m : meshes){
            m.complexDraw(true, false, true, true, false, false, false);
        }
    }
    
    public void drawFramebuffer(){
        for(Mesh m : meshes){
            m.complexDraw(true, false, true, false, false, false, false);
        }
    }
    
    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 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);
        }
    }
}
