//-
// ==========================================================================
// Copyright (C) 1995 - 2006 Autodesk, Inc. and/or its licensors.  All 
// rights reserved.
//
// The coded instructions, statements, computer programs, and/or related 
// material (collectively the "Data") in these files contain unpublished 
// information proprietary to Autodesk, Inc. ("Autodesk") and/or its 
// licensors, which is protected by U.S. and Canadian federal copyright 
// law and by international treaties.
//
// The Data is provided for use exclusively by You. You have the right 
// to use, modify, and incorporate this Data into other products for 
// purposes authorized by the Autodesk software license agreement, 
// without fee.
//
// The copyright notices in the Software and this entire statement, 
// including the above license grant, this restriction and the 
// following disclaimer, must be included in all copies of the 
// Software, in whole or in part, and all derivative works of 
// the Software, unless such copies or derivative works are solely 
// in the form of machine-executable object code generated by a 
// source language processor.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND. 
// AUTODESK DOES NOT MAKE AND HEREBY DISCLAIMS ANY EXPRESS OR IMPLIED 
// WARRANTIES INCLUDING, BUT NOT LIMITED TO, THE WARRANTIES OF 
// NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR 
// PURPOSE, OR ARISING FROM A COURSE OF DEALING, USAGE, OR 
// TRADE PRACTICE. IN NO EVENT WILL AUTODESK AND/OR ITS LICENSORS 
// BE LIABLE FOR ANY LOST REVENUES, DATA, OR PROFITS, OR SPECIAL, 
// DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES, EVEN IF AUTODESK 
// AND/OR ITS LICENSORS HAS BEEN ADVISED OF THE POSSIBILITY 
// OR PROBABILITY OF SUCH DAMAGES.
//
// ==========================================================================
//+

//
// 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/MDagPath.h>
#include <maya/MCallbackIdArray.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;

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)";
	}

	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 all callbacks
	//
	for (unsigned int i=0; i<callbackIds.length(); i++ ) {
		MMessage::removeCallback( (MCallbackId)callbackIds[i] );
	}

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