#include "stdafx.h"
#include "AssimpReader.h"

#include <LibQtGeoViewerCore/SceneMain.h>

#include "../FileUtil.h"

#include <C2/util/string_util.h>
#include <C2/lm/range1.h>
#include <LibGeo/Path.h>

#include <assimp/Importer.hpp>
#include <assimp/scene.h>
#include <assimp/postprocess.h>

#include "AssimpUtil.h"

#include <QFileInfo>
#include <sstream>



AssimpReader::AssimpReader(geom::GeomFileFormat fmt_)
{
	fmt = fmt_;
}

std::string AssimpReader::GetIndexedName(const std::string& base, int idx)
{
	std::ostringstream oss;
	oss << base << "-" << idx;
	return oss.str();
}

geom::GeomObject* AssimpReader::LoadGeom(SceneMain& scene, const std::string& filename)
{
	unsigned int ai_option =
		aiProcess_JoinIdenticalVertices;

	Assimp::Importer importer;
	const aiScene* ai_scene = importer.ReadFile(filename, ai_option);
	if (!ai_scene)
		return NULL;

	//PlotASNode(ai_scene->mRootNode);

	std::string dirpath = lib_geo::Path::GetParentDirPath(filename);
	std::string title = FileUtil::GetFileTitle(filename);

	GeomObject* geom = new GeomObject();
	geom->m_Name = title;

	for (unsigned int i = 0; i < ai_scene->mNumMeshes; ++i)
	{
		aiMesh* mesh_src = ai_scene->mMeshes[i];
		MeshBuf* mbuf = geom->CreateNewMeshBuf();

		std::string name = mesh_src->mName.C_Str();
		if (name.empty())
			name = GetIndexedName(title, i);

		mbuf->m_Name = name;

		ReadMeshShape(mesh_src, mbuf->m_Mesh);

		AssimpUtil::CopyMaterials(mbuf, ai_scene, scene.m_TexConfig, dirpath);

		ResetBone(mbuf, mesh_src);

		mbuf->m_Bones.m_SrcVertPos = mbuf->m_Mesh.m_Verts;
	}

	SetNodeTree(ai_scene, geom);

	geom::NodeMap nodemap;
	AssimpUtil::CreateBoneTree(nodemap, &geom->m_GTree.GetRoot());

	SetBoneNodeRef(nodemap, geom);

	LoadAnimation(nodemap, ai_scene, geom, scene.m_IOConfig);

	geom->SetFrameShape(scene.m_CurFrame);

	geom->InitializeBufferCommon();

	geom->InitAllBBox();

	geom->m_FileFormat = fmt;
	geom->m_FilePath   = filename;

	//geom->m_Nodes.TraceNodeTree();

	return geom;
}

bool AssimpReader::LoadCamera(SceneMain& scene, const std::string& filename)
{
	unsigned int ai_option =
		aiProcess_JoinIdenticalVertices;

	Assimp::Importer importer;
	const aiScene* ai_scene = importer.ReadFile(filename, ai_option);
	if (!ai_scene)
		return false;

	if (ai_scene->mNumAnimations <= 0)
		return false;
	if (ai_scene->mNumCameras <= 0)
		return false;

	aiNodeAnim* an = NULL;
	aiCamera* cam = NULL;
	const aiAnimation* anim = ai_scene->mAnimations[0];
	for (unsigned int i = 0; i < ai_scene->mNumCameras; ++i)
	{
		aiCamera* camera = ai_scene->mCameras[i];

		for (unsigned int j = 0; j < anim->mNumChannels; ++j)
		{
			aiNodeAnim* ch = anim->mChannels[j];
			if (ch->mNodeName == camera->mName)
			{
				an = ch;
				cam = camera;
				break;
			}
		}

		if (an != NULL)
			break;
	}

	if (an == NULL)
		return false;

	gl::Camera& camera = scene.m_Camera;

	ae::ConvertVec(camera.m_SeqPos, cam->mPosition);
	ae::ConvertVec(camera.m_SeqUp, cam->mUp);
	ae::ConvertVec(camera.m_SeqLook, cam->mLookAt);

	camera.m_Projection.m_Near = cam->mClipPlaneNear;
	camera.m_Projection.m_Far = cam->mClipPlaneFar;

	ReadKeys(an, camera.m_SeqTrans, anim, scene.m_IOConfig);

	camera.SetFrame(scene.m_CurFrame);

	scene.UpdateTransform();
	scene.ReportDoneEditGeometry();

	return true;
}

void AssimpReader::ReadMeshShape(aiMesh* mesh_src, lib_geo::BaseMesh& mesh_dst)
{
	AssimpUtil::CopyAIVerts(mesh_dst, mesh_src);
	AssimpUtil::CopyAIFaces(mesh_dst, mesh_src);

	if (!mesh_dst.HasNormal())
		mesh_dst.CreateNormalsEachVerts(true);

}

void AssimpReader::SetNodeTree(const aiScene* ai_scene, geom::GeomObject* geom)
{
	GeoNodeTree& nodes = geom->m_GTree;
	nodes.ResetRoot();

	const aiNode* root = ai_scene->mRootNode;
	SceneNode* n = &nodes.GetRoot();

	SetNodeTreeSub(ai_scene, root, geom, n);
}

void AssimpReader::SetNodeTreeSub(const aiScene* ai_scene, const aiNode* an, geom::GeomObject* geom, geom::SceneNode* gn)
{
	gn->m_Name = an->mName.C_Str();
	gn->m_MeshIDs.resize(an->mNumMeshes);
	for (unsigned int j = 0; j < an->mNumMeshes; ++j)
	{
		int mid = an->mMeshes[j];
		gn->m_MeshIDs[j] = mid;

		MeshBuf& mb = geom->m_MeshAry[mid];
		if (mb.m_Name.empty())
		{
			mb.m_Name = an->mName.C_Str();
		}
	}

	ae::ConvertMat(gn->m_Transform, an->mTransformation);

	for (unsigned int i = 0; i < an->mNumChildren; ++i)
	{
		aiNode* n = an->mChildren[i];
		SceneNode* gcn = new SceneNode();
		gn->m_Children.push_back(gcn);
		gcn->m_Parent = gn;
		SetNodeTreeSub(ai_scene, n, geom, gcn);
	}
}

void AssimpReader::ResetBone(MeshBuf* mbuf, aiMesh* mesh)
{
	mbuf->m_Bones.m_Bones.resize(mesh->mNumBones);
	for (unsigned int j = 0; j < mesh->mNumBones; ++j)
	{
		const aiBone* s_bone = mesh->mBones[j];
		Bone& d_bone = mbuf->m_Bones.m_Bones[j];
		d_bone.m_Name = s_bone->mName.C_Str();

		ConvertBoneWeight(s_bone, d_bone);

		lm::matrix4f offset;
		ae::ConvertMat(offset, s_bone->mOffsetMatrix);
		d_bone.SetOffset(offset);
	}
}

void AssimpReader::ConvertBoneWeight(const aiBone* s_bone, geom::Bone& d_bone)
{
	d_bone.m_Weights.resize(s_bone->mNumWeights);
	for (unsigned int k = 0; k < s_bone->mNumWeights; ++k)
	{
		BoneWeight& bw = d_bone.m_Weights[k];
		aiVertexWeight& vw = s_bone->mWeights[k];
		bw.m_Vid    = vw.mVertexId;
		bw.m_Weight = vw.mWeight;
	}
}

void AssimpReader::SetBoneNodeRef(geom::NodeMap& nodemap, geom::GeomObject* geom)
{
	for (MeshBuf& m : geom->m_MeshAry)
	{
		for (Bone& b : m.m_Bones.m_Bones)
		{
			b.m_RefNode = nodemap.Find(b.m_Name);
		}
	}
}

void AssimpReader::LoadAnimation(
	geom::NodeMap& nodemap,
	const aiScene* ai_scene,
	geom::GeomObject* geom,
	IOConfig& conf)
{
	if (!geom->m_GTree.HasRoot())
		return;

	if (ai_scene->mNumAnimations <= 0)
		return;

	const aiAnimation* anim = ai_scene->mAnimations[0];
	for (unsigned int j = 0; j < anim->mNumChannels; ++j)
	{
		aiNodeAnim* ch = anim->mChannels[j];
		SceneNode* sn = nodemap.Bones[ch->mNodeName.C_Str()];
		if (sn == NULL)
			continue;

		geom->UpdateBoneToNodeLink(sn);

		ReadKeys(ch, sn->m_SeqTrans, anim, conf);
	}
}

void AssimpReader::ReadKeys(
	aiNodeAnim* ch,
	geom::SeqTrans& st,
	const aiAnimation* anim,
	IOConfig& conf)
{
	if (conf.InterpolateAnimation)
		AssimpUtil::ReadKeys_Time(ch, st, anim);
	else
		AssimpUtil::ReadKeys_Raw(ch, st);
}

void AssimpReader::PlotASNode(aiNode* n)
{
	static int layer = 0;

	{
		std::ostringstream ss;
		for (int i = 0; i < layer; ++i)
		{
			ss << "|";
		}

		ss << "+";
		ss << n->mName.C_Str();
		ss << "[" << n->mNumMeshes << "]";
		ss << std::endl;
		
		OutputDebugStringA(ss.str().c_str());
	}

	for (unsigned int i = 0; i < n->mNumChildren; ++i)
	{
		layer++;
		PlotASNode(n->mChildren[i]);
		layer--;
	}
}
