//-
// ==========================================================================
// Copyright 1995,2006,2008 Autodesk, Inc. All rights reserved.
//
// Use of this software is subject to the terms of the Autodesk
// license agreement provided at the time of installation or download,
// or which otherwise accompanies this software in either electronic
// or hard copy form.
// ==========================================================================
//+

//
// dagMessageCmd.cpp
//
// Description:
//     Sample plug-in that demonstrates how to register/de-register
//     a callback with the MDagMessage class.
//
//     This plug-in will register a new command in maya called
//     "dagMessage" which adds a callback for the all nodes on
//     the active selection list. A message is printed to stdout 
//     whenever a connection is made or broken for those nodes. If
//		nothing is selected, the callback will be for all nodes.
//
//	   dagMessage -help will list the options.
//
#include <maya/MIOStream.h>
#include <maya/MPxCommand.h>
#include <maya/MFnPlugin.h>
#include <maya/MArgList.h>
#include <maya/MArgDatabase.h>
#include <maya/MIntArray.h>
#include <maya/MSelectionList.h>
#include <maya/MGlobal.h>
#include <maya/MPlug.h>
#include <maya/MSyntax.h>
#include <maya/MDagMessage.h>
#include <maya/MDGMessage.h>
#include <maya/MModelMessage.h>
#include <maya/MDagPath.h>
#include <maya/MFnDagNode.h>
#include <maya/MCallbackIdArray.h>
#include <maya/MObjectArray.h>

#define kCmdName					"dagMessage"

#define kAllDagFlag 				"-ad"
#define kAllDagFlagLong				"-allDag"
#define kParentAddedFlag			"-pa"
#define kParentAddedFlagLong		"-parentAdded"
#define kParentRemovedFlag			"-pr"
#define kParentRemovedFlagLong		"-parentRemoved"
#define kChildAddedFlag				"-ca"
#define kChildAddedFlagLong			"-childAdded"
#define kChildRemovedFlag			"-cr"
#define kChildRemovedFlagLong		"-childRemoved"
#define kChildReorderedFlag			"-cro"
#define kChildReorderedFlagLong		"-childReordered"

#define kHelpFlag					"-h"
#define kHelpFlagLong				"-help"

#define CheckErrorContinue(stat, msg)	\
	if (MS::kSuccess != stat) {			\
		displayError(msg);				\
		continue;						\
	}							

#define CheckErrorReturn(stat, msg)		\
	if (MS::kSuccess != stat) {			\
		displayError(msg);				\
		return;							\
	}

// This table will keep track of the registered callbacks
// so they can be removed when the plug-ins is unloaded.
//
MCallbackIdArray callbackIds;

// Node added to model callback.
static void userNodeRemovedCB(MObject& node,void *clientData)
{
	if (! node.isNull()) {
		bool doDisplay = true;

		MStatus status;
		MFnDagNode dagNode(node,&status);
		if ( status.error() ) {
			doDisplay = false;
			MGlobal::displayInfo("Error: failed to get dag node.");
		}

		if ( doDisplay ) {
			MString s = dagNode.name();
			MString info("DAG Model -  Node removed: ");
			info+= s;

			MGlobal::displayInfo(info);	
		}
	}

	// remove the callback
	MCallbackId id = MMessage::currentCallbackId();
	MMessage::removeCallback(id);
}

// Node added to model callback.
static void userNodeAddedCB(MObject& node,void *clientData)
{
	MStatus status;

	if (! node.isNull()) {
		bool doDisplay = true;

		MDagPath path;
		status = MDagPath::getAPathTo(node,path);
		if ( status.error() ) {
			doDisplay = false;
			MGlobal::displayInfo("Error: failed to get dag path to node.");
		}

		if ( doDisplay ) {
			MString s = path.fullPathName();
			MString info("DAG Model -  Node added: ");
			info+= s;

			path.transform(&status);
			if (MS::kInvalidParameter == status) {
				info += "(WORLD)";
			}

			MGlobal::displayInfo(info);	
		}
	}

	// remove the callback
	MCallbackId id = MMessage::currentCallbackId();
	MMessage::removeCallback(id);

	// listen for removal message
	/* MCallbackId id = */ MModelMessage::addNodeRemovedFromModelCallback( node, userNodeRemovedCB, 0, &status );
	if ( status.error() ) {
		MGlobal::displayError("Failed to install node removed from model callback.\n");
		return;
	}
}

// Install a node added callback for the node specified
// by dagPath.
static void installNodeAddedCallback( MDagPath& dagPath )
{
	MStatus status;

	MObject dagNode = dagPath.node();
	if ( dagNode.isNull() )
		return;

	/* MCallbackId id = */ MModelMessage::addNodeAddedToModelCallback( dagNode, userNodeAddedCB, 0, &status );
	if ( status.error() ) {
		MGlobal::displayError("Failed to install node added to model callback.\n");
		return;
	}
}

// Decide if the dag is in the model. Dag paths and names
// may not be setup if the dag has not been added to
// the model.
static bool dagNotInModel( MDagPath& dagPath )
{
	MStatus status;
	MFnDagNode dagFn( dagPath, &status );
	if ( status.error() )
		return false;
	bool inModel = dagFn.inModel( &status );
	if ( status.error() )
		return false;
	return ( inModel == false );
}

void userDAGGenericCB(MDagMessage::DagMessage msg, MDagPath &child,
					  MDagPath &parent, void *)
{	
	MString dagStr("DAG Changed - ");
	switch (msg) {
		case MDagMessage::kParentAdded:
			dagStr += "Parent Added: ";
			break;
		case MDagMessage::kParentRemoved:
			dagStr += "Parent Removed: ";
			break;
		case MDagMessage::kChildAdded:
			dagStr += "Child Added: ";
			break;
		case MDagMessage::kChildRemoved:
			dagStr += "Child Removed: ";
			break;
		case MDagMessage::kChildReordered:
			dagStr += "Child Reordered: ";
			break;
		default:
			dagStr += "Unknown Type: ";
			break;
	}

	dagStr += "child = ";
	dagStr += child.fullPathName();
	dagStr += ", parent = ";
	dagStr += parent.fullPathName();

	// Check to see if the parent is the world object.
	//
	MStatus pStat;
	parent.transform(&pStat);
	if (MS::kInvalidParameter == pStat) {
		dagStr += "(WORLD)";
	}

	// Install callbacks if node is not in the model.
	// Callback is for node added to model.
	bool incomplete = false;
	if ( dagNotInModel( child ) ) {
		installNodeAddedCallback( child );
		incomplete = true;
	}
	if ( dagNotInModel( parent ) ) {
		installNodeAddedCallback( parent);
		incomplete = true;
	}

	// Warn user that dag path info may be
	// incomplete
	if (incomplete)
		dagStr += "\t// May be incomplete!";	

	MGlobal::displayInfo(dagStr);
}


//////////////////////////////////////////////////////////////////////////
//
// Command class declaration
//
//////////////////////////////////////////////////////////////////////////

class dagMessageCmd : public MPxCommand
{
public:
					dagMessageCmd() {};
	virtual			~dagMessageCmd(); 
	MStatus			doIt( const MArgList& args );

	static MSyntax	newSyntax();
	static void*	creator();

private:

	MStatus			addGenericCallback(MDagPath *dagPath, 
									   MDagMessage::DagMessage msg, 
									   MString cbName);	
};

//////////////////////////////////////////////////////////////////////////
//
// Command class implementation
//
//////////////////////////////////////////////////////////////////////////

dagMessageCmd::~dagMessageCmd() {}

void* dagMessageCmd::creator()
{
	return new dagMessageCmd();
}

MSyntax dagMessageCmd::newSyntax()
{
	MSyntax syntax;

	syntax.useSelectionAsDefault(true);
	syntax.setObjectType(MSyntax::kSelectionList);
	syntax.setMinObjects(0);

	syntax.addFlag(kAllDagFlag, kAllDagFlagLong);
	syntax.addFlag(kParentAddedFlag, kParentAddedFlagLong);
	syntax.addFlag(kParentRemovedFlag, kParentRemovedFlagLong);
	syntax.addFlag(kChildAddedFlag, kChildAddedFlagLong);
	syntax.addFlag(kChildRemovedFlag, kChildRemovedFlagLong);
	syntax.addFlag(kChildReorderedFlag, kChildReorderedFlagLong);
	syntax.addFlag(kHelpFlag, kHelpFlagLong);
	return syntax;
}

MStatus dagMessageCmd::addGenericCallback(MDagPath *dagPath, 
										  MDagMessage::DagMessage msg, 
										  MString cbName)
{
	MStatus status = MS::kFailure;

	if (NULL == dagPath) {
		MCallbackId id = MDagMessage::addDagCallback(	msg,
														userDAGGenericCB, 
														NULL, 
														&status);
		if (MS::kSuccess == status) {
			MString info("Adding a callback for");
			info += cbName;
			info += "on all nodes";
			MGlobal::displayInfo(info);			
			callbackIds.append( id );
		} else {
			MString err("Could not add callback to");
			err += dagPath->fullPathName();
			MGlobal::displayError(err);
		}
	} else {
		MCallbackId id = MDagMessage::addDagCallback(*dagPath,
													msg,
													userDAGGenericCB, 
													NULL, 
													&status);
		if (MS::kSuccess == status) {
			MString info("Adding a callback for");
			info += cbName;
			info += "on ";
			info += dagPath->fullPathName();
			MGlobal::displayInfo(info);			
			callbackIds.append( id );
		} else {
			MString err("Could not add callback to");
			err += dagPath->fullPathName();
			MGlobal::displayError(err);
		}
	}

	return status;
}

MStatus dagMessageCmd::doIt( const MArgList& args)
//
// Takes the  nodes that are on the active selection list and adds an
// attriubte changed callback to each one.
//
{
	MStatus 		status;
	MSelectionList 	list;
    MArgDatabase argData(syntax(), args);

	status = argData.getObjects(list);
	if (MS::kSuccess != status) {
		MGlobal::displayError("Error getting objects");
		return status;
	}

	//	Get the flags
	//
	bool allDagUsed = argData.isFlagSet(kAllDagFlag);
	bool parentAddedUsed = argData.isFlagSet(kParentAddedFlag);
	bool parentRemovedUsed = argData.isFlagSet(kParentRemovedFlag);
	bool childAddedUsed = argData.isFlagSet(kChildAddedFlag);
	bool childRemovedUsed = argData.isFlagSet(kChildRemovedFlag);
	bool childReorderedUsed = argData.isFlagSet(kChildReorderedFlag);
	bool helpUsed = argData.isFlagSet(kHelpFlag);

	bool nothingSet = (	!allDagUsed && !parentAddedUsed && 
						!parentRemovedUsed && !childAddedUsed && 
						!childRemovedUsed && !childReorderedUsed && 
						!helpUsed);

	if (nothingSet) {
		MGlobal::displayError("A flag must be used. dagMessage -help for availible flags.");
		return MS::kFailure;
	}

	if (argData.isFlagSet(kHelpFlag)) {
		MGlobal::displayInfo("dagMessage -help");
		MGlobal::displayInfo("\tdagMessage adds a callback to the selected nodes,");
		MGlobal::displayInfo("\tor if no nodes are selected, to all nodes. The callback");
		MGlobal::displayInfo("\tprints a message when called. When the plug-in is unloaded");
		MGlobal::displayInfo("\tthe callbacks are removed.");
		MGlobal::displayInfo("");
		MGlobal::displayInfo("\t-h -help : This message is printed");
		MGlobal::displayInfo("\t-ad -allDag : parent changes and child reorders");
		MGlobal::displayInfo("\t-pa -parentAdded : A parent is added");
		MGlobal::displayInfo("\t-pr -parentRemoved : A parent is removed");
		MGlobal::displayInfo("\t-ca -childAdded : A child is added (only for individual nodes)");
		MGlobal::displayInfo("\t-cr -childRemoved : A child is removed (only for individual nodes)");
		MGlobal::displayInfo("\t-cro -childReordered : A child is reordered");
		MGlobal::displayInfo("");
	}

	unsigned nObjs = list.length();
	if (nObjs == 0) {
		//	Add the callback for all changes of the specified type.
		//
		if (allDagUsed) {
			MCallbackId id = MDagMessage::addAllDagChangesCallback(userDAGGenericCB, NULL, &status);
			if (status) {
				callbackIds.append( id );
				MGlobal::displayInfo("Added a callback for all Dag changes on all nodes.\n");
			} else {
				MGlobal::displayError("Could not add a -allDag callback");
				return status;
			}
		}

		if (parentAddedUsed) {
			status = addGenericCallback(NULL, 
										MDagMessage::kParentAdded,
										MString(" parent added "));
			if (MS::kSuccess != status) {
				return status;
			}
		}

		if (parentRemovedUsed) {
			status = addGenericCallback(NULL, 
										MDagMessage::kParentRemoved,
										MString(" parent removed "));
			if (MS::kSuccess != status) {
				return status;
			}
		}		

		if (childAddedUsed) {
			MGlobal::displayError("-childAdded can only be used when a node is selected");
			status = MS::kFailure;
			return status;
		}

		if (childRemovedUsed) {
			MGlobal::displayError("-childRemoved can only be used when a node is selected");
			status = MS::kFailure;
			return status;
		}	

		if (childReorderedUsed) {
			status = addGenericCallback(NULL, 
										MDagMessage::kChildReordered,
										MString(" child reordered "));
			if (MS::kSuccess != status) {
				return status;
			}
		}	
	} else {
		for (unsigned int i=0; i< nObjs; i++) {
			MDagPath dagPath;
			list.getDagPath(i, dagPath);

			if (!dagPath.isValid()) {
				continue;
			}

			//	Add the callback for all changes of the specified type.
			//
			if (allDagUsed) {
				MCallbackId id = MDagMessage::addAllDagChangesCallback(dagPath, userDAGGenericCB, NULL, &status);
				if (status) {
					callbackIds.append( id );
					MString infoStr("Added a callback for all Dag changes on ");
					infoStr += dagPath.fullPathName();
					MGlobal::displayInfo(infoStr);
				} else {
					MGlobal::displayError("Could not add a -allDag callback");
					return status;
				}
			}

			if (parentAddedUsed) {
				status = addGenericCallback(&dagPath, 
											MDagMessage::kParentAdded,
											MString(" parent added "));
				if (MS::kSuccess != status) {
					return status;
				}
			}

			if (parentRemovedUsed) {
				status = addGenericCallback(&dagPath, 
											MDagMessage::kParentRemoved,
											MString(" parent removed "));
				if (MS::kSuccess != status) {
					return status;
				}
			}		

			if (childAddedUsed) {
				status = addGenericCallback(&dagPath, 
											MDagMessage::kChildAdded,
											MString(" child added "));
				if (MS::kSuccess != status) {
					return status;
				}
			}

			if (childRemovedUsed) {
				status = addGenericCallback(&dagPath, 
											MDagMessage::kChildRemoved,
											MString(" child removed "));
				if (MS::kSuccess != status) {
					return status;
				}
			}	

			if (childReorderedUsed) {
				status = addGenericCallback(&dagPath, 
											MDagMessage::kChildReordered,
											MString(" child reordered "));
				if (MS::kSuccess != status) {
					return status;
				}
			}	
		}
	}
	
	return status;
}

//////////////////////////////////////////////////////////////////////////
//
// Plugin registration
//
//////////////////////////////////////////////////////////////////////////

MStatus initializePlugin( MObject obj )
{
	MFnPlugin plugin( obj );
	return plugin.registerCommand( kCmdName,
									dagMessageCmd::creator,
									dagMessageCmd::newSyntax);
}

MStatus uninitializePlugin( MObject obj)
{
	// Remove callbacks
	//
	for (unsigned int i=0; i<callbackIds.length(); i++ ) {
		MMessage::removeCallback( callbackIds[i] );
	}

	MFnPlugin plugin( obj );
	return plugin.deregisterCommand( kCmdName );
}