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

#ifdef _WIN32
#include <time.h>
#else
#include <sys/times.h>
#endif

#include <maya/MArgDatabase.h>
#include <maya/MArgList.h>
#include <maya/MDagPath.h>
#include <maya/MFnAttribute.h>
#include <maya/MFnNumericAttribute.h>
#include <maya/MFnTypedAttribute.h>
#include <maya/MGlobal.h>
#include <maya/MPlug.h>
#include <maya/MPlugArray.h>
#include <maya/MPxCommand.h>
#include <maya/MSelectionList.h>
#include <maya/MSyntax.h>
#include <maya/MString.h>
#include <maya/M3dView.h>

#include "shaveCallbacks.h"
#include "shaveGlobals.h"
#include "shaveUtil.h"
#include "shaveUtilCmd.h"
#include "shaveNode.h"
#include "shaveHairShape.h"
#include "shaveSDK.h"

static const char*	flAttrType					= "-attrType";
static const char*	fsAttrType					= "-at";
static const char*	flConvertComponentSelections = "-convertComponentSelections";
static const char*	fsConvertComponentSelections = "-ccs";
static const char*	flCopyAttr					= "-copyAttr";
static const char*	fsCopyAttr					= "-ca";
static const char*	flDuplicateInConnections	= "-duplicateInConnections";
static const char*	fsDuplicateInConnections	= "-dic";
static const char*	flForce						= "-force";
static const char*	fsForce						= "-f";
static const char*	flInitPlugin				= "-initPlugin";
static const char*	fsInitPlugin				= "-ip";
static const char*	flMakeTempFileName			= "-makeTempFileName";
static const char*	fsMakeTempFileName			= "-mtf";
static const char*	flMoveOutConnections		= "-moveOutConnections";
static const char*	fsMoveOutConnections		= "-moc";
#ifdef PROFILE
static const char*	flProfile 					= "-profile";
static const char*	fsProfile					= "-p";
#endif
static const char*	flSilent					= "-silent";
static const char*	fsSilent					= "-s";
static const char*	flTest						= "-test";
static const char*	fsTest						= "-t";
static const char*	flTimestamp					= "-timestamp";
static const char*	fsTimestamp					= "-ts";
static const char*  flScaleAll					= "-scaleAll";
static const char*  fsScaleAll					= "-sa";
static const char*  flScaleCurrent				= "-scaleCur";
static const char*  fsScaleCurrent				= "-sc";



const MString shaveUtilCmd::commandName("shaveUtil");

#ifdef PROFILE
bool shaveUtilCmd::gProfiling = false;

extern "C" void pct_on(char** argv);
extern "C" void pct_off();
#endif


shaveUtilCmd::shaveUtilCmd()
: mForce(false)
, mSilent(false)
{}

shaveUtilCmd::~shaveUtilCmd()	{}


void* shaveUtilCmd::createCmd()
{
	return new shaveUtilCmd();
}


MSyntax shaveUtilCmd::createSyntax()
{
	MSyntax	syntax;

	syntax.enableEdit(false);
	syntax.enableQuery(false);

	syntax.addFlag(fsAttrType, flAttrType, MSyntax::kString);
	syntax.addFlag(fsConvertComponentSelections, flConvertComponentSelections);
	syntax.addFlag(fsCopyAttr, flCopyAttr, MSyntax::kString, MSyntax::kString);
	syntax.addFlag(fsDuplicateInConnections, flDuplicateInConnections);
	syntax.addFlag(fsForce, flForce);
	syntax.addFlag(fsInitPlugin, flInitPlugin);

	syntax.addFlag(
		fsMakeTempFileName,
		flMakeTempFileName,
		MSyntax::kString,
		MSyntax::kString,
		MSyntax::kString
	);

	syntax.addFlag(fsMoveOutConnections, flMoveOutConnections);
#ifdef PROFILE
	syntax.addFlag(fsProfile, flProfile, MSyntax::kBoolean);
#endif
	syntax.addFlag(fsSilent, flSilent);
	syntax.addFlag(fsTest, flTest, MSyntax::kString, MSyntax::kString, MSyntax::kString);
	syntax.addFlag(fsTimestamp, flTimestamp, MSyntax::kString);

	syntax.addFlag(fsScaleAll,		flScaleAll,		MSyntax::kDouble);
	syntax.addFlag(fsScaleCurrent,	flScaleCurrent,	MSyntax::kDouble);

	return syntax;
}


static void threadCalc(unsigned threadID, void* data)
{
	int			i;
	unsigned	seed = *(unsigned*)&data;
	double		val = (double)seed;

	for (i = 0; i < 100000000; i++)
		val = 1.0 / val + seed;

	printf("thread %d, seed %d, result %lf\n", threadID, seed, val);
}


MStatus shaveUtilCmd::doIt(const MArgList& argList)
{
	MStatus 		st;
	MArgDatabase	args(syntax(), argList, &st);

	if (!st) return st;

	mForce = args.isFlagSet(fsForce);
	mSilent = args.isFlagSet(fsSilent);

	if (args.isFlagSet(fsAttrType))
	{
		MString	plugName;

		args.getFlagArgument(fsAttrType, 0, plugName);
		setResult(getAttrType(plugName));
	}
	else if (args.isFlagSet(fsConvertComponentSelections))
	{
		shaveUtil::convertComponentSelections();
	}
	else if (args.isFlagSet(fsCopyAttr))
	{
		MString	srcPlugName;
		MString	destPlugName;
		bool	dupInConns = args.isFlagSet(fsDuplicateInConnections);
		bool	moveOutConns = args.isFlagSet(fsMoveOutConnections);

		args.getFlagArgument(fsCopyAttr, 0, srcPlugName);
		args.getFlagArgument(fsCopyAttr, 1, destPlugName);

		st = copyAttr(srcPlugName, destPlugName, dupInConns, moveOutConns);
	}
	else if (args.isFlagSet(fsInitPlugin))
	{
		finishPluginInitialization();
	}
#ifdef PROFILE
	else if (args.isFlagSet(fsProfile))
	{
		bool	turnItOn;

		args.getFlagArgument(fsProfile, 0, turnItOn);

		if (turnItOn && !gProfiling)
		{
			static char* argv[] = { "maya.bin", NULL, NULL };

			pct_on(argv);
		}
		else if (!turnItOn && gProfiling)
		{
			pct_off();
		}

		gProfiling = turnItOn;
	}
#endif
	else if (args.isFlagSet(fsTimestamp))
	{
		MString	msg;
		args.getFlagArgument(fsTimestamp, 0, msg);
		shaveUtil::timestamp(msg);
	}
	else if (args.isFlagSet(fsMakeTempFileName))
	{
		MString	dir;
		MString	prefix;
		MString	suffix;

		args.getFlagArgument(fsMakeTempFileName, 0, dir);
		args.getFlagArgument(fsMakeTempFileName, 1, prefix);
		args.getFlagArgument(fsMakeTempFileName, 2, suffix);

		MString	fileName = shaveUtil::makeUniqueTempFileName(dir, prefix, suffix);
		setResult(fileName);
	}
	else if (args.isFlagSet(fsTest))
	{
	}
	else if (args.isFlagSet(fsScaleAll))
	{
		double s;
		args.getFlagArgument(fsScaleAll, 0, s);
		return scaleAll((float)s);
	}
	else if (args.isFlagSet(fsScaleCurrent))
	{
		double s;
		args.getFlagArgument(fsScaleCurrent, 0, s);
		return scaleCurrent((float)s);
	}

	return st;
}


MStatus shaveUtilCmd::copyAttr(
	MString	origPlugName,
	MString	dupPlugName,
	bool	dupInConns,
	bool	moveOutConns
) const
{
	MStatus			st;

	MPlug			dupPlug;
	MSelectionList	list;
	MString			msg;
	MPlug			origPlug;

	list.add(origPlugName);
	list.getPlug(0, origPlug);

	list.clear();
	list.add(dupPlugName);
	list.getPlug(0, dupPlug);

	if (origPlug.isNull())
	{
		msg = MString("Cannot find source plug '") + origPlugName + "'.";
		st = MS::kInvalidParameter;
	}
	else if (dupPlug.isNull())
	{
		msg = MString("Cannot find destination plug '") + dupPlugName + "'.";
		st = MS::kInvalidParameter;
	}
	else
	{
		st = copyPlug(origPlug, dupPlug, dupInConns, moveOutConns, true);

		if (st) st = copyPlug(origPlug, dupPlug, dupInConns, moveOutConns, false);
	}

	if (!mSilent && (msg.length() > 0)) MGlobal::displayError(msg);

	return st;
}


MStatus shaveUtilCmd::copyInboundConnection(MPlug& origPlug, MPlug& dupPlug)
	const
{
	MStatus	st;

	//
	// Does the original plug have an incoming connection which
	// needs to be duplicated?
	//
	MPlugArray	conns;

	origPlug.connectedTo(conns, true, false);

	if (conns.length() > 0)
	{
		MDGModifier	mod;

		st = mod.connect(conns[0], dupPlug);
		if (st) st = mod.doIt();

		if (!st && !mSilent)
		{
			MGlobal::displayError(
				MString("Cannot connect '") + conns[0].name()
				+ "' to '" + dupPlug.name() + "': "
				+ st.errorString()
			);
		}
	}

	return st;
}


MStatus shaveUtilCmd::copyPlug(
	MPlug	origPlug,
	MPlug	dupPlug,
	bool	dupInConns,
	bool	moveOutConns,
	bool	checkOnly
) const
{
	MStatus	st = MS::kSuccess;
	MString	msg;

	MObject	origAttr = origPlug.attribute();
	MObject	dupAttr = dupPlug.attribute();

	if (origAttr.apiType() != dupAttr.apiType())
	{
		msg = "The source and destination attribute types do not match.";
		st = MS::kInvalidParameter;
	}
	else if (origPlug.isArray() && !dupPlug.isArray())
	{
		msg = "The source attribute is a multi, but the destination is not.";
		st = MS::kInvalidParameter;
	}
	else if (!origPlug.isArray() && dupPlug.isArray())
	{
		msg = "The destination attribute is a multi, but the source is not.";
		st = MS::kInvalidParameter;
	}
	else if (origPlug.isCompound() && !dupPlug.isCompound())
	{
		msg = "The source attribute is compound, but the destination is not.";
		st = MS::kInvalidParameter;
	}
	else if (!origPlug.isCompound() && dupPlug.isCompound())
	{
		msg = "The destination attribute is compound, but the source is not.";
		st = MS::kInvalidParameter;
	}
	else if (origPlug.isArray())
	{
		unsigned	i;
		MIntArray	indices;

		origPlug.getExistingArrayAttributeIndices(indices);

		for (i = 0; i < indices.length(); i++)
		{
			st = copyPlug(
					origPlug.elementByLogicalIndex(indices[i]),
					dupPlug.elementByLogicalIndex(indices[i]),
					dupInConns,
					moveOutConns,
					checkOnly
				);

			if (!st || checkOnly) break;
		}

		//
		// There may be connections to the entire array, so deal with those
		// as well.
		//
		if (moveOutConns) moveOutboundConnections(origPlug, dupPlug);
		if (dupInConns) copyInboundConnection(origPlug, dupPlug);
	}
	else if (origPlug.isCompound())
	{
		unsigned	i;
		unsigned	n = origPlug.numChildren();

		if (dupPlug.numChildren() != n)
		{
			msg = "Source and destination attributes have different children.";
			st = MS::kInvalidParameter;
		}
		else
		{
			for (i = 0; i < n; i++)
			{
				st = copyPlug(
						origPlug.child(i),
						dupPlug.child(i),
						dupInConns,
						moveOutConns,
						checkOnly
					);

				if (!st) break;
			}

			//
			// There may be connections to the entire compound, so deal
			// with those as well.
			//
			if (st && moveOutConns)
				st = moveOutboundConnections(origPlug, dupPlug);

			if (st && dupInConns)
				st = copyInboundConnection(origPlug, dupPlug);
		}
	}
	else if (!checkOnly)
	{
		//
		// Move output connections from the original plug over to the
		// copy.
		//
		if (moveOutConns)
			st = moveOutboundConnections(origPlug, dupPlug);

		if (st)
		{
			//
			// If the copy plug already has an incoming connection then we
			// won't be able to make another incoming connection and won't
			// be able to set a value.
			//
			MPlugArray	conns;
			dupPlug.connectedTo(conns, true, false);

			if (mForce && (conns.length() > 0))
			{
				//	Break the existing incoming connection.
				//
				MDGModifier	mod;
				mod.disconnect(conns[0], dupPlug);
				mod.doIt();

				conns.clear();
			}

			if (conns.length() == 0)
			{
				//
				// Does the original plug have an incoming connection which
				// needs to be duplicated?
				//
				origPlug.connectedTo(conns, true, false);

				if (dupInConns && (conns.length() > 0))
				{
					st = copyInboundConnection(origPlug, dupPlug);
				}
				else
				{
					//
					// Copy the original plug's value to the copy.
					//
					switch (origAttr.apiType())
					{
						case MFn::kDoubleAngleAttribute:
						case MFn::kDoubleLinearAttribute:
						{
							double	val;
							origPlug.getValue(val);
							dupPlug.setValue(val);
						}
						break;

						case MFn::kFloatAngleAttribute:
						case MFn::kFloatLinearAttribute:
						case MFn::kTimeAttribute:
						{
							float	val;
							origPlug.getValue(val);
							dupPlug.setValue(val);
						}
						break;

						case MFn::kEnumAttribute:
						{
							short	val;
							origPlug.getValue(val);
							dupPlug.setValue(val);
						}
						break;

						case MFn::kCompoundAttribute:
						case MFn::kMessageAttribute:
							// %%% give an error?
						break;

						case MFn::kNumericAttribute:
						{
							MFnNumericAttribute	nAttr(origAttr);

							switch (nAttr.unitType())
							{
								case MFnNumericData::kBoolean:
								{
									bool	val;
									origPlug.getValue(val);
									dupPlug.setValue(val);
								}
								break;

								case MFnNumericData::kByte:
								case MFnNumericData::kChar:
								{
									char	val;
									origPlug.getValue(val);
									dupPlug.setValue(val);
								}
								break;

								case MFnNumericData::kShort:
								{
									short	val;
									origPlug.getValue(val);
									dupPlug.setValue(val);
								}
								break;

								case MFnNumericData::kLong:
								{
									int	val;
									origPlug.getValue(val);
									dupPlug.setValue(val);
								}
								break;

								case MFnNumericData::kFloat:
								{
									float	val;
									origPlug.getValue(val);
									dupPlug.setValue(val);
								}
								break;

								case MFnNumericData::kDouble:
								{
									double	val;
									origPlug.getValue(val);
									dupPlug.setValue(val);
								}
								break;

								default:
								{
									MObject	data;
									origPlug.getValue(data);
									dupPlug.setValue(data);
								}
								break;
							}
						}
						break;

						default:
						{
							MObject	data;
							origPlug.getValue(data);
							dupPlug.setValue(data);
						}
						break;
					}
				}
			}
		}
	}

	if (!mSilent && (msg.length() > 0)) MGlobal::displayError(msg);

	return st;
}


//	Now that we know Maya has loaded shaveUI.mel and executed shaveUI(), we
//	can enable all the bits and pieces which rely on procedures in those
//	scripts.
void shaveUtilCmd::finishPluginInitialization()
{
	shaveCallbacks::registerCallbacks();
	shaveCallbacks::setCleanUpMELOnExit(true);
}


MString shaveUtilCmd::getAttrType(MString plugName)
{
	MString			attrType;
	MSelectionList	list;
	MPlug			plug;

	list.add(plugName);
	list.getPlug(0, plug);

	if (!plug.isNull())
	{
		MObject			attr = plug.attribute();

		switch (attr.apiType())
		{
			case MFn::kDoubleAngleAttribute:
				attrType = "doubleAngle";
			break;

			case MFn::kFloatAngleAttribute:
				attrType = "floatAngle";
			break;

			case MFn::kDoubleLinearAttribute:
				attrType = "doubleLinear";
			break;

			case MFn::kFloatLinearAttribute:
				attrType = "floatLinear";
			break;

			case MFn::kTimeAttribute:
				attrType = "time";
			break;

			case MFn::kEnumAttribute:
				attrType = "enum";
			break;

			case MFn::kCompoundAttribute:
				attrType = "compound";
			break;

			case MFn::kGenericAttribute:
				attrType = "generic";
			break;

			case MFn::kLightDataAttribute:
				attrType = "lightData";
			break;

			case MFn::kMatrixAttribute:
				attrType = "matrix";
			break;

			case MFn::kFloatMatrixAttribute:
				attrType = "floatMatrix";
			break;

			case MFn::kMessageAttribute:
				attrType = "message";
			break;

			case MFn::kAttribute2Double:
				attrType = "double2";
			break;

			case MFn::kAttribute2Float:
				attrType = "float2";
			break;

			case MFn::kAttribute2Short:
				attrType = "short2";
			break;

			case MFn::kAttribute2Int:
				attrType = "int2";
			break;

			case MFn::kAttribute3Double:
				attrType = "double3";
			break;

			case MFn::kAttribute3Float:
				attrType = "float3";
			break;

			case MFn::kAttribute3Short:
				attrType = "short3";
			break;

			case MFn::kAttribute3Int:
				attrType = "int3";
			break;

			case MFn::kAttribute4Double:
				attrType = "double4";
			break;

			case MFn::kNumericAttribute:
			{
				MFnNumericAttribute	nAttr(attr);

				switch (nAttr.unitType())
				{
					case MFnNumericData::kBoolean:
						attrType = "bool";
					break;

					case MFnNumericData::kByte:
						attrType = "byte";
					break;

					case MFnNumericData::kChar:
						attrType = "char";
					break;

					case MFnNumericData::kShort:
						attrType = "short";
					break;

					case MFnNumericData::k2Short:
						attrType = "short2";
					break;

					case MFnNumericData::k3Short:
						attrType = "short3";
					break;

					//
					// The 'int' types have the same enum values as the
					// 'long' types, so we only specify one here.  We've
					// chosen to use the name 'long' because that's what
					// the 'setAttr' command accepts.
					//
					case MFnNumericData::kLong:
						attrType = "long";
					break;

					case MFnNumericData::k2Long:
						attrType = "long2";
					break;

					case MFnNumericData::k3Long:
						attrType = "long3";
					break;

					case MFnNumericData::kFloat:
						attrType = "float";
					break;

					case MFnNumericData::k2Float:
						attrType = "float2";
					break;

					case MFnNumericData::k3Float:
						attrType = "float3";
					break;

					case MFnNumericData::kDouble:
						attrType = "double";
					break;

					case MFnNumericData::k2Double:
						attrType = "double2";
					break;

					case MFnNumericData::k3Double:
						attrType = "double3";
					break;

					default:
					break;
				}
			}
			break;

			case MFn::kTypedAttribute:
			{
				MFnTypedAttribute	tAttr(attr);

				switch (tAttr.attrType())
				{
					case MFnData::kNumeric:
						attrType = "numeric";
					break;

					case MFnData::kPlugin:
						//
						// This doesn't work.  For plugin data attrType()
						// returns MFnData::kInvalid.
						//
						attrType = "pluginData";
					break;

					case MFnData::kPluginGeometry:
						attrType = "pluginGeometry";
					break;

					case MFnData::kString:
						attrType = "string";
					break;

					case MFnData::kMatrix:
						attrType = "matrix";
					break;

					case MFnData::kStringArray:
						attrType = "stringArray";
					break;

					case MFnData::kDoubleArray:
						attrType = "doubleArray";
					break;

					case MFnData::kIntArray:
						attrType = "intArray";
					break;

					case MFnData::kPointArray:
						attrType = "pointArray";
					break;

					case MFnData::kVectorArray:
						attrType = "vectorArray";
					break;

					case MFnData::kComponentList:
						attrType = "componentList";
					break;

					case MFnData::kMesh:
						attrType = "mesh";
					break;

					case MFnData::kLattice:
						attrType = "lattice";
					break;

					case MFnData::kNurbsCurve:
						attrType = "nurbsCurve";
					break;

					case MFnData::kNurbsSurface:
						attrType = "nurbsSurface";
					break;

					case MFnData::kDynArrayAttrs:
						attrType = "arrayAttrs";
					break;

					case MFnData::kDynSweptGeometry:
						attrType = "sweptGeometry";
					break;

					case MFnData::kSubdSurface:
						attrType = "subd";
					break;

					default:
					break;
				}
			}
			break;

			default:
			break;
		}
	}

	return attrType;
}


MStatus shaveUtilCmd::moveOutboundConnections(MPlug& origPlug, MPlug& dupPlug)
	const
{
	MStatus	st;

	MPlugArray	conns;
	unsigned	i;
	MDGModifier	mod;

	origPlug.connectedTo(conns, false, true);

	for (i = 0; i < conns.length(); i++)
	{
		st = mod.disconnect(origPlug, conns[i]);
		if (st) st = mod.connect(dupPlug, conns[i]);
		if (!st) break;
	}

	if (st) st = mod.doIt();

	if (!st && !mSilent)
	{
		MGlobal::displayError(
			MString("Cannot move output connections from '")
			+ origPlug.name() + "' to '" + dupPlug.name() + "': "
			+ st.errorString()
		);
	}

	return st;
}


MStatus shaveUtilCmd::scaleAll(float s)
{
	MStatus stat = MS::kSuccess;
	MGlobal::displayInfo(MString("Scale all ") + s);

	MObjectArray nodes;
	shaveUtil::getShaveNodes(nodes);

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

	for (i = 0; i < numShaveNodes; i++)
	{
		MFnDependencyNode dFn(nodes[i]);

		shaveHairShape* shShape = (shaveHairShape*)dFn.userNode();

		SHAVENODE* sh = shShape->getHairNode();
		shShape->doXform();
		SOFTscale_select(s);
		SHAVEfetch_node(sh);
		shShape->applyEdits(true);


		float trigger;
		MPlug triggerPlug = dFn.findPlug("trigger");
		triggerPlug.getValue(trigger);
		triggerPlug.setValue(trigger+1.0f);


		M3dView::active3dView().refresh(true,true);

	}
	return stat;
}
MStatus	shaveUtilCmd::scaleCurrent(float s)
{
	MStatus stat = MS::kSuccess;
	//MGlobal::displayInfo(MString("Scale current ") + s);

	MSelectionList list;
	MDagPath nodePath;
	MObject component;
	MFnDagNode sFn;

	MGlobal::getActiveSelectionList( list );
	for ( MItSelectionList listIter( list ); !listIter.isDone(); listIter.next() )
	{
		

		listIter.getDagPath( nodePath, component );
		nodePath.extendToShape();
		sFn.setObject( nodePath );

		MGlobal::displayInfo(sFn.name());

		MObjectArray nodes;
		shaveUtil::getShaveNodes(nodes);

		unsigned int	numShaveNodes = nodes.length();
		unsigned int	i;
		for (i = 0; i < numShaveNodes; i++)
		{
			MFnDependencyNode dFn(nodes[i]);
			if(dFn.name() == sFn.name())
			{
				MGlobal::displayInfo(MString("Scale ") + dFn.name() + " " + s);
				
				shaveHairShape* shShape = (shaveHairShape*)dFn.userNode();

				SHAVENODE* sh = shShape->getHairNode();
				shShape->doXform();
				SOFTscale_select(s);
				SHAVEfetch_node(sh);
				shShape->applyEdits(true);

				float trigger;
				MPlug triggerPlug = dFn.findPlug("trigger");
				triggerPlug.getValue(trigger);
				triggerPlug.setValue(trigger+1.0f);
			}

		}
	}
	M3dView::active3dView().refresh(true,true);
	return stat;
}