#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <vector>

#define GLM_FORCE_RADIANS
#define GLM_FORCE_DEPTH_ZERO_TO_ONE
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <glm/gtc/type_ptr.hpp>

#define TINYGLTF_IMPLEMENTATION
#define STB_IMAGE_IMPLEMENTATION
#define TINYGLTF_NO_STB_IMAGE_WRITE

#ifdef VK_USE_PLATFORM_ANDROID_KHR
	#define TINYGLTF_ANDROID_LOAD_FROM_ASSETS
#endif
#include "tiny_gltf.h"

#include "vulkanexamplebase.h"

#define ENABLE_VALIDATION false


// Contains everything required to render a glTF model in Vulkan
// This class is heavily simplified (compared to glTF's feature set) but retains the basic glTF structure
class VulkanglTFModel
{
public:
	// The class requires some Vulkan objects so it can create it's own resources
	vks::VulkanDevice* vulkanDevice;
	VkQueue copyQueue;
	uint32_t nodeCount;

	// The vertex layout for the samples' model
	struct Vertex {
		glm::vec3 pos;
		glm::vec3 normal;
		glm::vec2 uv;
		glm::vec3 color;
		glm::vec3 tangent;
		glm::vec3 jointIndices;
		glm::vec3 jointWeights;
	};

	// Single vertex buffer for all primitives
	struct Vertices {
		VkBuffer buffer;
		VkDeviceMemory memory;
	} vertices;

	// Single index buffer for all primitives
	struct Indices {
		int count;
		VkBuffer buffer;
		VkDeviceMemory memory;
	} indices;

	// The following structures roughly represent the glTF scene structure
	// To keep things simple, they only contain those properties that are required for this sample
	struct Node;

	// A primitive contains the data for a single draw call
	struct Primitive {
		uint32_t firstIndex;
		uint32_t indexCount;
		int32_t materialIndex;
	};

	// Contains the node's (optional) geometry and can be made up of an arbitrary number of primitives
	struct Mesh {
		std::vector<Primitive> primitives;
	};

	// A node represents an object in the glTF scene graph
	struct Node {
		Node* parent;
		uint32_t index;
		std::vector<Node*>  children;
		Mesh mesh;
		glm::vec3 translation{};
		glm::vec3 scale{ 1.0f };
		glm::quat rotation{};
		int32_t             skin = -1;
		glm::mat4 getLocalMatrix()
		{
			return bAnimateNode ? glm::translate(glm::mat4(1.0f), translation) * glm::mat4(rotation) * glm::scale(glm::mat4(1.0f), scale) : matrix;
		}
		glm::mat4 matrix;
		bool bAnimateNode = false;
		
		~Node() {
			for (auto& child : children) {
				delete child;
			};
		}
		
		

	};
	// material data for pbr
	struct MaterialData
	{
		vks::Buffer buffer;
		VkDescriptorSet descriptorSet;
		struct Values
		{
			glm::vec3 emissiveFactor;
			glm::vec4 baseColorFactor;
		}values;

	};
	// A glTF material stores information in e.g. the texture that is attached to it and colors
	struct Material {
		glm::vec4 baseColorFactor = glm::vec4(1.0f);
		uint32_t baseColorTextureIndex;
		uint32_t normalMapTextureIndex;
		uint32_t matalicRoughTextureIndex;
		int32_t emissiveTextureIndex = -1;
		MaterialData materialData;
	};

	// Contains the texture for a single glTF image
	// Images may be reused by texture objects and are as such separated
	struct Image {
		vks::Texture2D texture;
		// We also store (and create) a descriptor set that's used to access this texture from the fragment shader
		VkDescriptorSet descriptorSet;
	};

	// A glTF texture stores a reference to the image and a sampler
	// In this sample, we are only interested in the image
	struct Texture {
		int32_t imageIndex;
	};

	// structure of skin
	struct Skin {
		std::string name;
		Node* skeletonRoot = nullptr;
		std::vector<glm::mat4> inverseBindMatrices;
		std::vector<Node*> joints;
		vks::Buffer ssbo;
		VkDescriptorSet descriptorSet;
		
	};

	struct AnimationSampler
	{
		std::string interpolation;
		std::vector<float> inputs;
		std::vector<glm::vec4> outputsVec4;

	};

	struct AnimationChannel
	{
		std::string path;
		Node* node;
		uint32_t samplerIndex;

	};

	struct Animation
	{
		std::string name;
		std::vector<AnimationSampler> samplers;
		std::vector<AnimationChannel> channels;

		float start = std::numeric_limits<float>::max();
		float end = std::numeric_limits<float>::min();
		float currentTime = 0.0f;
	};

	/*
		Model data
	*/
	std::vector<Image> images;
	std::vector<Texture> textures;
	std::vector<Material> materials;
	std::vector<Node*> nodes;
	std::vector<Skin> skins;
	std::vector<Animation> animations;

	uint32_t activeAnimation = 0;



	 //VulkanglTFModel();
	~VulkanglTFModel()
	{
		for (auto node : nodes) {
			delete node;
		}
		// Release all Vulkan resources allocated for the model
		vkDestroyBuffer(vulkanDevice->logicalDevice, vertices.buffer, nullptr);
		vkFreeMemory(vulkanDevice->logicalDevice, vertices.memory, nullptr);
		vkDestroyBuffer(vulkanDevice->logicalDevice, indices.buffer, nullptr);
		vkFreeMemory(vulkanDevice->logicalDevice, indices.memory, nullptr);
		for (auto& material : materials)
		{
			material.materialData.buffer.destroy();
		}
		for (Image image : images) {
			vkDestroyImageView(vulkanDevice->logicalDevice, image.texture.view, nullptr);
			vkDestroyImage(vulkanDevice->logicalDevice, image.texture.image, nullptr);
			vkDestroySampler(vulkanDevice->logicalDevice, image.texture.sampler, nullptr);
			vkFreeMemory(vulkanDevice->logicalDevice, image.texture.deviceMemory, nullptr);
		}
	}
	void      loadImages(tinygltf::Model& input);
	void      loadTextures(tinygltf::Model& input);
	void      loadMaterials(tinygltf::Model& input);
	Node* findNode(Node* parent, uint32_t index);
	Node* nodeFromIndex(uint32_t index);
	void      loadSkins(tinygltf::Model& input);
	void      loadAnimations(tinygltf::Model& input);
	void      loadNode(const tinygltf::Node& inputNode, const tinygltf::Model& input, VulkanglTFModel::Node* parent, uint32_t nodeIndex, std::vector<uint32_t>& indexBuffer, std::vector<VulkanglTFModel::Vertex>& vertexBuffer);
	glm::mat4 getNodeMatrix(VulkanglTFModel::Node* node);
	void updateNodeMatrix(Node* node, std::vector<glm::mat4>& nodeMatrics);
	void      updateJoints(VulkanglTFModel::Node* node);
	void      updateAnimation(float deltaTime,vks::Buffer buffer);
	void drawNode(VkCommandBuffer commandBuffer, VkPipelineLayout pipelineLayout, VulkanglTFModel::Node node);
	void      draw(VkCommandBuffer commandBuffer, VkPipelineLayout pipelineLayout);
};

class VulkanExample : public VulkanExampleBase
{
public:
	bool wireframe = false;
	bool normalMapping = true;
	bool ToneMapping = true;
	bool pbrEnabled = true;

	VulkanglTFModel glTFModel;

	struct ShaderData {
		vks::Buffer buffer;
		struct Values {
			glm::mat4 projection;
			glm::mat4 model;
			glm::vec4 lightPos = glm::vec4(5.0f, 5.0f, 5.0f, 1.0f);
			glm::vec4 viewPos;
			glm::vec4 bFlagSet = glm::vec4(0.0f, 0.0f, 0.0f, 0.0f);
		} values;
		vks::Buffer skinSSBO;
	} shaderData;

	struct Pipelines {
		VkPipeline solid;
		VkPipeline wireframe = VK_NULL_HANDLE;
		VkPipeline toneMapping = VK_NULL_HANDLE;
	} pipelines;

	struct  PipelineLayouts
	{
		VkPipelineLayout pbrLayout;
		VkPipelineLayout tonemappingLayout;
	} pipelineLayouts;

	VkPipelineLayout pipelineLayout;
	
	VkDescriptorSet descriptorSet;
	VkDescriptorSet skinDescriptorSet;
	VkDescriptorSet tonemappingDescriptorSet = VK_NULL_HANDLE;

	struct FrameBufferAttachment
	{
		VkImage image;
		VkDeviceMemory deviceMemory;
		VkImageView imageView;
		VkFormat format;
		

		void destroy(VkDevice device)
		{
			vkDestroyImage(device, image, nullptr);
			vkDestroyImageView(device, imageView,nullptr);
			vkFreeMemory(device, deviceMemory, nullptr);

		}
	};

	struct FrameBuffer
	{
		int32_t width, height;
		VkFramebuffer frameBuffer;
		VkRenderPass renderPass;
		void setSize(int32_t w, int32_t h)
		{
			this->width = w;
			this->height = h;
		}
		void destroy(VkDevice device)
		{
			vkDestroyFramebuffer(device, frameBuffer, nullptr);
			vkDestroyRenderPass(device, renderPass, nullptr);
		}
	};

	struct PBRFrameBuffer {
		FrameBufferAttachment color, depth;
		FrameBuffer fbo;
		bool bCreate = false;
	} pbrFrameBuffer;

	VkSampler colorSampler;

	struct DescriptorSetLayouts {
		VkDescriptorSetLayout matrices;
		VkDescriptorSetLayout textures;
		VkDescriptorSetLayout materialUniform;
		VkDescriptorSetLayout ssbo;
		VkDescriptorSetLayout jointMatrices;
	} descriptorSetLayouts;

	struct IBLTextures
	{
		vks::TextureCubeMap skyboxCube;
		vks::TextureCubeMap irradianceCube;
		vks::TextureCubeMap prefilteredCube;
		vks::Texture2D lutBrdf;
	} ibltextures;

	VulkanglTFModel skyboxModel;

	VulkanExample();
	~VulkanExample()
	{
		// Clean up used Vulkan resources
		// Note : Inherited destructor cleans up resources stored in base class
		vkDestroyPipeline(device, pipelines.solid, nullptr);
		vkDestroyPipeline(device, pipelines.toneMapping, nullptr);
		if (pipelines.wireframe != VK_NULL_HANDLE) {
			vkDestroyPipeline(device, pipelines.wireframe, nullptr);
		}

		vkDestroyPipelineLayout(device, pipelineLayouts.pbrLayout, nullptr);
		vkDestroyPipelineLayout(device, pipelineLayouts.tonemappingLayout, nullptr);
		vkDestroyDescriptorSetLayout(device, descriptorSetLayouts.matrices, nullptr);
		vkDestroyDescriptorSetLayout(device, descriptorSetLayouts.textures, nullptr);
		vkDestroyDescriptorSetLayout(device, descriptorSetLayouts.materialUniform, nullptr);
		vkDestroyDescriptorSetLayout(device, descriptorSetLayouts.ssbo, nullptr);
		ibltextures.irradianceCube.destroy();
		ibltextures.skyboxCube.destroy();
		ibltextures.prefilteredCube.destroy();
		ibltextures.lutBrdf.destroy();
		pbrFrameBuffer.color.destroy(device);
		pbrFrameBuffer.depth.destroy(device);
		pbrFrameBuffer.fbo.destroy(device);
		vkDestroySampler(device, colorSampler, nullptr);

		shaderData.buffer.destroy();
		shaderData.skinSSBO.destroy();
	}
	void         loadglTFFile(std::string filename);
	virtual void getEnabledFeatures();
	void createAttachment(VkFormat format, VkImageUsageFlagBits usage, FrameBufferAttachment* attachment, uint32_t width, uint32_t height);
	virtual void setupFrameBuffer();
	void         buildCommandBuffers();
	void         loadAssets();
	void         setupDescriptors();
	void         preparePipelines();
	void         prepareUniformBuffers();
	void         updateUniformBuffers();
	void         prepare();
	virtual void render();
	virtual void viewChanged();
	virtual void OnUpdateUIOverlay(vks::UIOverlay* overlay);
};