#include <iostream>
#include <vector>
#include <fstream>
#include <limits>
#include <stdint.h> 

#include <assimp/Importer.hpp>      // C++ importer interface
#include <assimp/scene.h>           // Output data structure
#include <assimp/postprocess.h>     // Post processing flags
#include <assimp/ai_assert.h>

#pragma comment (lib, "assimp.lib")

// One instance per aiMesh in the globally loaded asset
MeshHelper** apcMeshes;

const aiScene* scene;

////////////////////////////////////////////////////////////
// Load
////////////////////////////////////////////////////////////
bool Load(const std::string& pFile)
{
	// Create an instance of the Importer class
	Assimp::Importer importer;
	// And have it read the given file with some example postprocessing
	// Usually - if speed is not the most important aspect for you - you'll 
	// propably to request more postprocessing than we do in this example.
	scene = importer.ReadFile(pFile, 
							aiProcess_CalcTangentSpace       | 
							aiProcess_Triangulate            |
							aiProcess_JoinIdenticalVertices  |
							aiProcess_FlipUVs);
  
	// If the import failed, report it
	if(!scene)
	{
		std::string s = "Error in load the " + pFile;
		MessageBoxA(NULL, s.c_str(), "Error mesh", MB_OK);
		return false;
	}

	// allocate a new MeshHelper array and build a new instance
	// for each mesh in the original asset
	apcMeshes = new MeshHelper*[scene->mNumMeshes]();
	for (unsigned int i = 0; i < scene->mNumMeshes;++i)
		apcMeshes[i] = new MeshHelper();

	for (unsigned int i = 0; i < scene->mNumMeshes;++i)
	{
		const aiMesh* mesh = scene->mMeshes[i];

		/*
		// create the material for the mesh
		if (!apcMeshes[i]->piEffect)	{
			CMaterialManager::Instance().CreateMaterial(apcMeshes[i], mesh);
		}
		*/

		// create vertex buffer
		if(FAILED( devices.d3ddev->CreateVertexBuffer(sizeof(Vertex) *
														mesh->mNumVertices,
														D3DUSAGE_WRITEONLY,
														0,
														D3DPOOL_DEFAULT, &apcMeshes[i]->piVB,NULL)))	
		{
			MessageBoxA(NULL, "Failed to create vertex buffer", "Error Model", MB_OK);
			return false;
		}

		DWORD dwUsage = 0;
		if (apcMeshes[i]->piOpacityTexture || 1.0f != apcMeshes[i]->fOpacity)
			dwUsage |= D3DUSAGE_DYNAMIC;

		unsigned int nidx;
		switch (mesh->mPrimitiveTypes) 
		{
			case aiPrimitiveType_POINT:
				nidx = 1;
				break;
			case aiPrimitiveType_LINE:
				nidx = 2;
				break;
			case aiPrimitiveType_TRIANGLE:
				nidx = 3;
				break;
			default: 
				ai_assert(false);
		};


		// check whether we can use 16 bit indices
		if (mesh->mNumFaces * 3 >= 65536)	{
			// create 32 bit index buffer
			if(FAILED( devices.d3ddev->CreateIndexBuffer( 4 * mesh->mNumFaces * nidx,
														D3DUSAGE_WRITEONLY | dwUsage,
														D3DFMT_INDEX32,
														D3DPOOL_DEFAULT, 
														&apcMeshes[i]->piIB,
														NULL)))
			{
				MessageBoxA(NULL, "Failed to create 32 Bit index buffer", "Error model", MB_OK);
				return false;
			}

			// now fill the index buffer
			unsigned int* pbData;
			apcMeshes[i]->piIB->Lock(0,0,(void**)&pbData,0);
			for (unsigned int x = 0; x < mesh->mNumFaces;++x)
			{
				for (unsigned int a = 0; a < nidx;++a)
				{
					*pbData++ = mesh->mFaces[x].mIndices[a];
				}
			}
		}
		else	{
			// create 16 bit index buffer
			if(FAILED( devices.d3ddev->CreateIndexBuffer( 2 * mesh->mNumFaces * nidx,
														D3DUSAGE_WRITEONLY | dwUsage,
														D3DFMT_INDEX16,
														D3DPOOL_DEFAULT,
														&apcMeshes[i]->piIB,
														NULL)))
			{
				MessageBoxA(NULL, "Failed to create 16 Bit index buffer", "Error model", MB_OK);
				return false;
			}

			// now fill the index buffer
			uint16_t* pbData;
			apcMeshes[i]->piIB->Lock(0,0,(void**)&pbData,0);
			for (unsigned int x = 0; x < mesh->mNumFaces;++x)
			{
				for (unsigned int a = 0; a < nidx;++a)
				{
					*pbData++ = (uint16_t)mesh->mFaces[x].mIndices[a];
				}
			}
		}
		apcMeshes[i]->piIB->Unlock();

		// collect weights on all vertices. Quick and careless
		std::vector<std::vector<aiVertexWeight>> weightsPerVertex( mesh->mNumVertices);
		for( unsigned int a = 0; a < mesh->mNumBones; a++)	{
			const aiBone* bone = mesh->mBones[a];
			for( unsigned int b = 0; b < bone->mNumWeights; b++)
				weightsPerVertex[bone->mWeights[b].mVertexId].push_back( aiVertexWeight( a, bone->mWeights[b].mWeight));
		}

		// now fill the vertex buffer
		Vertex* pbData2;
		apcMeshes[i]->piVB->Lock(0,0,(void**)&pbData2,0);
		for (unsigned int x = 0; x < mesh->mNumVertices;++x)
		{
			pbData2->vPosition = mesh->mVertices[x];

			if (NULL == mesh->mNormals)
				pbData2->vNormal = aiVector3D(0.0f,0.0f,0.0f);
			else pbData2->vNormal = mesh->mNormals[x];

			if (NULL == mesh->mTangents)	{
				pbData2->vTangent = aiVector3D(0.0f,0.0f,0.0f);
				pbData2->vBitangent = aiVector3D(0.0f,0.0f,0.0f);
			}
			else	{
				pbData2->vTangent = mesh->mTangents[x];
				pbData2->vBitangent = mesh->mBitangents[x];
			}

			if (mesh->HasVertexColors( 0 ))	{
				using std::min;
				using std::max;
				unsigned char m_argb[4];
				m_argb[0] = (unsigned char)max( min( mesh->mColors[0][x].a * 255.0f, 255.0f), 0.0f);
				m_argb[1] = (unsigned char)max( min( mesh->mColors[0][x].r * 255.0f, 255.0f), 0.0f);
				m_argb[2] = (unsigned char)max( min( mesh->mColors[0][x].g * 255.0f, 255.0f), 0.0f);
				m_argb[3] = (unsigned char)max( min( mesh->mColors[0][x].b * 255.0f, 255.0f), 0.0f);

				pbData2->dColorDiffuse = D3DCOLOR_ARGB(m_argb[0],m_argb[1],m_argb[2],m_argb[3]);
			}
			else pbData2->dColorDiffuse = D3DCOLOR_ARGB(0xFF,0xff,0xff,0xff);

			// ignore a third texture coordinate component
			if (mesh->HasTextureCoords( 0))	{
				pbData2->vTextureUV = aiVector2D(
					mesh->mTextureCoords[0][x].x,
					mesh->mTextureCoords[0][x].y);
			}
			else pbData2->vTextureUV = aiVector2D(0.5f,0.5f);

			if (mesh->HasTextureCoords( 1))	{
				pbData2->vTextureUV2 = aiVector2D(
					mesh->mTextureCoords[1][x].x,
					mesh->mTextureCoords[1][x].y);
			}
			else pbData2->vTextureUV2 = aiVector2D(0.5f,0.5f);

			// Bone indices and weights
			if( mesh->HasBones())	{
				unsigned char boneIndices[4] = { 0, 0, 0, 0 };
				unsigned char boneWeights[4] = { 0, 0, 0, 0 };
				ai_assert( weightsPerVertex[x].size() <= 4);
				for( unsigned int a = 0; a < weightsPerVertex[x].size(); a++)
				{
					boneIndices[a] = weightsPerVertex[x][a].mVertexId;
					boneWeights[a] = (unsigned char) (weightsPerVertex[x][a].mWeight * 255.0f);
				}

				memcpy( pbData2->mBoneIndices, boneIndices, sizeof( boneIndices));
				memcpy( pbData2->mBoneWeights, boneWeights, sizeof( boneWeights));
			} else
			{
				memset( pbData2->mBoneIndices, 0, sizeof( pbData2->mBoneIndices));
				memset( pbData2->mBoneWeights, 0, sizeof( pbData2->mBoneWeights));
			}

			++pbData2;
		}
		apcMeshes[i]->piVB->Unlock();

		// now generate the second vertex buffer, holding all normals
		if (!apcMeshes[i]->piVBNormals)	{
			//GenerateNormalsAsLineList(apcMeshes[i],mesh);

			ai_assert(NULL != apcMeshes[i]);
			ai_assert(NULL != mesh);

			if (!mesh->mNormals)return 0;

			// create vertex buffer
			if(FAILED( devices.d3ddev->CreateVertexBuffer(sizeof(LineVertex) *
														mesh->mNumVertices * 2,
														D3DUSAGE_WRITEONLY,
														LineVertex::GetFVF(),
														D3DPOOL_DEFAULT, &apcMeshes[i]->piVBNormals,NULL)))
			{
				MessageBoxA(NULL,"Failed to create vertex buffer for the normal list", "Error model", MB_OK);
				return false;
			}

			// now fill the vertex buffer with data
			LineVertex* pbData2;
			apcMeshes[i]->piVBNormals->Lock(0,0,(void**)&pbData2,0);
			for (unsigned int x = 0; x < mesh->mNumVertices;++x)
			{
				pbData2->vPosition = mesh->mVertices[x];

				++pbData2;

				aiVector3D vNormal = mesh->mNormals[x];
				vNormal.Normalize();

				// scalo with the inverse of the world scaling to make sure
				// the normals have equal length in each case
				// TODO: Check whether this works in every case, I don't think so
				//vNormal.x /= g_mWorld.a1*4;
				//vNormal.y /= g_mWorld.b2*4;
				//vNormal.z /= g_mWorld.c3*4;

				pbData2->vPosition = mesh->mVertices[x] + vNormal;

				++pbData2;
			}
			apcMeshes[i]->piVBNormals->Unlock();
		}
	}

	// get the number of vertices/faces in the model
	unsigned int iNumVert = 0;
	unsigned int iNumFaces = 0;
	for (unsigned int i = 0; i < scene->mNumMeshes;++i)
	{
		iNumVert += scene->mMeshes[i]->mNumVertices;
		iNumFaces += scene->mMeshes[i]->mNumFaces;
	}

	meshrender = new MeshRender(devices, scene, apcMeshes);

	return true;
}

////////////////////////////////////////////////////////////
// Render
////////////////////////////////////////////////////////////
void Render(bool wireframe)
{
	if(wireframe) 
	{
		devices.d3ddev->SetRenderState(D3DRS_FILLMODE, D3DFILL_WIREFRAME);
	}

	aiMatrix4x4 m;
	RenderNode(scene->mRootNode, true);
	// render all child nodes
	for (unsigned int i = 0; i < scene->mRootNode->mNumChildren;++i)
		RenderNode(scene->mRootNode->mChildren[i], true );

	if(wireframe) 
	{
		devices.d3ddev->SetRenderState(D3DRS_FILLMODE, D3DFILL_SOLID);
	}
}

////////////////////////////////////////////////////////////
// Render a single node
////////////////////////////////////////////////////////////
bool RenderNode(aiNode* piNode, bool bAlpha)
{
	aiMatrix4x4 aiMe;
	aiMatrix4x4* piMatrix = new aiMatrix4x4();

	for (unsigned int i = 0; i < piNode->mNumMeshes;++i)
	{
		if (bAlpha)
			meshrender->DrawSorted(piNode->mMeshes[i], aiMe);
		else 
			meshrender->DrawUnsorted(piNode->mMeshes[i]);
	}


	return true;
}

////////////////////////////////////////////////////////////
// Release
////////////////////////////////////////////////////////////
void Release()
{
	bool bNoMaterials = false;
	if (!scene)
		return;

	// TODO: Move this to a proper destructor
	for (unsigned int i = 0; i < scene->mNumMeshes;++i)
	{
		if(apcMeshes[i]->piVB)
		{
			apcMeshes[i]->piVB->Release();
			apcMeshes[i]->piVB = NULL;
		}
		if(apcMeshes[i]->piVBNormals)
		{
			apcMeshes[i]->piVBNormals->Release();
			apcMeshes[i]->piVBNormals = NULL;
		}
		if(apcMeshes[i]->piIB)
		{
			apcMeshes[i]->piIB->Release();
			apcMeshes[i]->piIB = NULL;
		}

		if (!bNoMaterials)
		{
			if(apcMeshes[i]->piEffect)
			{
				apcMeshes[i]->piEffect->Release();
				apcMeshes[i]->piEffect = NULL;
			}
			if(apcMeshes[i]->piDiffuseTexture)
			{
				apcMeshes[i]->piDiffuseTexture->Release();
				apcMeshes[i]->piDiffuseTexture = NULL;
			}
			if(apcMeshes[i]->piNormalTexture)
			{
				apcMeshes[i]->piNormalTexture->Release();
				apcMeshes[i]->piNormalTexture = NULL;
			}
			if(apcMeshes[i]->piSpecularTexture)
			{
				apcMeshes[i]->piSpecularTexture->Release();
				apcMeshes[i]->piSpecularTexture = NULL;
			}
			if(apcMeshes[i]->piAmbientTexture)
			{
				apcMeshes[i]->piAmbientTexture->Release();
				apcMeshes[i]->piAmbientTexture = NULL;
			}
			if(apcMeshes[i]->piEmissiveTexture)
			{
				apcMeshes[i]->piEmissiveTexture->Release();
				apcMeshes[i]->piEmissiveTexture = NULL;
			}
			if(apcMeshes[i]->piOpacityTexture)
			{
				apcMeshes[i]->piOpacityTexture->Release();
				apcMeshes[i]->piOpacityTexture = NULL;
			}
			if(apcMeshes[i]->piShininessTexture)
			{
				apcMeshes[i]->piShininessTexture->Release();
				apcMeshes[i]->piShininessTexture = NULL;
			}
		}
	}
}