// Copyright (C) 2014-2015 ChaosForge Ltd
// http://chaosforge.org/
//
// This file is part of Nova libraries. 
// For conditions of distribution and use, see copying.txt file in root folder.

#include "nv/formats/assimp_loader.hh"

#include "nv/interface/data_channel_access.hh"
#include "nv/stl/hash_store.hh"
#include "nv/lib/assimp.hh"

using namespace nv;

namespace nv {

	struct assimp_data
	{
		const aiScene* scene;
		vector< const aiBone* > bones;
		vector< const aiNode* > nodes;
		vector< const aiMesh* > meshes;
		vector< const aiNode* > skeletons;

		hash_store< shash64, const aiBone* > bone_by_name;
		hash_store< shash64, const aiNode* > node_by_name;
	};

}

const unsigned MAX_BONES = 64;

struct assimp_plain_vtx 
{
	vec3 position;
	vec3 normal;
	vec2 texcoord;
	vec4 tangent;

	assimp_plain_vtx() {}
	assimp_plain_vtx( const vec3& v, const vec2& t, const vec3& n, const vec4& g )
	{
		position = v;
		texcoord = t;
		normal   = n;
		tangent  = g;
	}
};

struct assimp_skinned_vtx 
{
	vec3  position;
	vec3  normal;
	vec2  texcoord;
	vec4  tangent;
	ivec4 boneindex;
	vec4  boneweight;

	assimp_skinned_vtx() {}
	assimp_skinned_vtx( const vec3& v, const vec2& t, const vec3& n, const vec4& g )
	{
		position = v;
		texcoord = t;
		normal   = n;
		tangent  = g;
	}
};

struct assimp_key_p  { float time; vec3 translation; };
struct assimp_key_r  { float time; quat rotation; };
struct assimp_key_s  { float time; vec3 scale; };
struct assimp_key_tr { transform tform; };


nv::assimp_loader::assimp_loader( string_table* strings, const string_view& a_ext, uint32 a_assimp_flags /*= 0 */ )
	: mesh_loader( strings ), m_mesh_count(0)
{
	m_ext   = a_ext;
	m_assimp_flags = a_assimp_flags;
	if ( m_assimp_flags == 0 )
	{
		m_assimp_flags = ( 
			aiProcess_CalcTangentSpace				|  
			aiProcess_GenSmoothNormals				|  
			aiProcess_JoinIdenticalVertices			|   
			aiProcess_ImproveCacheLocality			|  
			aiProcess_LimitBoneWeights				|  
			aiProcess_RemoveRedundantMaterials      |  
			aiProcess_SplitLargeMeshes				|  
			aiProcess_Triangulate					|  
			aiProcess_GenUVCoords                   |  
			aiProcess_SortByPType                   |  
			aiProcess_FindDegenerates               |  
			aiProcess_FindDegenerates |
			0 );
	}
	m_data = new assimp_data;
	m_data->scene = nullptr;
}


bool nv::assimp_loader::load( stream& source )
{
	load_assimp_library();
	if ( m_data->scene != nullptr ) aiReleaseImport( m_data->scene );
	m_data->scene = nullptr;
	m_mesh_count = 0;
	NV_LOG_NOTICE( "AssImp loading file..." );
	size_t size = source.size();
	char* data  = new char[ size ];
	source.read( data, size, 1 );
	const aiScene* scene = aiImportFileFromMemory( data, size, m_assimp_flags, m_ext.data() );

	if( !scene)
	{
		NV_LOG_ERROR( aiGetErrorString() );
		return false;
	}
	m_data->scene = scene;
	m_mesh_count  = scene->mNumMeshes;
	NV_LOG_NOTICE( "Loading successfull" );

	scan_nodes( scene->mRootNode );

	for ( unsigned i = 0; i < scene->mRootNode->mNumChildren; ++i )
	{
		const aiNode* node = scene->mRootNode->mChildren[i];
		if ( node->mNumMeshes == 0 )
			m_data->skeletons.push_back( node );
	}

	for ( nv::uint32 i = 0; i < m_mesh_count; i++ )
	{
 		data_node_info info;
		data_channel_set* mdata = data_channel_set_creator::create_set( 2 );
		load_mesh_data( mdata, i, info );
		m_meshes.push_back( mdata );
		m_mesh_info.push_back( info );
	}

	scene_report();
	return true;
}

data_channel_set* nv::assimp_loader::release_mesh_data( size_t index, data_node_info& info )
{
	if ( index >= m_mesh_count ) return nullptr;
	info = m_mesh_info[index];
	return m_meshes[index];
}
void nv::assimp_loader::load_mesh_data( data_channel_set* data, size_t index, data_node_info& info )
{
	const aiMesh*  mesh  = m_data->scene->mMeshes[ index ];

	bool skinned = mesh->mNumBones > 0;

	data_descriptor desc;
	if ( skinned )
		desc.initialize< assimp_skinned_vtx >();
	else
		desc.initialize< assimp_plain_vtx >();
	data_channel_set_creator maccess( data );
	string64 name( mesh->mName.data, mesh->mName.length );
	if ( mesh->mName.length == 0 )
	{
		for ( auto node : m_data->nodes )
		{
			if ( node->mMeshes )
				for ( uint32 i = 0; i < node->mNumMeshes; ++i )
					if ( node->mMeshes[i] == index )
					{
						name.assign( node->mName.data, node->mName.length );
						if ( i != 0 )
						{
							name.append( "#0" );
							name.append( i );
						}
					}

		}
	}
	info.name = make_name( name );
	info.parent_id = -1;
	int hack_for_node_anim;
	if ( is_node_animated() )
		info.parent_id = sint16( index );


	uint8*  cdata   = maccess.add_channel( desc, mesh->mNumVertices ).raw_data();
	uint32* indices = reinterpret_cast<uint32*>( maccess.add_channel< index_u32 >( mesh->mNumFaces * 3 ).raw_data() );

	if ( mesh->mTangents && mesh->mBitangents )
		for ( unsigned int i = 0; i < mesh->mNumVertices; i++ )
		{
			vec3 v = assimp_vec3_cast( mesh->mVertices[i] );
			vec3 n = math::normalize( assimp_vec3_cast( mesh->mNormals[i] ) );
			vec3 t = math::normalize( assimp_vec3_cast( mesh->mTangents[i] ) );
			vec3 b = math::normalize( assimp_vec3_cast( mesh->mBitangents[i] ) );
			vec2 s = assimp_st_cast( mesh->mTextureCoords[0][i] );

			vec3 t_i = math::normalize( t - n * math::dot( n, t ) );
			float det = ( math::dot( math::cross( n, t ), b ) );
			det = ( det < 0.0f ? -1.0f : 1.0f );
			nv::vec4 vt( t_i[0], t_i[1], t_i[2], det );
			if ( skinned )
				reinterpret_cast<assimp_skinned_vtx*>( cdata )[i] = assimp_skinned_vtx( v, s, n, vt );
			else
				reinterpret_cast<assimp_plain_vtx*>( cdata )[i] = assimp_plain_vtx( v, s, n, vt );
		}

	if ( skinned )
	{
		assimp_skinned_vtx* vtx = reinterpret_cast< assimp_skinned_vtx* >( cdata );
		for (unsigned int m=0; m<mesh->mNumBones; m++)
		{
			aiBone* bone  = mesh->mBones[m];
			for ( size_t w=0; w<bone->mNumWeights; w++)
			{
				assimp_skinned_vtx& v = vtx[ bone->mWeights[w].mVertexId ];
				bool found = false;
				for ( size_t i = 0 ; i < 4; ++i )
				{
					if ( v.boneweight[i] <= 0.0f ) 
					{
						v.boneindex[i]  = int( m );
						v.boneweight[i] = bone->mWeights[w].mWeight;
						found = true;
						break;
					}
				}
				NV_ASSERT( found, "Too many weights!" );
			}
		}
	}

	for (unsigned int i=0; i<mesh->mNumFaces; i++)
	{
		const aiFace* face = &mesh->mFaces[i];
		for (unsigned int j=0; j<face->mNumIndices; j++)
		{
			indices[ i*3 + j ] = uint32( face->mIndices[j] );
		}
	}

}

nv::assimp_loader::~assimp_loader()
{
	if ( m_data->scene != nullptr ) aiReleaseImport( m_data->scene );
	delete m_data;
}

void nv::assimp_loader::scene_report() const
{
	if ( m_data->scene == nullptr ) return;

	NV_LOG_NOTICE( "------------------------" );
	NV_LOG_NOTICE( "Texture   count - ", m_data->scene->mNumTextures );
	NV_LOG_NOTICE( "Animation count - ", m_data->scene->mNumAnimations );
	NV_LOG_NOTICE( "Material  count - ", m_data->scene->mNumMaterials );
	NV_LOG_NOTICE( "Meshes    count - ", m_data->scene->mNumMeshes );
	NV_LOG_NOTICE( "------------------------" );

	aiNode* root = m_data->scene->mRootNode;
	if (root)
	{
		NV_LOG_NOTICE( "Root node  - ", root->mName.data );
		NV_LOG_NOTICE( "  meshes   - ", root->mNumMeshes );
		NV_LOG_NOTICE( "  children - ", root->mNumChildren );
	}
	else
	{
		NV_LOG_NOTICE( "No root node!" );
	}
	NV_LOG_NOTICE( "------------------------" );

	if ( m_data->scene->mNumMeshes > 0 )
	{
		for ( nv::uint32 mc = 0; mc < m_data->scene->mNumMeshes; mc++ )
		{
			aiMesh* mesh = m_data->scene->mMeshes[mc];

			NV_LOG_NOTICE( "Mesh #", mc, "   - ", string_view( static_cast<char*>( mesh->mName.data ), mesh->mName.length ) );
			NV_LOG_NOTICE( "  bones   - ", mesh->mNumBones );
			NV_LOG_NOTICE( "  uvs     - ", mesh->mNumUVComponents[0] );
			NV_LOG_NOTICE( "  verts   - ", mesh->mNumVertices );
			NV_LOG_NOTICE( "  faces   - ", mesh->mNumFaces );

			// 			NV_LOG_NOTICE(  "Bones:" );
			// 			for (unsigned int m=0; m<mesh->mNumBones; m++)
			// 			{
			// 				aiBone* bone  = mesh->mBones[m];
			// 				NV_LOG_NOTICE( bone->mName.C_Str() );
			// 			}
		}
	}
	else
	{
		NV_LOG_NOTICE( "No meshes!" );
	}
	NV_LOG_NOTICE( "------------------------" );

	for ( auto node : m_data->nodes )
	{
		NV_LOG_NOTICE( "Node : ", string_view( node->mName.data, node->mName.length ) );
	}

	for ( auto skeleton : m_data->skeletons )
	{
		NV_LOG_NOTICE( "Skeleton : ", string_view( skeleton->mName.data, skeleton->mName.length ) );
	}

	// 	if ( scene->mNumMaterials > 0 )
	// 	{
	// 		for (unsigned int m=0; m < scene->mNumMaterials; m++)
	// 		{
	// 			int texIndex = 0;
	// 			aiReturn texFound = aiReturn_SUCCESS;
	// 			aiString path;	// filename
	//  			while (texFound == aiReturn_SUCCESS)
	//  			{
	// 				texFound = scene->mMaterials[m]->GetTexture(aiTextureType_DIFFUSE, texIndex, &path);
	// 				NV_LOG_NOTICE( "  material - ", path.data );
	// 				texIndex++;
	// 			}
	// 		}
	// 	}
	// 	else
	// 	{
	// 		NV_LOG_NOTICE( "No materials" );
	// 	}
	// 	NV_LOG_NOTICE( "------------------------" );

}

bool nv::assimp_loader::is_node_animated()
{
	return is_animated() && m_data->skeletons.empty();
}

void nv::assimp_loader::build_skeleton( vector< data_node_info >& skeleton, const void* node, sint16 parent_id )
{
	const aiNode* ainode = reinterpret_cast<const aiNode*>( node );

	if ( ainode->mNumMeshes > 0 )
	{
		int error;
		int bug_this_works_only_if_before_releasing_meshes;

		nv::uint32 mid = ainode->mMeshes[0];
		m_mesh_info[mid].parent_id = parent_id;
		return;
	}

	string_view name( ainode->mName.data, ainode->mName.length );
	if ( name.starts_with( '_' ) ) return;

	data_node_info info;
	info.name      = make_name( name );
	info.parent_id = parent_id;

	sint16 this_id = sint16( skeleton.size() );
	skeleton.push_back( info );
	for ( unsigned i = 0; i < ainode->mNumChildren; ++i )
	{
		build_skeleton( skeleton, ainode->mChildren[i], this_id );
	}
}

data_node_list* nv::assimp_loader::release_merged_bones()
{
	if ( m_data->skeletons.empty() ) return nullptr;
	vector< data_node_info > bone_data;
	hash_store< shash64, uint16 > bone_map;

	{
		const aiNode* skeleton = m_data->skeletons[0];
		build_skeleton( bone_data, skeleton, -1 );

		for ( uint32 i = 0; i < bone_data.size(); ++i )
		{
			bone_map[bone_data[i].name] = uint16(i);
		}
	}

	for ( unsigned int m = 0; m < m_mesh_count; ++m )
	{
		uint16 translate[MAX_BONES];
		vector< data_node_info > bones;
		const aiMesh*  mesh = m_data->scene->mMeshes[m];
		if ( mesh->mNumBones != 0 )
		{
			for ( unsigned int b = 0; b < mesh->mNumBones; b++ )
			{
				aiBone* bone = mesh->mBones[b];
				mat4    offset = assimp_mat4_cast( bone->mOffsetMatrix );
				shash64 bone_name( bone->mName.data );

				int remove_this;

				NV_ASSERT( bone_map.find( shash64( bone->mName.data ) ) != bone_map.end(), "BONE NOT FOUND!" );
				uint16 index = bone_map[bone_name];
				bone_data[index].transform = offset;
				translate[b] = index;
			}

			data_channel_access< assimp_skinned_vtx > channel( const_cast<raw_data_channel*>( m_meshes[m]->get_channel( 0 ) ) );
			for ( unsigned v = 0; v < channel.size(); ++v )
			{
				assimp_skinned_vtx& vertex = channel.data()[v];
				for ( size_t i = 0; i < 4; ++i )
				{
					if ( vertex.boneweight[i] > 0.0f )
					{
						vertex.boneindex[i] = int( translate[vertex.boneindex[i]] );
					}
				}
			}

		}
	}

	for ( uint32 i = 0; i < bone_data.size(); ++i )
	{
		int error; // not really, just
		int check_this_shit;
		if ( bone_data[i].transform == mat4() )
		{
			mat4 tr = nv::math::inverse( assimp_mat4_cast( m_data->node_by_name[bone_data[i].name]->mTransformation ) );
			int pid = bone_data[i].parent_id;
			if ( pid >= 0 )
				bone_data[i].transform = tr * bone_data[ size_t( pid ) ].transform;
			else
				bone_data[i].transform = tr;
		}
//		list->append( bone_data[i] );
	}


	data_node_list* list = new data_node_list( make_name( "bones" ) );
	for ( uint32 i = 0; i < bone_data.size(); ++i )
	{
		list->append( bone_data[i] );
	}

	return list;
}

mesh_nodes_data* nv::assimp_loader::release_mesh_nodes_data( size_t index /*= 0*/ )
{
	if ( m_data->scene == nullptr ) return nullptr;
	const aiScene* scene = reinterpret_cast<const aiScene*>( m_data->scene );
	if ( scene->mRootNode == nullptr || scene->mAnimations == nullptr || scene->mAnimations[index] == nullptr) return nullptr;

	const aiAnimation* anim = scene->mAnimations[index];

	uint32 count = count_nodes( scene->mRootNode );
	count = count - 1;

	uint16 frame_rate     = static_cast<uint16>( anim->mTicksPerSecond );
	uint16 duration       = static_cast<uint16>( anim->mDuration )+1;

	data_channel_set** temp  = new data_channel_set*[ count ];
	data_node_info*    temp2 = new data_node_info[count ];
	array_ref< data_channel_set* > temp_ref( temp, count );
	array_ref< data_node_info >    temp2_ref( temp2, count );

	nv::sint16 next = 0;
	for ( unsigned i = 0; i < scene->mRootNode->mNumChildren; ++i )
	{
		next = load_node( index, temp_ref, temp2_ref, scene->mRootNode->mChildren[i], next, -1 );
	}
//	load_node( index, temp_ref, temp2_ref, scene->mRootNode, 0, -1 );

	mesh_nodes_data* result = new mesh_nodes_data( make_name( static_cast<const char*>( anim->mName.data ) ), frame_rate, duration );
	for ( nv::uint32 i = 0; i < count; ++i )
	{
		result->append( temp_ref[i], temp2_ref[i] );
	}
	result->initialize();
	delete temp;
	delete temp2;
	return result;
}

data_node_list* nv::assimp_loader::release_data_node_list( size_t /*= 0 */ )
{
	return release_merged_bones();
}

bool nv::assimp_loader::is_animated( size_t /*= 0 */ )
{
	int this_is_incorrect;
	return m_mesh_count == 0 || ( m_data->scene->mNumAnimations > 0 && m_data->skeletons.size() == 0 );
}

void nv::assimp_loader::scan_nodes( const void* node ) const
{
	const aiNode* ainode = reinterpret_cast<const aiNode*>( node );
	m_data->nodes.push_back( ainode );
	m_data->node_by_name[shash64(ainode->mName.data)] = ainode;

	for ( unsigned i = 0; i < ainode->mNumChildren; ++i )
	{
		scan_nodes( ainode->mChildren[i] );
	}
}

nv::uint32 nv::assimp_loader::count_nodes( const void* node ) const
{
	const aiNode* ainode = reinterpret_cast< const aiNode* >( node );
	nv::uint32 count = 1;
	for ( unsigned i = 0; i < ainode->mNumChildren; ++i )
	{
		count += count_nodes( ainode->mChildren[i] );
	}
	return count;
}

nv::sint16 nv::assimp_loader::load_node( uint32 anim_id, array_ref< data_channel_set* > nodes, array_ref< data_node_info > infos, const void* vnode, sint16 this_id, sint16 parent_id )
{
	const aiNode*  node  = reinterpret_cast<const aiNode*>( vnode );
	string_view name( static_cast< const char* >( node->mName.data ) );
	const aiAnimation* anim  = m_data->scene->mAnimations[anim_id];
	const aiNodeAnim*  anode = nullptr;

	for ( unsigned i = 0 ; i < anim->mNumChannels ; i++ )
	{
		anode = anim->mChannels[i];
		if ( string_view( static_cast< const char* >( anode->mNodeName.data ) ) == name ) break;
		anode = nullptr;
	}


	transform t = nv::transform( nv::assimp_mat4_cast( node->mTransformation ) );

	nodes[ uint32( this_id ) ] = anode ? create_keys( anode, t ) : data_channel_set_creator::create_set( 0 );
	infos[ uint32( this_id ) ].name      = make_name( name );
	infos[ uint32( this_id ) ].parent_id = parent_id;
	// This value is ignored by the create_transformed_keys, but needed by create_direct_keys!
	// TODO: find a common solution!
	//       This is bad because create_transformed_keys never uses node-transformations for
	//       node's without keys
	// TODO: this can probably be deleted
//   	infos[this_id].transform = nv::assimp_mat4_cast( node->mTransformation );
//   	if ( this_id == 0 ) infos[this_id].transform = mat4();


	nv::sint16 next = this_id + 1;
	for ( unsigned i = 0; i < node->mNumChildren; ++i )
	{
		next = load_node( anim_id, nodes, infos, node->mChildren[i], next, this_id );
	}

	return next;
}

data_channel_set* nv::assimp_loader::create_keys( const void* vnode, const transform& tr )
{
	const aiNodeAnim* node  = reinterpret_cast< const aiNodeAnim* >( vnode );
	if ( node->mNumPositionKeys == 0 && node->mNumRotationKeys == 0 && node->mNumScalingKeys == 0 )
	{
		return data_channel_set_creator::create_set( 0 );
	}

/* OLD
	data_channel_set* set = data_channel_set_creator::create_set( 2 );
	data_channel_set_creator key_set( set );

	assimp_key_p* pchannel = key_set.add_channel< assimp_key_p >( node->mNumPositionKeys ).data();
	assimp_key_r* rchannel = key_set.add_channel< assimp_key_r >( node->mNumRotationKeys ).data();
//	assimp_key_s* schannel = key_set.add_channel< assimp_key_s >( node->mNumScalingKeys ).data();

	for ( unsigned np = 0; np < node->mNumPositionKeys; ++np )
	{
		pchannel[np].time        = static_cast<float>( node->mPositionKeys[np].mTime );
		pchannel[np].translation = assimp_vec3_cast(node->mPositionKeys[np].mValue);
	}
	for ( unsigned np = 0; np < node->mNumRotationKeys; ++np )
	{
		rchannel[np].time     = static_cast<float>( node->mRotationKeys[np].mTime );
		rchannel[np].rotation = assimp_quat_cast(node->mRotationKeys[np].mValue );
	}
	*/
	data_channel_set* set = data_channel_set_creator::create_set( 1 );
	data_channel_set_creator key_set( set );

	assimp_key_tr* channel = key_set.add_channel< assimp_key_tr >( node->mNumPositionKeys ).data();
	for ( unsigned np = 0; np < node->mNumPositionKeys; ++np )
	{
		channel[np].tform.set_position( assimp_vec3_cast( node->mPositionKeys[np].mValue ) );
		channel[np].tform.set_orientation( assimp_quat_cast( node->mRotationKeys[np].mValue ) );
		if ( is_node_animated() )
			channel[np].tform = tr.inverse() * channel[np].tform ;
	}

// 	if ( node->mNumScalingKeys > 0 )
// 	{
// 		vec3 scale_vec0 = assimp_vec3_cast( node->mScalingKeys[0].mValue );
// 		float scale_value   = math::length( math::abs( scale_vec0 - vec3(1,1,1) ) );
// 		if ( node->mNumScalingKeys > 1 || scale_value > 0.001 ) 
// 		{
// 			NV_LOG( nv::LOG_WARNING, "scale key significant!" );
// 			for ( unsigned np = 0; np < node->mNumRotationKeys; ++np )
// 			{
// 				schannel[np].time  = (float)node->mScalingKeys[np].mTime;
// 				schannel[np].scale = assimp_vec3_cast(node->mScalingKeys[np].mValue);
// 			}
// 		}
// 		else 
// 		{
// 			schannel[0].time  = (float)node->mScalingKeys[0].mTime;
// 			schannel[0].scale = assimp_vec3_cast(node->mScalingKeys[0].mValue);
// 		}
// 	}
	return set;
}

// mesh_data_pack* nv::assimp_loader::release_mesh_data_pack()
// {
// 	if ( m_scene == nullptr || m_mesh_count == 0 ) return nullptr;
// 	const aiScene* scene = reinterpret_cast<const aiScene*>( m_scene );
// 	bool has_bones = false;
// 	data_channel_set* meshes = data_channel_set_creator::create_set_array( m_mesh_count, 2 );
// 	for ( size_t m = 0; m < m_mesh_count; ++m )
// 	{
// 		const aiMesh* mesh = scene->mMeshes[ m ];
// 		data_channel_set_creator( &meshes[m] ).set_name( make_name( static_cast<const char*>( mesh->mName.data ) ) );
// 		if ( mesh->mNumBones > 0 ) has_bones = true;
// 		load_mesh_data(&meshes[m],m);
// 	}
// 
// 	mesh_nodes_data* nodes = ( has_bones ? release_merged_bones( meshes ) : release_mesh_nodes_data(0) );
// 	return new mesh_data_pack( m_mesh_count, meshes, nodes );
// }

nv::size_t nv::assimp_loader::get_nodes_data_count() const
{
	if ( m_data->scene == nullptr ) return 0;
	return m_data->scene->mNumAnimations;
}