// Shave and a Haircut
// (c) 2019 Epic Games
// US Patent 6720962

#include <maya/MDGMessage.h>
#include <maya/MDGModifier.h>
#include <maya/MEventMessage.h>
#include <maya/MFileIO.h>
#include <maya/MFnDependencyNode.h>
#include <maya/MGlobal.h>
#include <maya/MObjectArray.h>
#include <maya/MPlug.h>
#include <maya/MSceneMessage.h>
#include <maya/MSelectionList.h>

#include "shaveCallbacks.h"
#include "shaveDebug.h"
#include "shaveGlobals.h"
#include "shaveHairShape.h"
#include "shaveRender.h"
#include "shaveSDK.h"
#include "shaveUtil.h"


//--------------------------------------------------------------------
//
//	External Interface
//
//--------------------------------------------------------------------

void shaveCallbacks::endLoad(bool isImport, bool isReference)
{
	ENTER();

	shaveUtil::setLoadingFile(false);

	//
	// The output mesh of a shaveHairShape can depend upon other objects in
	// the scene: hair objects, skull objects, shaveGlobals, etc.  So the
	// shaveHairShape cannot properly initialize its output mesh until
	// after the scene is fully loaded.
	//
	// If we're done loading all files, then run through all the
	// shaveHairShapes and initialize any new ones.
	//
	if (!shaveUtil::isLoadingFile()) initializeHairNodes();

	//if default renderer is changed in Maya prefs
	//we need to make sure that correct MEL callbacks are set
	if (!shaveUtil::isLoadingFile())
		MGlobal::executeCommand("shave_selectedRendererChanged");

	LEAVE();
}


void shaveCallbacks::mayaExiting(void* clientData)
{
	if (mCleanUpMELOnExit)
    {
		MGlobal::executeCommand("shaveCleanup(true)");
    }
//fprintf (stdout,"why are we calling SHAVEcleanup()?\n");fflush(stdout);
	SHAVEcleanup();

	mCleanUpMELOnExit = false;

	// You might think that we should call removeCallbacks() here, but it's
	// possible that this method is executing within one of those callbacks
	// and Maya does not behave well when a callback function removes
	// itself.  So we leave the callback removal to others.
	//
}


void shaveCallbacks::prepareForNewScene(bool firstTime)
{
	ENTER();

	SHAVEclear_stack();

	LEAVE();
}


void shaveCallbacks::registerCallbacks()
{
	//
	// Now that we know the plugin has loaded successfully, register
	// the callbacks.
	//
	beforeExportID =	MSceneMessage::addCallback(
									MSceneMessage::kBeforeExport, beforeExport
								);

	afterExportID =		MSceneMessage::addCallback(
									MSceneMessage::kAfterExport, afterExport
								);

	beforeImportID =	MSceneMessage::addCallback(
									MSceneMessage::kBeforeImport, beforeImport
								);

	afterImportID =		MSceneMessage::addCallback(
									MSceneMessage::kAfterImport, afterImport
								);

	beforeImportRefID =	MSceneMessage::addCallback(
									MSceneMessage::kBeforeImportReference,
									beforeImportReference
								);

	afterImportRefID =	MSceneMessage::addCallback(
									MSceneMessage::kAfterImportReference,
									afterImportReference
								);

	beforeOpenID =		MSceneMessage::addCallback(
									MSceneMessage::kBeforeOpen, beforeOpen
								);

	afterOpenID =		MSceneMessage::addCallback(
									MSceneMessage::kAfterOpen, afterOpen
								);

	beforeReferenceID =	MSceneMessage::addCallback(
									MSceneMessage::kBeforeReference,
									beforeReference
								);

	afterReferenceID =	MSceneMessage::addCallback(
									MSceneMessage::kAfterReference,
									afterReference
								);

	beforeSaveID =		MSceneMessage::addCallback(
									MSceneMessage::kBeforeSave, beforeSave
								);

	afterSaveID =		MSceneMessage::addCallback(
									MSceneMessage::kAfterSave, afterSave
								);

	afterNewID =		MSceneMessage::addCallback(
									MSceneMessage::kAfterNew, afterNew
								);

	mayaExitingID =		MSceneMessage::addCallback(
									MSceneMessage::kMayaExiting,
									mayaExiting
								);

	nodeCreatedID =		MDGMessage::addNodeAddedCallback(nodeCreated);

	shaveNodeDeletedID=	MDGMessage::addNodeRemovedCallback(
									shaveNodeDeleted,
									shaveHairShape::nodeTypeName
								);

	selectionChangedID=	MEventMessage::addEventCallback(
									"SelectionChanged", selectionChanged
								);

	//	The node added/removed callbacks above keep a count of the number
	//	of hair nodes in the scene. There may already be some so let's get
	//	the counter properly initialized.
	if (shaveHairShape::initNumShaveNodes() > 0)
	{
		//	There is additional initialization required when a scene
		//	contains shaveNodes.
		shaveCallbacks::firstHairNode();
	}

	//	Several of the callbacks above keep track of whether we're in the
	//	middle of loading a file or not. A file load may already be in
	//	progress, so let's get the state of that set correctly.
	//
	// (It would be nice to just use MFileIO::isReadingFile() everywhere,
	// instead of having to track it, but Maya sets isReadingFile to false
	// before the newly-loaded nodes have had a chance to run their first
	// compute cycle, so that's no good.)
	//
	if (MFileIO::isReadingFile())
		shaveCallbacks::startLoad();
	else
	{
		//	The callbacks above are responsible for initializing any
		//	existing hair nodes once all file loading has completed. File
		//	loading may have completed before the callbacks were set up,
		//	though, so we must initialize any hair nodes which are already
		//	in the scene.
		initializeHairNodes();
	}
}


void shaveCallbacks::removeCallbacks()
{
	removeCallback(beforeExportID);
	removeCallback(afterExportID);
	removeCallback(beforeImportID);
	removeCallback(afterImportID);
	removeCallback(beforeImportRefID);
	removeCallback(afterImportRefID);
	removeCallback(beforeOpenID);
	removeCallback(afterOpenID);
	removeCallback(beforeReferenceID);
	removeCallback(afterReferenceID);
	removeCallback(beforeSaveID);
	removeCallback(afterSaveID);

	removeCallback(afterNewID);
	removeCallback(mayaExitingID);
	removeCallback(nodeCreatedID);
	removeCallback(shaveNodeDeletedID);
	removeCallback(selectionChangedID);
}


void shaveCallbacks::startLoad()
{
	ENTER();
	shaveUtil::setLoadingFile(true);
	LEAVE();
}



//--------------------------------------------------------------------
//
//	Internal Interface
//
//--------------------------------------------------------------------

bool	shaveCallbacks::mCleanUpMELOnExit = true;

//
// Callback IDs
//
MCallbackId	shaveCallbacks::beforeExportID = 0;
MCallbackId	shaveCallbacks::afterExportID = 0;
MCallbackId	shaveCallbacks::beforeImportID = 0;
MCallbackId	shaveCallbacks::afterImportID = 0;
MCallbackId	shaveCallbacks::beforeImportRefID = 0;
MCallbackId	shaveCallbacks::afterImportRefID = 0;
MCallbackId	shaveCallbacks::beforeOpenID = 0;
MCallbackId	shaveCallbacks::afterOpenID = 0;
MCallbackId	shaveCallbacks::beforeReferenceID = 0;
MCallbackId	shaveCallbacks::afterReferenceID = 0;
MCallbackId	shaveCallbacks::beforeSaveID = 0;
MCallbackId	shaveCallbacks::afterSaveID = 0;

MCallbackId	shaveCallbacks::afterNewID = 0;
MCallbackId	shaveCallbacks::mayaExitingID = 0;
MCallbackId	shaveCallbacks::nodeCreatedID = 0;
MCallbackId	shaveCallbacks::shaveNodeDeletedID = 0;
MCallbackId	shaveCallbacks::selectionChangedID = 0;


void shaveCallbacks::beforeExport(void* clientData)
{
	ENTER();

	updateNodeVersions();

	//
	// The renderer gets called during some types of export, so make sure
	// that it knows about it.
	//
	shaveRender::exportStart();

    //
    // There's a bug in Maya's ASCII format when saving locked, shared
    // nodes.  shaveGlobals is a shared node so if this is an ASCII save,
    // be sure to unlock it first.
    //
    // We can't use MFileIO::fileType() because that reflects the type of
    // the main scene, not the export file.  Instead, we'll look at the
    // filename and if ends in anything other than '.mb' we'll unlock
    // the shaveGlobals node.
    //
    MString filename = MFileIO::beforeExportFilename().toLowerCase();
    int     len = filename.length();

    if ((len < 3) || (filename.substring(len-2, len) != ".mb"))
        unlockShaveGlobals();

	LEAVE();
}


void shaveCallbacks::afterExport(void* clientData)
{
	ENTER();

	//
	// The renderer gets called during some types of export, so make sure
	// that it knows about it.
	//
	shaveRender::exportEnd();

	LEAVE();
}


void shaveCallbacks::beforeImport(void* clientData)
{
	ENTER();
	startLoad();
	LEAVE();
}


void shaveCallbacks::afterImport(void* clientData)
{
	ENTER();
	endLoad(true, false);
	LEAVE();
}


void shaveCallbacks::beforeImportReference(void* clientData)
{
	ENTER();
	startLoad();
	LEAVE();
}


void shaveCallbacks::afterImportReference(void* clientData)
{
	ENTER();
	endLoad(true, false);
	LEAVE();
}


void shaveCallbacks::afterNew(void* clientData)
{
	ENTER();

	prepareForNewScene(false);

	MGlobal::executeCommand("shaveSceneCleared");
	MGlobal::executeCommand("shave_resetSelectedRenderer");

	LEAVE();
}


void shaveCallbacks::beforeOpen(void* clientData)
{
	ENTER();

	prepareForNewScene(false);
	startLoad();

	LEAVE();
}


void shaveCallbacks::afterOpen(void* clientData)
{
	ENTER();


	endLoad(false, false);

	//MObjectArray	shaveHairShapes;
	//shaveUtil::getShaveNodes(shaveHairShapes);
	//for(unsigned int i = 0; i < shaveHairShapes.length(); i++)
	//{
	//	MFnDependencyNode dFn(shaveHairShapes[i]);
	//	MPlug plug = dFn.findPlug("clumps");
	//	int c; plug.getValue(c); plug.setValue(c);

	//	////does not make any effect
	//	//shaveHairShape* sh  = (shaveHairShape*)dFn.userNode();
	//	//sh->mDisplayHairCacheDirty = true;
	//}

	LEAVE();
}


void shaveCallbacks::beforeReference(void* clientData)
{
	ENTER();
	startLoad();
	LEAVE();
}


void shaveCallbacks::afterReference(void* clientData)
{
	ENTER();
	endLoad(false, true);
	if(shaveRender::rendererIsVray())
	{
		MGlobal::executeCommand("shave_removeShaveRenderCallbacks;");
		MGlobal::executeCommand("shaveVraySetRenderCallbacks;");
	}
	LEAVE();
}


void shaveCallbacks::beforeSave(void* clientData)
{
	ENTER();

	//
	// If there aren't any shaveHairShapes then get rid of the shaveGlobals
	// node as well.
	//
	bool			stillHaveNodes = false;
	MObjectArray	nodes;

	shaveUtil::getShaveNodes(nodes);

	if (nodes.length() == 0)
	{
		unsigned int	i;

		shaveUtil::getShaveGlobalsNodes(nodes);

		if (nodes.length() > 0)
		{
			MDGModifier		dgMod;

			for (i = 0; i < nodes.length(); i++)
			{
				MFnDependencyNode	nodeFn(nodes[i]);
				shaveGlobals*		nodePtr = (shaveGlobals*)nodeFn.userNode();

				if (!nodePtr->deleteMe(dgMod)) stillHaveNodes = true;
			}

			nodes.clear();

			if (!dgMod.doIt()) stillHaveNodes = true;
		}
	}
	else
		stillHaveNodes = true;

	if (stillHaveNodes)
	{
		updateNodeVersions();

    	//
    	// There's a bug in Maya's ASCII format when saving locked, shared
    	// nodes.  shaveGlobals is a shared node so if this is an ASCII save,
    	// be sure to unlock it first.
    	//
    	if (MFileIO::fileType() == "mayaAscii") unlockShaveGlobals();
	}

	LEAVE();
}


void shaveCallbacks::afterSave(void* clientData)
{
	ENTER();
	MGlobal::executeCommand("shaveSceneWritten");
	LEAVE();
}


void shaveCallbacks::nodeCreated(MObject& node, void* clientData)
{
	//
	// This is a bit too expensive to log ENTER/LEAVE on every node
	// created, even in debug mode, so we'll leave the macros until we know
	// if it's a node we care about.
	//
	MFnDependencyNode	nodeFn(node);

	if (nodeFn.typeId() == shaveHairShape::id)
	{
		ENTER();

		//
		// If this is the first shaveHairShape to be added to the scene,
		// enable various bits of functionality.
		//
		if (shaveHairShape::getNumShaveNodes() == 0)
			firstHairNode();

		shaveHairShape::nodeAdded();

		//if default renderer is changed in Maya prefs
		//we need to make sure that correct MEL callbacks are set
		if (!shaveUtil::isLoadingFile())
			MGlobal::executeCommand("shave_selectedRendererChanged");

		LEAVE();
	}
	else if (node.hasFn(MFn::kField))
	{
		//	If a new field node has been added, mark the existing list as
		//	dirty.
		shaveUtil::setFieldsDirty();
	}
}


void shaveCallbacks::shaveNodeDeleted(MObject& node, void* clientData)
{
	ENTER();

	shaveHairShape::nodeRemoved();

	//
	// If this is the last shaveHairShape, remove the rendering callbacks
	// so that their overhead is not incurred.
	//
	if (shaveHairShape::getNumShaveNodes() == 0)
	{
		shaveRender::cleanupCallbacks();

		//
		// Some menu items should be disabled if there are no shave nodes.
		//
		// The evalDeferred is required because the node is not fully
		// removed yet and we don't want the command to execute until it
		// is.
		//
		MGlobal::executeCommand("evalDeferred shave_enableMenus");
	}
	else if (shaveHairShape::getNumShaveNodes() < 0)
	{
		MGlobal::displayWarning(
			"Internal Shave problem: negative shaveHairShape count."
		);
	}

	LEAVE();
}


void shaveCallbacks::selectionChanged(void* clientData)
{
	//
	// If we're supposed to ignore this selection change, then ignore it
	// and reset the flag.
	//
	if (shaveUtil::isIgnoringNextSelectionChange())
	{
		shaveUtil::setIgnoringNextSelectionChange(false);
		return;
	}

	//
	// Apply the selections to all the hair shapes in the scene.
	//
	MObjectArray	shaveNodes;
	shaveUtil::getShaveNodes(shaveNodes);

	if (shaveNodes.length() > 0)
	{
		MDagPath			currentHair = shaveUtil::getCurrentHairShape();
		unsigned			i;
		MFnDependencyNode	nodeFn;
		shaveHairShape*		nodePtr;
		MSelectionList		selections;

		//
		// Have the non-current hair nodes clear their selection caches.
		//
		for (i = 0; i < shaveNodes.length(); i++)
		{
			if (!currentHair.isValid() || (shaveNodes[i] != currentHair.node()))
			{
				nodeFn.setObject(shaveNodes[i]);
				nodePtr = (shaveHairShape*)nodeFn.userNode();
				nodePtr->clearSelections();
			}
		}

		//
		// Have the current hair node update its selection cache to match
		// the selection list.
		//
		if (currentHair.isValid())
		{
			MGlobal::getActiveSelectionList(selections);

			nodeFn.setObject(currentHair.node());
			nodePtr = (shaveHairShape*)nodeFn.userNode();
			nodePtr->updateSelections(selections);
		}
	}

	//
	// Let the MEL scripts do their stuff.
	//
	MGlobal::executeCommand("shave_selectionChanged");
}


void shaveCallbacks::firstHairNode()
{
	shaveRender::setupCallbacks();

	// Some menu items should be enabled if there are now shave nodes.
	//
	// The evalDeferred is required because the node is not fully
	// created yet and we don't want the command to execute until it
	// is.
	MGlobal::executeCommand("evalDeferred shave_enableMenus");
}


void shaveCallbacks::initializeHairNodes()
{
	bool	warningRequired = false;

	//	Get the plugin version.
	const char*	pluginVersionStr = SHAVEquery_version();

	// If we have any shaveHairShapes in the scene which are from a
	// newer version of Shave, delete them and give a warning.
	MObjectArray	shaveNodes;
	shaveUtil::getShaveNodes(shaveNodes);

	MString			badSceneVersionStr;
	MDGModifier		dgMod;
	unsigned int	i;
	unsigned int	numShaveNodes = shaveNodes.length();

	for (i = numShaveNodes; i > 0; --i)
	{
		MPlug	sceneVersionPlug(
					shaveNodes[i-1], shaveHairShape::aShaveVersion
				);
		MString	sceneVersionStr;

		sceneVersionPlug.getValue(sceneVersionStr);

		if (shaveUtil::compareShaveVersions(sceneVersionStr, pluginVersionStr)
				> 0)
		{
			MFnDependencyNode	nodeFn(shaveNodes[i-1]);
			shaveHairShape*		sn = (shaveHairShape*)nodeFn.userNode();

			sn->deleteMe(dgMod);
			shaveNodes.remove(i-1);
			warningRequired = true;
			badSceneVersionStr = sceneVersionStr;
			--numShaveNodes;
		}
	}

	dgMod.doIt();

	if (warningRequired)
	{
		MGlobal::displayError(
			MString("Shave: plugin version is ") + pluginVersionStr
			+ " but scene contains incompatible elements from a"
			+ " later version of Shave (" + badSceneVersionStr
			+ "). The incompatible elements have been deleted so"
			+ " DO NOT SAVE OUT THIS SCENE or you may lose data."
		);
	}

	//
	//	Parameters override
	//
	//int	pluginVerParts[3];
	//splitShaveVersion(pluginVersionStr, pluginVerParts);
	for (i = 0; i < numShaveNodes; i++)
	{
		MString	sceneVersionStr;
		int sceneVerParts[3];
		MPlug	sceneVersionPlug(shaveNodes[i],  shaveHairShape::aShaveVersion);
		sceneVersionPlug.getValue(sceneVersionStr);
		shaveUtil::splitShaveVersion(sceneVersionStr, sceneVerParts);
		if(sceneVerParts[0] < 6)
		{
			float val;
			MGlobal::displayInfo("Shave: 'Value Variation' was overriden");
			MPlug	valVarPlug(shaveNodes[i],  shaveHairShape::shaveParamValueVariation);
			valVarPlug.getValue(val);
			valVarPlug.setValue(val/3.0f);
		}
	}

	//
	// Let MEL do its stuff.
	//
	MGlobal::executeCommand("shaveSceneLoaded");

	//
	// If we have any shaveHairShapes in the scene, then make sure that
	// we have a shaveGlobals node as well.
	//
	if (numShaveNodes > 0) MGlobal::executeCommand("shaveGlobals()");

	//
	// Tell the shaveHairShapes that they can calculate their output
	// meshes now.
	//
	for (i = 0; i < numShaveNodes; i++)
	{
		MFnDependencyNode	nodeFn(shaveNodes[i]);
		shaveHairShape*		nodePtr = (shaveHairShape*)nodeFn.userNode();

		if (!nodePtr->nodeIsInitialized())
			nodePtr->initializeNode();
	}
}


void shaveCallbacks::removeCallback(MCallbackId& id)
{
	if (id != 0)
	{
		MMessage::removeCallback(id);
		id = 0;
	}
}


void shaveCallbacks::unlockShaveGlobals()
{
    MObject node = shaveGlobals::getDefaultNode();

    if (!node.isNull())
    {
        MFnDependencyNode   nodeFn(node);

        if (nodeFn.isLocked()) nodeFn.setLocked(false);
    }
}


void shaveCallbacks::updateNodeVersions()
{
	ENTER();

	//
	// Make sure that the version numbers on all of the shaveGlobals nodes
	// are up to date.
	//
	MObjectArray	nodes;
	int				nodeVersion;

	shaveUtil::getShaveGlobalsNodes(nodes);

	unsigned int	numNodes = nodes.length();
	unsigned int	i;

	for (i = 0; i < numNodes; i++)
	{
		MPlug	plug(nodes[i], shaveGlobals::aNodeVersion);

		plug.getValue(nodeVersion);

		if (nodeVersion != shaveGlobals::kNodeVersion)
			plug.setValue(shaveGlobals::kNodeVersion);
	}

	//
	// Ditto the shaveHairShapes.
	//
	MString	pluginVersion = SHAVEquery_version();
	MString	sceneVersion;

	nodes.clear();
	shaveUtil::getShaveNodes(nodes);
	numNodes = nodes.length();

	for (i = 0; i < numNodes; i++)
	{
		MPlug	plug(nodes[i], shaveHairShape::aNodeVersion);

		plug.getValue(nodeVersion);

		if (nodeVersion != shaveHairShape::kNodeVersion)
			plug.setValue(shaveHairShape::kNodeVersion);

		plug.setAttribute(shaveHairShape::aShaveVersion);
		plug.getValue(sceneVersion);

		if (sceneVersion != pluginVersion)
			plug.setValue(pluginVersion);
	}

	nodes.clear();

	LEAVE();
}