/*
    QGL Loader Assimp
    Ivan Pelipenko peri4ko@yandex.ru

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU Lesser General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Lesser General Public License for more details.

    You should have received a copy of the GNU Lesser General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "loader_assimp.h"

#include "glmaterial.h"
#include "glmesh.h"
#include "globject.h"
#include "glscene.h"

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

QString fromAiString(const aiString & s) {
	return QString::fromLocal8Bit(s.C_Str());
}
QColor fromAiColor(const aiColor4D & c) {
	return QColor::fromRgbF(piClamp(c.r, 0.f, 1.f), piClamp(c.g, 0.f, 1.f), piClamp(c.b, 0.f, 1.f));
}
QVector3D fromAiVector3D(const aiVector3D & v) {
	return QVector3D(v.x, v.y, v.z);
}
Vector3i fromAiFace(const aiFace & v) {
	return Vector3i(v.mIndices[0], v.mIndices[1], v.mIndices[2]);
}
QMatrix4x4 fromAiMatrix4D(const aiMatrix4x4 & v) {
	return QMatrix4x4(v.a1, v.a2, v.a3, v.a4, v.b1, v.b2, v.b3, v.b4, v.c1, v.c2, v.c3, v.c4, v.d1, v.d2, v.d3, v.d4);
}
bool isAiMeshTriangles(const aiMesh * m) {
	return (m->mPrimitiveTypes & aiPrimitiveType_TRIANGLE) == aiPrimitiveType_TRIANGLE;
}


Mesh * assimpMesh(const aiMesh * m) {
	if (!m) return 0;
	if (!isAiMeshTriangles(m)) return 0;
	Mesh * ret = new Mesh();
	int vcnt = m->mNumVertices, tcnt = m->mNumFaces;

	QVector<QVector3D> & v(ret->vertices());
	v.resize(vcnt);
	QVector<QVector2D> & t(ret->texcoords());
	t.resize(vcnt);
	QVector<Vector3i> & ind(ret->indicesTriangles());

	for (int i = 0; i < vcnt; ++i)
		v[i] = fromAiVector3D(m->mVertices[i]);

	if (m->HasNormals()) {
		QVector<QVector3D> & n(ret->normals());
		n.resize(vcnt);
		for (int i = 0; i < vcnt; ++i)
			n[i] = fromAiVector3D(m->mNormals[i]);
	}

	if (m->HasTextureCoords(0)) {
		for (int i = 0; i < vcnt; ++i)
			t[i] = fromAiVector3D(m->mTextureCoords[0][i]).toVector2D();
	}

	if (m->HasFaces()) {
		ind.resize(tcnt);
		for (int i = 0; i < tcnt; ++i)
			ind[i] = fromAiFace(m->mFaces[i]);
	} else {
		tcnt = vcnt / 3;
		ind.resize(tcnt);
		int si = 0;
		for (int i = 0; i < tcnt; ++i) {
			si     = i + i + i;
			ind[i] = Vector3i(si, si + 1, si + 2);
		}
	}

	// qDebug() << "add mesh" << v.size() << ret->normals().size() << ret->texcoords().size() << ret->indices().size();

	return ret;
}


QColor aiMatColor(const aiMaterial * m, const char * key, uint s0, uint s1, const QColor & def = Qt::white) {
	aiColor4D col;
	aiReturn r = m->Get(key, s0, s1, col);
	// qDebug() << key << r << col.r << col.g << col.b;
	if (r != aiReturn_SUCCESS) return def;
	return fromAiColor(col);
}
float aiMatFloat(const aiMaterial * m, const char * key, uint s0, uint s1, float def = 0.f) {
	float ret;
	aiReturn r = m->Get(key, s0, s1, ret);
	if (r != aiReturn_SUCCESS) return def;
	return ret;
}
QString aiMatString(const aiMaterial * m, const char * key, uint s0, uint s1) {
	aiString p;
	aiReturn r = const_cast<aiMaterial *>(m)->Get(key, s0, s1, p);
	// qDebug() << "GetTexture" << key << s0 << r << fromAiString(p);
	if (r != aiReturn_SUCCESS) return QString();
	return fromAiString(p);
}
Material * assimpMaterial(const aiMaterial * m) {
	if (!m) return 0;
	Material * ret = new Material();
	/// WARNING: no function GetName() in aiMaterial in stable release
	// ret->name = fromAiString(const_cast<aiMaterial*>(m)->GetName());
	aiString name;
	const_cast<aiMaterial *>(m)->Get(AI_MATKEY_NAME, name);
	ret->name           = fromAiString(name);
	// qDebug() << "mat" << ret->name;
	// for (int i = 0; i < m->mNumProperties; ++i) {
	//	qDebug()<< fromAiString(m->mProperties[i]->mKey);// << "=" << aiMatFloat(m, m->mProperties[i]->mKey.C_Str(), 0, 0);
	// }
	ret->color_diffuse  = aiMatColor(m, AI_MATKEY_COLOR_DIFFUSE);
	ret->color_emission = aiMatColor(m, AI_MATKEY_COLOR_EMISSIVE);
	float shine         = aiMatFloat(m, AI_MATKEY_SHININESS, -1.f);
	if (shine >= 0) {
		ret->map_roughness.color_amount = 0.8f - (shine / 100.f * 0.6f);
		// qDebug() << "shine" << shine;
	}
	ret->map_diffuse.bitmap_path   = aiMatString(m, AI_MATKEY_TEXTURE_DIFFUSE(0));
	ret->map_normal.bitmap_path    = aiMatString(m, AI_MATKEY_TEXTURE_NORMALS(0));
	// ret->map_metalness.bitmap_path = aiMatString(m, AI_MATKEY_TEXTURE_SPECULAR(0));
	ret->map_roughness.bitmap_path = aiMatString(m, AI_MATKEY_TEXTURE_SHININESS(0));
	ret->map_emission.bitmap_path  = aiMatString(m, AI_MATKEY_TEXTURE_EMISSIVE(0));
	ret->transparency              = 1.f - aiMatFloat(m, AI_MATKEY_OPACITY, 1.f);
	ret->detectMaps();
	ret->setTypes();
	return ret;
}


Light * assimpLight(const aiLight * l) {
	if (!l) return 0;
	if (l->mType != aiLightSource_POINT && l->mType != aiLightSource_SPOT) return 0;
	Light * ret = new Light();
	ret->setName(fromAiString(l->mName));
	ret->setPos(fromAiVector3D(l->mPosition));
	ret->setDirection(fromAiVector3D(l->mDirection));
	ret->decay_const     = l->mAttenuationConstant;
	ret->decay_linear    = l->mAttenuationLinear;
	ret->decay_quadratic = l->mAttenuationQuadratic;
	ret->angle_start     = l->mAngleInnerCone * rad2deg;
	ret->angle_end       = l->mAngleOuterCone * rad2deg;
	if (l->mType == aiLightSource_SPOT) ret->light_type = Light::Cone;
	QVector3D col(l->mColorDiffuse.r, l->mColorDiffuse.g, l->mColorDiffuse.b);
	ret->intensity = col.length();
	col /= ret->intensity;
	ret->setColor(QColor::fromRgbF(col[0], col[1], col[2]));
	return ret;
}


ObjectBaseList assimpObject(const aiNode * n,
                            const QVector<Mesh *> & meshes,
                            aiMesh ** ai_meshes,
                            const QVector<Material *> & materials,
                            const QMap<QString, Light *> & light_by_name,
                            QVector<Light *> & out_lights) {
	if (!n) return ObjectBaseList();
	ObjectBaseList ret;
	ObjectBase * obj = 0;
	QString name     = fromAiString(n->mName);
	Light * light    = light_by_name.value(name, 0);
	if (light) {
		obj = light->clone();
		out_lights << (Light *)obj;
	} else
		obj = new ObjectBase();
	obj->setName(name);
	obj->setMatrix(fromAiMatrix4D(n->mTransformation));
	ret << obj;
	// qDebug() << "add object" << ret << ret->name();
	if (!light) {
		// qDebug() << name << "has" << n->mNumMeshes << "meshes";
		for (uint i = 0; i < n->mNumMeshes; ++i) {
			int mi = n->mMeshes[i];
			if (meshes[mi]) {
				if (obj->mesh()) {
					obj = obj->clone(false);
					ret << obj;
				}
				obj->setMesh(meshes[mi]);
				int mati = ai_meshes[mi]->mMaterialIndex;
				if (mati >= 0 || mati < materials.size()) obj->setMaterial(materials[mati]);
				// ret->setMesh(meshes[mi]);
				// qDebug() << "set mesh" << mi << ret->mesh();
				// break;
			}
		}
	}
	for (uint i = 0; i < n->mNumChildren; ++i) {
		ObjectBaseList cl = assimpObject(n->mChildren[i], meshes, ai_meshes, materials, light_by_name, out_lights);
		foreach(ObjectBase * c, cl)
			obj->addChild(c);
	}

	return ret;
}


Scene * loadScene(const QString & filepath) {
	if (filepath.isEmpty()) return 0;
	qDebug() << "[Loader Assimp] Import" << filepath << "...";
	Assimp::Importer importer;
	const aiScene * ais =
		importer.ReadFile(filepath.toUtf8(),
	                      aiProcess_Triangulate | aiProcess_SortByPType | aiProcess_GenUVCoords | aiProcess_TransformUVCoords);
	if (!ais) {
		qDebug() << "[Loader Assimp] Error: \"" + QString(importer.GetErrorString()) + "\"";
		return 0;
	}
	qDebug() << "[Loader Assimp] Imported" << ais->mNumMeshes << "meshes";
	QVector<Mesh *> meshes;
	for (uint i = 0; i < ais->mNumMeshes; ++i)
		meshes << assimpMesh(ais->mMeshes[i]);
	QVector<Material *> materials;
	for (uint i = 0; i < ais->mNumMaterials; ++i)
		materials << assimpMaterial(ais->mMaterials[i]);
	QVector<Light *> lights;
	for (uint i = 0; i < ais->mNumLights; ++i)
		lights << assimpLight(ais->mLights[i]);
	QMap<QString, Light *> light_by_name;
	foreach(Light * l, lights)
		if (l) light_by_name[l->name()] = l;

	QVector<Light *> out_lights;
	ObjectBaseList rootl = assimpObject(ais->mRootNode, meshes, ais->mMeshes, materials, light_by_name, out_lights);
	if (rootl.isEmpty()) return 0;
	ObjectBase * root = rootl[0];
	root->transferTransformToChildren(true);

	ObjectBaseList rcl = root->children(true);
	foreach(ObjectBase * c, rcl) {
		foreach(Light * l, out_lights) {
			if (c->name() == (l->name() + ".Target")) {
				l->setDistance((l->worldPos() - c->worldPos()).length());
				delete c;
				break;
			}
		}
	}
	root->cleanTree();

	Scene * scene = new Scene();
	scene->setName(root->name());
	foreach(ObjectBase * o, root->children())
		scene->addObject(o);
	lights.removeAll(nullptr);
	qDeleteAll(lights);

	return scene;
}


QStringList supportedFormats() {
	Assimp::Importer importer;
	aiString ret;
	importer.GetExtensionList(ret);
	return fromAiString(ret).toLower().split(";");
}
