////////////////////////////////////////////////////////////////////////////////
// MW3MechImporter.cpp
// MechWarrior 3 Mech and vehicle import
// Copyright Bjoern Ganster
// Developed since 2007
////////////////////////////////////////////////////////////////////////////////

#include "MechWarrior3MechImporter.h"

#include <iostream>

#define FILE_NUMBER FILE_NO_MECHWARRIOR3_MECH_IMPORTER_CPP

////////////////////////////////////////////////////////////////////////////////
// Import MW3 mechs from ZBD files

MW3MechImporter::MW3MechImporter (bool piratesMoon)
: m_textureLib(nullptr),
  m_tc(nullptr),
  m_mechOutputDir(),
  m_targetPath(),
  mPiratesMoon(piratesMoon)
{
}

// Experiment: obtain texture list from materials file inside mechlib.zbd
void parseMaterials(MechWarrior3ZBD& mechLib, vector<BGCPString>& textureFNs)
{
	SET_SCOPE;
	// Parse materials directly from the so-named file in mechlib.zbd
	size_t fileNum = mechLib.getEntryNumberByName("materials");
	MemoryBlock materialsFile;
	auto fileSize = mechLib.getFileSize(fileNum);
	materialsFile.resize(fileSize);
	if (mechLib.getFileChunk(fileNum, 0, materialsFile.getPtr<unsigned char>(), materialsFile.getSize())) {
		size_t offset = 0x2c;
		while (offset < materialsFile.getSize()) {
			auto nameLen = materialsFile.get<WORD>(offset);
			offset += 4;
			if (nameLen > 0 && nameLen < 30) {
				BGCPString name(materialsFile.getPtr<char>(offset), nameLen, CP_ASCII);
				offset += nameLen;
				textureFNs.emplace_back(name);
			} else {
				textureFNs.emplace_back(BGCPString());
			}

			while (materialsFile.get<BYTE>(offset) != 0x3f && offset < materialsFile.getSize()) {
				++offset;
			}
			offset += 13;
		}

		/*COUT << "Using materials:" << std::endl;
		for (size_t i = 0; i < textureFNs.size(); ++i) {
			COUT << i << ": ";
			textureFNs[i].printLn();
		}*/
	} else {
		COUT << "Odd, unable to obtain materials" << std::endl;
	}
	RETURN;
}

void MW3MechImporter::importSingle(XMLTreeConstIterator& xmlIter, const PathChar* outputDir)
{
	// Import geometry
	SET_SCOPE;
	BGCPString sourcePath;
	xmlIter.getAttributeValueByName ("source", sourcePath);
	xmlIter.getAttributeValueByName ("target", m_targetPath);

	if (m_targetPath.empty()) {
		RETURN;
	}

	COUT << "Importing ";
	m_targetPath.printLn();
	m_mechOutputDir = BGCPString(outputDir, CP_SYSTEM);
	m_mechOutputDir.addPathSep(m_targetPath);
	bgmkdir (m_mechOutputDir.getAsPath());

	RenderableObject meshes;
	if (!sourcePath.empty()) {
		SET_OP;
		size_t fileNum = mMechLib.getEntryNumberByName(sourcePath);
		MemoryBlock flt;
		auto fileSize = mMechLib.getFileSize(fileNum);
		flt.resize(fileSize);
		mMechLib.getFileChunk(fileNum, 0, flt.getPtr<unsigned char>(), flt.getSize());
		exportFLT (&flt, meshes);
	}

	int configsCopied = 0;
	for (size_t i = 0; i < xmlIter.getChildCount(); ++i) {
		// Copy config
		SET_OP;
		xmlIter.goToChild(i);
		if (xmlIter.nodeNameEquals("copyConfig")) {
			// Copy config file (Mech configuration)
			BGCPString filename;
			xmlIter.getAttributeValueByName("filename", filename);
			BGCPString ifn (getDataDir());
			ifn.addPathSep("fixed configs");
			ifn.addPathSep(filename);
			BGCPString ofn (m_mechOutputDir);
			ofn.addPathSep(filename);
			copyFile(ifn.getAsPath(), ofn.getAsPath());
			++configsCopied;
		} else if (xmlIter.nodeNameEquals("addGeometry")) {
			// Copy geometry from external file
			BGCPString filename;
			xmlIter.getAttributeValueByName("filename", filename);
			BGCPString ifn (getDataDir());
			ifn.addPathSep("mech-import");
			ifn.addPathSep(filename);
			meshes.loadFromObj(ifn.getAsPath());
		} else if (xmlIter.nodeNameEquals("importZBD")) {
			// Select gamez.zbd file for import
			BGCPString filename;
			xmlIter.getAttributeValueByName("filename", filename);
			BGCPString fileNoStr = filename.copyPart(4, 10);
			size_t fileNo = (size_t)fileNoStr.toInt();
			fileNo += 3; // account for the fake files ...
			Mesh* newMesh = new Mesh("Mesh", MeshID, &meshes);
			mGameZ.decodeGameZMesh(fileNo, newMesh);
		} else if (xmlIter.nodeNameEquals("shadowFix")) {
			BGCPString method, partName;
			xmlIter.getAttributeValueByName("method", method);
			xmlIter.getAttributeValueByName("part", partName);
			for (size_t i = 0; i < meshes.getObjectCount(); ++i) {
				Mesh* part = dynamic_cast<Mesh*>(meshes.getObject(i));
				if (part != nullptr && part->getName() == partName) {
					if (method == "1") {
						part->findMeshEdges();
						part->fillHoles();
					} else {
						part->findMeshEdges();
						part->fillHoles2();
					}
				}
			}
		} else {
			// Copy description into new mvmk file
			SET_OP;
			BGCPString mvmkFN = m_mechOutputDir;
			mvmkFN.addPathSep(m_targetPath + ".mvmk");

			BGCPString nodeName;
			xmlIter.getNodeName(nodeName);
			XMLTree mvmkTree(nodeName);
			XMLTreeIterator mvmkIter = XMLTreeIterator(&mvmkTree, mvmkTree.getRoot());
			mvmkIter.copyFromChilds(xmlIter);
			mvmkTree.saveToFile(mvmkFN.getAsPath(), CP_ASCII);
		}
		xmlIter.goToParent();
	}

	/*if (configsCopied == 0) {
		// Copy fake mech config if there isn't one
		SET_OP;
		BGCPString input = getDataDir();
		input.addPathSep("fake.mek");
		BGCPString output = m_mechOutputDir;
		output.addPathSep(m_targetPath + "-fake.MEK");
		MemoryBlock mb;
		if (mb.loadFromFile (input.getAsPath()))
			mb.saveToFile(output.getAsPath());
	}*/

	SET_OP;
	BGCPString fullExportPath(m_mechOutputDir);
	fullExportPath.addPathSep(BGCPString("geometry.obj"));

	meshes.saveToObj(fullExportPath.getAsPath(), true);

	RETURN;
}

void MW3MechImporter::execute(const PathChar* mw3Dir, const PathChar* outputDir, const PathChar* importXML)
{
	SET_SCOPE;
	XMLTree importXMLTree;
	if (!importXMLTree.loadFromFile(importXML)) {
		COUT << "Failed to read MW3 mech ist from " << mw3Dir << std::endl;
		RETURN;
	}

	XMLTreeConstIterator xmlIter(importXMLTree.getRoot());

	// Open mech library file mechlib.zbd
	BGCPString mechLibZBDFN(mw3Dir, CP_SYSTEM);
	mechLibZBDFN.addPathSep("zbd");
	mechLibZBDFN.addPathSep("mechlib.zbd");

	// Open mech texture file
	SET_OP;
	BGCPString textureLibFN(mw3Dir, CP_SYSTEM);
	textureLibFN.addPathSep("zbd");
	textureLibFN.addPathSep("rmechtex.zbd");
	if (fileExists(textureLibFN.getAsPath())) {
		m_textureLib = new MechWarrior3ZBD ();
		if (!m_textureLib->open(textureLibFN.getAsPath())) {
			delete m_textureLib;
			m_textureLib = NULL;
		}
	} else {
		m_textureLib = NULL;
	}

	SET_OP;
	if (mMechLib.open (mechLibZBDFN.getAsPath())) {

		// Obtain texture list from materials file inside mechlib.zbd
		parseMaterials(mMechLib, textureFNs);

		for (size_t i = 0; i < xmlIter.getChildCount(); i++) {
			SET_OP;
			xmlIter.goToChild(i);

			if (xmlIter.nodeNameEquals("import")) {
				SET_OP;
				importSingle(xmlIter, outputDir);
			} else if (xmlIter.nodeNameEquals("openZBDs")) {
				SET_OP;
				BGCPString path(mw3Dir, CP_SYSTEM);
				BGCPString subPath;
				xmlIter.getAttributeValueByName("path", subPath);
				path.addPathSep("zbd");
				path.addPathSep(subPath);
				BGCPString gameZPath = path;
				gameZPath.addPathSep("gamez.zbd");
				mGameZ.open(gameZPath.getAsPath(), true);
			}

			SET_OP;
			xmlIter.goToParent();
		}
	}

	RETURN;
}

MW3MechImporter::~MW3MechImporter ()
{
	safeDelete(m_textureLib);
	//safeDelete(m_tc); // deleted in non-rec. MW3MechImporter::exportFLT
}

////////////////////////////////////////////////////////////////////////////////
// MW3 mech import

// Get texture if already loaded or load the texture and then return it ...
ManagedTexture* MW3MechImporter::getTexture(size_t textureNum, bool hasUVs)
{
	SET_SCOPE;
	ManagedTexture* manText = NULL;

	// Select texture or load it
	if (m_textureLib != NULL && hasUVs && textureNum < textureFNs.size()) {
		const BGCPString& textureFN = textureFNs[textureNum];
		auto fileNum = m_textureLib->getEntryNumberByName(textureFN);
		manText = m_tc->findTextureByID(textureNum);
		if (manText == NULL 
		&&  fileNum < m_textureLib->getEntryCount()) 
		{
			MemoryBlock mb;
			size_t textureSize = m_textureLib->getFileSize(fileNum);
			mb.resize(textureSize);
			if (m_textureLib->getFileChunk(fileNum, 0, mb.getPtr<unsigned char>(0), mb.getSize())) {
				Texture* newText = new Texture;
				if (Decode_MW3_555 (mb.getPtr<unsigned char>(0), mb.getSize(), *newText, false)) {
					manText = m_tc->add(newText);
					manText->setID (textureNum);
				} else {
					delete newText;
				}
			}
		}
	}

	RETURN manText;
}

bool gInvertPolygons = false;

void setPolyVertex(bool invertPolygons, MeshPolygon* poly, size_t indexInPoly, size_t vertexIndexInMesh)
{
	if (invertPolygons) {
		size_t vertexCount = poly->getVertexCount();
		poly->setVertexIndex(vertexCount-indexInPoly-1, vertexIndexInMesh);
	} else {
		poly->setVertexIndex(indexInPoly, vertexIndexInMesh);
	}
}

// todo: is there a scale factor hidden somewhere here?
struct FltNodeHeader {
	char partName[36];
	uint8_t flags1, flags2, flags3, flags4;
	uint32_t data[6];
	char unknown[28];
	uint8_t childNodeCount;
	char unknown2[11];
	float unknownFloats[32];
	float rx, ry, rz;
	float unknownFloats2[12];
	float tx, ty, tz;
	float unknownFloats3[5];
	char unknown4[28];
}; // 352 bytes

struct FltMeshHeader {
	char unknown[16];
	DWORD polygonCount, vertexCount, pointCount2;
	char unknown2[62];
}; // 92 bytes

// Read texture coordinates
void MW3MechImporter::readTextureCoords(MemoryBlock* flt, Mesh* mesh, MeshPolygon* poly, size_t polyData, bool invertPolygons)
{
	SET_SCOPE;
	auto vertexCount = poly->getVertexCount();
	for (size_t j = 0; j < vertexCount; j++) {
		double u = 0, v = 0;
		if (polyData + 8 < flt->getSize()) {
			u = flt->get<float>(polyData);
			polyData += 4;
			v = flt->get<float>(polyData);
			polyData += 4;
		}

		auto ti = mesh->addTexCoords(u,v);

		SET_OP;
		if (invertPolygons) {
			poly->setTexCoordIndex(vertexCount-j-1, ti);
		} else {
			poly->setTexCoordIndex(j, ti);
		}

		// Check if textures need to be repeated ...
		if (poly->m_manText != nullptr) {
			if (u > poly->m_manText->xreps) {
				poly->m_manText->xreps = std::min<int>((int)(u + 0.99999), 10);
			}
			if (v > poly->m_manText->yreps) {
				poly->m_manText->yreps = std::min<int>((int)(v + 0.99999), 10);
			}
		}
	}
	RETURN;
}

// Extract nodes recursively
void MW3MechImporter::extractNodes(MemoryBlock* flt, size_t& offset, Matrix m, RenderableObject& meshes)
{
	SET_SCOPE;
	int UnknownPolygonDataSize = 12;
	double xfactor = 1.0;

	FltNodeHeader* nodeHeader = flt->getPtr<FltNodeHeader>(offset);
	if (nodeHeader == nullptr) {
		RETURN;
	}
	bool hasPolygons = ((nodeHeader->flags2 & 2) != 0);
	offset += sizeof(FltNodeHeader);

	if (hasPolygons) {
		SET_OP;
		m = m * TranslationMatrix (nodeHeader->tx, nodeHeader->ty, nodeHeader->tz);
		m = m * RotationMatrix (nodeHeader->rx, -1, 0, 0);
		m = m * RotationMatrix (nodeHeader->ry, 0, -1, 0);
		m = m * RotationMatrix (nodeHeader->rz, 0, 0, -1);

		FltMeshHeader* meshHeader = flt->getPtr<FltMeshHeader>(offset);
		if (meshHeader == nullptr) {
			RETURN;
		}

		Mesh* mesh = new Mesh();
		mesh->setName(nodeHeader->partName);
		offset += sizeof(FltMeshHeader);

		// Extract vertices
		size_t i = 0;
		while (offset + 12 < flt->getSize() && i < meshHeader->vertexCount) {
			SET_OP;
			Point3D p;
			p.x = flt->get<float> (offset) * xfactor;
			p.y = flt->get<float> (offset+4);
			p.z = flt->get<float> (offset+8);
			offset += 12;
			p = m.Transform(p);
			mesh->addPoint(p);
			i++;
		}

		// Skip point set 2
		offset += 12*meshHeader->pointCount2;

		// Extract polygons
		size_t polyData = offset + sizeof(LocalPolygonHeader) * meshHeader->polygonCount;
		for (size_t i = 0; i < meshHeader->polygonCount && offset + 36 < flt->getSize(); i++) {
			SET_OP;
			LocalPolygonHeader* polygonHeader = flt->getPtr<LocalPolygonHeader>(offset);
			offset += sizeof(LocalPolygonHeader);
			MeshPolygon* poly = mesh->addPolygon(polygonHeader->vertexCount);
			poly->m_manText = getTexture(polygonHeader->textureID, (polygonHeader->flags > 0));

			for (size_t j = 0; j < polygonHeader->vertexCount && polyData + 4 < flt->getSize(); j++) {
				DWORD index = flt->get<DWORD> (polyData);
				polyData += 4;
				setPolyVertex(gInvertPolygons, poly, j, index);
			}

			// Skip data
			if (polygonHeader->hasPointIndex2)
				polyData += 4*polygonHeader->vertexCount; // PointIndex2

			// Read texture coordinates
			if (polygonHeader->flags > 0) {
				readTextureCoords(flt, mesh, poly, polyData, gInvertPolygons);
				polyData += 8 * polygonHeader->vertexCount;
			}

			// Skip unknown data
			polyData += UnknownPolygonDataSize*polygonHeader->vertexCount;
		}
		offset = polyData;

		// Store results until final UVs are known (no files are written yet)
		if (meshHeader->vertexCount > 0 && meshHeader->polygonCount > 0) {
			mesh->findMeshEdges();
			meshes.addChild(mesh);
		}
	}

	// Extract child nodes - currently ignoes the node hierarchy, simply visiting the nodes in the order they appear in the file
	SET_OP;
	for (size_t i = 0; i < nodeHeader->childNodeCount && offset < flt->getSize(); ++i) {
		extractNodes (flt, offset, m, meshes);
	}

	RETURN;
}

struct FltMeshHeader_PM {
	char unknown[16];
	DWORD polygonCount, vertexCount, pointCount2;
	char unknown2[64];
	DWORD unknownCount;
	char unknown3[2];
}; // 100 bytes

struct FltPolygonHeader_PM {
	uint8_t vertexCount;
	uint8_t polygonMode; // Or just flags?
	uint8_t data1[12];
	uint8_t hasPointIndex2;
	uint8_t data2[25];
}; // 39 bytes?

// Extract nodes recursively
void MW3MechImporter::extractNodes_PM (MemoryBlock* flt, size_t& offset, Matrix m, RenderableObject& meshes)
{
	SET_SCOPE;
	int AdditionalPolygonDataSize = 4;
	int UnknownPolygonDataSize = 20;
	double xfactor = -1.0;

	// todo: track matrix transforms outside of meshes? -> needs node hierarchy ...

	FltNodeHeader* nodeHeader = flt->getPtr<FltNodeHeader>(offset);
	if (nodeHeader == nullptr) {
		RETURN;
	}
	bool hasPolygons = ((nodeHeader->flags2 & 2) > 0);
	offset += sizeof(FltNodeHeader);

	if (hasPolygons) {
		SET_OP;
		FltMeshHeader_PM* meshHeader = flt->getPtr<FltMeshHeader_PM>(offset);

		if (meshHeader == nullptr) {
			RETURN;
		}

		offset += sizeof(FltMeshHeader_PM);
		m = m * TranslationMatrix (nodeHeader->tx, nodeHeader->ty, nodeHeader->tz);
		m = m * RotationMatrix (nodeHeader->rx, -1, 0, 0);
		m = m * RotationMatrix (nodeHeader->ry, 0, -1, 0);
		m = m * RotationMatrix (nodeHeader->rz, 0, 0, -1);

		Mesh* mesh = new Mesh();
		mesh->setName(nodeHeader->partName);

		// Extract vertices
		size_t i = 0;
		while (offset + 12 < flt->getSize() && i < meshHeader->vertexCount) {
			SET_OP;
			Point3D p;
			p.x = flt->get<float> (offset) * xfactor;
			p.y = flt->get<float> (offset+4);
			p.z = flt->get<float> (offset+8);
			offset += 12;
			p = m.Transform(p);
			mesh->addPoint(p);
			i++;
		}

		// Skip point set 2 - normals?
		offset += 12 * meshHeader->pointCount2;

		// Extract polygons
		size_t polyData = offset + sizeof(FltPolygonHeader_PM) * meshHeader->polygonCount;
		for (size_t i = 0; i < meshHeader->polygonCount && offset + 36 < flt->getSize(); i++) {
			SET_OP;
			FltPolygonHeader_PM* polygonHeader = flt->getPtr<FltPolygonHeader_PM>(offset);
			bool hasUVs = true; // (flt->get<uint8_t>(offset+18) > 0); // 19?
			offset += sizeof(FltPolygonHeader_PM);
			size_t textureNumOffset = polyData + 4*polygonHeader->vertexCount;

			if (polygonHeader->hasPointIndex2) {
				textureNumOffset += 4*polygonHeader->vertexCount;
			}

			auto textureNum = flt->get<uint16_t>(textureNumOffset); // most likely location?
			ManagedTexture* manText = getTexture(textureNum, hasUVs);
			size_t texCoordsOffset = textureNumOffset + AdditionalPolygonDataSize;

			if (polygonHeader->polygonMode != 48) {
				// Single polygon - geometry extracted correctly, only some texture/texture coordinate data seems to come out correctly here
				MeshPolygon* poly = mesh->addPolygon(polygonHeader->vertexCount);
				poly->m_manText = manText;
				for (size_t j = 0; j < polygonHeader->vertexCount && polyData + 4 < flt->getSize(); j++) {
					DWORD index = flt->get<DWORD> (polyData + 4 * j);
					setPolyVertex(!gInvertPolygons, poly, j, index);
				}
				readTextureCoords(flt, mesh, poly, texCoordsOffset, gInvertPolygons);
			} else {
				// Triangle strip
				for (size_t j = 0; j + 2 < polygonHeader->vertexCount && polyData + 4 < flt->getSize(); j++) {
					MeshPolygon* poly = mesh->addPolygon(3);
					poly->m_manText = manText;
					for (size_t k = 0; k < 3; ++k) {
						DWORD index = flt->get<DWORD> (polyData+4*(j+k));
						if (j % 2 == 0) {
							setPolyVertex(gInvertPolygons, poly, 2-k, index);
						} else {
							setPolyVertex(gInvertPolygons, poly, k, index);
						}
					}
					if (j % 2 == 0) {
						readTextureCoords(flt, mesh, poly, texCoordsOffset, true);
					} else {
						readTextureCoords(flt, mesh, poly, texCoordsOffset, false);
					}
					texCoordsOffset += 8;
				}
			}

			polyData = textureNumOffset + AdditionalPolygonDataSize + UnknownPolygonDataSize*polygonHeader->vertexCount;
		}
		offset = polyData + 12 * meshHeader->unknownCount; // ??? not used ???

		// Store results until final UVs are known (no files are written yet)
		if (meshHeader->vertexCount > 0 && meshHeader->polygonCount > 0) {
			meshes.addChild(mesh);
		}
	}

	// Extract child nodes
	SET_OP;
	for (size_t i = 0; i < nodeHeader->childNodeCount && offset < flt->getSize(); ++i) {
		extractNodes_PM (flt, offset, m, meshes);
	}

	RETURN;
}

void MW3MechImporter::exportFLT (MemoryBlock* flt, RenderableObject& meshes) 
{
	SET_SCOPE;
	if (m_textureLib != NULL) {
		m_tc = new TextureCompiler;
	}

	size_t offset = 0;
	Matrix m = RotationMatrix(M_PI, 0, 1, 0);
	if (mPiratesMoon) {
		SET_OP;
		extractNodes_PM (flt, offset, m, meshes);
	} else {
		SET_OP;
		extractNodes (flt, offset, m, meshes);
	}

	// Prepare mech texture
	if (m_tc != NULL) {
		COUT << "Used textures: ";
		for (size_t i = 0; i < textureFNs.size(); ++i) {
			ManagedTexture* manText = m_tc->getTextureByID(i);
			if (manText != nullptr) {
				textureFNs[i].print();
				COUT << ", ";
			}
		}
		COUT << std::endl;

		SET_OP;
		m_tc->reassignPatches();
		Texture* finalMechTexture = m_tc->buildAtlas();
		TextureDim w = finalMechTexture->getWidth();
		TextureDim h = finalMechTexture->getHeight();
		BGCPString textFN = m_mechOutputDir;
		textFN.addPathSep(m_targetPath + ".BMP");
		if (finalMechTexture->getWidth() > 0 
		&&  finalMechTexture->getHeight() > 0)
		{
			finalMechTexture->saveToFile(textFN.getAsPath(), 0, 0, w, h);
		}

		SET_OP;
		//safeDelete(finalMechTexture); //TextureCompiler handles cleanup of the texture
		safeDelete(m_tc);

		// Adjust UVs
		for (size_t i = 0; i < meshes.getObjectCount(); i++) {
			Mesh* mesh = dynamic_cast<Mesh*>(meshes.getObject(i));
			for (size_t j = 0; j < mesh->getPolygonCount(); j++) {
				MeshPolygon* poly = mesh->getPolygon(j);
				ManagedTexture* manText = poly->m_manText;
				if (manText != NULL) {
					for (size_t k = 0; k < poly->getVertexCount(); k++) {
						SET_OP;
						auto texCoordIndex = poly->getTexCoordIndex(k);
						double u, v;
						mesh->getTexCoords(texCoordIndex, u, v);
						u = manText->transformU(u);
						v = manText->transformV(v);
						mesh->setTexCoords(texCoordIndex, u, v);
					}
				}
			}
		}
	}

	// Save .OBJs
	/*SET_OP;
	BGCPString textFN = m_mechOutputDir;
	textFN.addPathSep("geometry.obj");
	meshes.saveToObj(textFN.getAsPath(), true);*/

	RETURN;
}
