/*
===========================================================================
maya2q3 - export .md3 files from maya
Copyright (C) 2005 HermitWorks Entertainment Corporation

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License
as published by the Free Software Foundation; either version 2
of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
===========================================================================
*/

#include "common.h"

namespace shader
{

class Q3ShaderNode : public MPxHwShaderNode
{
private:
	// Attributes
	static MObject attrDispColor;
	static MObject attrShaderName;
	static MObject attrAnimTime;

	// Internal data
	MPoint m_cameraPosWS;

	MFloatVector m_dispColor;
	bool m_attribsDirty;

	// Callbacks that we monitor so we can release OpenGL-dependant
	// resources before their context gets destroyed.
	MCallbackId m_cbBeforeNew;
	MCallbackId m_cbBeforeOpen;
	MCallbackId m_cbBeforeRemoveRef;
	MCallbackId m_cbMayaExiting;

#if MAYA_API_VERSION >= 700
	static size_t cbScriptCommandRefCount;
	static MCallbackId cbScriptCommand;

	//list of all instances of this type, used to notify when the
	//search path for shader data is changed...should this be
	//scrapped in favor of searching the scene for all instances?
	static std::vector< Q3ShaderNode* > allNodes;
#endif

	Shader *m_shader;
protected:
	static void CDECL releaseCallback( void* clientData )
	{
		static_cast< Q3ShaderNode* >( clientData )->ReleaseAllGL();
	}

#if MAYA_API_VERSION >= 700
	static void CDECL commandCallback( const MString &cmd, void * )
	{
		//check to see if the command is changing the search path for the shaders

		const char *cmdStr = cmd.asChar();

		const char *s0 = strstr( cmdStr, "workspace" );			//this is the only command that can change the path
		const char *s1 = strstr( cmdStr, "Q3GameData" );		//this is the path's key in the paths list

		const char *r = strstr( cmdStr, "-rt" );				//the path is stored as a render type
		if( !r )
			r = strstr( cmdStr, "-renderType" );

		const char *q = strstr( cmdStr, "-q" );					//query mode can't change the value
		if( !q )
			q = strstr( cmdStr, "-query" );

		if( s0 && s1 && r && !q )
		{
			//someone's touching the game data
			//flush the folders and refresh on the next load
			q3::ResourceLoader::ClearDirectories();

			//ToDo: force all instances of this node to flush their data
			for( size_t i = 0; i < allNodes.size(); i++ )
				allNodes[i]->m_attribsDirty = true;
		}
	}
#endif

	void OutputGLError( const char *call )
	{
		GLenum error;

		while( (error = glGetError()) != GL_NO_ERROR )
			std::cerr << call << ":" << error << " is " << (const char *)gluErrorString( error ) << "\n";
	}

	void ReleaseAllGL( void )
	{
		MStatus err;
		M3dView view = M3dView::active3dView( &err );
		  
		if( err == MS::kSuccess )
		{
			view.beginGL();

			if( m_shader )
				m_shader->UnloadGLResources();
		
			view.endGL();
		}		
	}

public:
	static MTypeId id;

	static void* CDECL Creator( void ) { return new Q3ShaderNode; }

	static MStatus CDECL Initialize( void )
	{
		MFnTypedAttribute attr;
		MFnNumericAttribute numAttr; 

		attrShaderName = attr.create( "shaderName", "sn", MFnData::kString );
		attr.setStorable( true );
		attr.setKeyable( false );
		//attr.setDefault( MString( "" ) ); --can't do this! is there a better way? do I even need to?
		attr.setCached( true );
		attr.setInternal( true );

		attrAnimTime = numAttr.create( "animTime", "at", MFnNumericData::kFloat );
		numAttr.setStorable( true );
		numAttr.setKeyable( true );
		numAttr.setCached( true );
		numAttr.setInternal( true );
		numAttr.setMin( 0.0F );
		numAttr.setMax( 60.0F );
		numAttr.setDefault( 0.0F );

		// Create input attributes
		attrDispColor = numAttr.createColor( "dispColor", "dc" );
		numAttr.setStorable( true );
		numAttr.setKeyable( false );
		numAttr.setDefault( 0.5F, 0.5F, 0.5F );
		numAttr.setCached( true );
		numAttr.setInternal( true );

		addAttribute( attrShaderName );
		addAttribute( attrDispColor );
		addAttribute( attrAnimTime );

		attributeAffects( attrDispColor, outColor );
		attributeAffects( attrShaderName, outColor );
		attributeAffects( attrAnimTime, outColor );

		return MS::kSuccess;
	}

	Q3ShaderNode( void ) : m_dispColor( 0.5F, 0.5F, 0.5F ), m_attribsDirty( false ), m_shader( null ),
		m_uvTex( null )
	{
		m_cbBeforeNew = MSceneMessage::addCallback( MSceneMessage::kBeforeNew,
			releaseCallback, this );
		m_cbBeforeOpen = MSceneMessage::addCallback( MSceneMessage::kBeforeOpen,
			releaseCallback, this );
		m_cbBeforeRemoveRef = MSceneMessage::addCallback( MSceneMessage::kBeforeRemoveReference,
			releaseCallback, this);
		m_cbMayaExiting = MSceneMessage::addCallback( MSceneMessage::kMayaExiting,
			releaseCallback, this );

#if MAYA_API_VERSION >= 700
		if( !cbScriptCommandRefCount )
		{
			MStatus status;
			cbScriptCommand = MCommandMessage::addCommandCallback( commandCallback, this, &status );

			if( status == MS::kSuccess )
				cbScriptCommandRefCount++;
		}
		else
			cbScriptCommandRefCount++;

		allNodes.push_back( this );
#endif
	}

	virtual /* override */ ~Q3ShaderNode( void )
	{
		if( m_cbBeforeNew )
			MMessage::removeCallback( m_cbBeforeNew );
		if( m_cbBeforeOpen )
			MMessage::removeCallback( m_cbBeforeOpen );
		if( m_cbBeforeRemoveRef )
			MMessage::removeCallback( m_cbBeforeRemoveRef );
		if( m_cbMayaExiting )
			MMessage::removeCallback( m_cbMayaExiting );

#if MAYA_API_VERSION >= 700
		if( cbScriptCommandRefCount )
			if( --cbScriptCommandRefCount == 0 )
				MMessage::removeCallback( cbScriptCommand );

		for( size_t i = 0; i < allNodes.size(); i++ )
			if( allNodes[i] == this )
			{
				allNodes.erase( allNodes.begin() + i );
				break;
			}
#endif
	}

	virtual void /* override */ postConstructor( void ) { setMPSafe( false ); }

	virtual MStatus /* override */ compute( const MPlug &plug, MDataBlock &block )
	{
		if( plug != outColor && plug.parent() != outColor )
			return MS::kUnknownParameter;

		MFloatVector &color = block.inputValue( attrDispColor ).asFloatVector();

		MDataHandle outColorHandle = block.outputValue( outColor );
		MFloatVector& outColor = outColorHandle.asFloatVector();

		outColor = color;

		outColorHandle.setClean();

		return MS::kSuccess;
	}

	virtual bool /* override */ getInternalValueInContext( const MPlug& plug,
		MDataHandle &handle, MDGContext & )
	{
		bool didGet = false;

		if( plug == attrDispColor )
		{
			didGet = true;
			handle.set( m_dispColor );
		}

		return didGet;
	}

	virtual bool /* override */ setInternalValueInContext( const MPlug& plug,
		const MDataHandle& handle, MDGContext& )
	{
		bool didStore = false;

		if( plug == attrDispColor )
		{
			didStore = true;

			MFloatVector &val = handle.asFloatVector();
			if( m_dispColor != val )
				m_dispColor = val;
		}
		else if( plug == attrShaderName )
		{
			//let Maya store this, we just want to know when it changes
			didStore = false; 
			m_attribsDirty = true;			
		}

		return didStore;
	}

	virtual int /* override */ normalsPerVertex( void ) { return 1; }
	virtual int /* override */ texCoordsPerVertex( void ) { return 1; }
	
	virtual bool /* override */ hasTransparency( void )
	{
		if( m_shader )
			return m_shader->HasTransparency();

		return false;
	}

private:
	gl::Context m_shaderCtx;
public:
	MStatus EnsureState( void )
	{
		//do GL setup
		gl::Context ctx = gl::Context::GetCurrent();

		if( ctx != m_shaderCtx )
		{
			//Maya makes no promises about current context and it may change
			//so keep the extensions bound to the latest one!!!
 			glewInit();

			//m_shaderCtx will be updated towards the end of this function
		}

		if( m_attribsDirty )
		{	
			//kill the uv editor texture
			if( ctx != m_uvCtx )
			{
				if( m_uvCtx.MakeCurrent() )
				{
					glDeleteTextures( 1, &m_uvTex );
					m_uvTex = null;

					ctx.MakeCurrent();
				}
			}
			else
			{
				glDeleteTextures( 1, &m_uvTex );
				m_uvTex = null;
			}

			//kill the shader
			if( m_shader )
			{
				m_shader->UnloadGLResources();
				delete m_shader;
				m_shader = null;
			}

			//see if we should get a shader
			MString shName = "";
			MPlug plug( thisMObject(), attrShaderName );
			plug.getValue( shName );

			if( shName != "" )
			{
				const char *shaderName = shName.asChar();

				//load up new shader!
				const char **shaderFiles = q3::ResourceLoader::FindFiles( "scripts", ".shader" );
				for( const char **cch = shaderFiles; *cch; cch++ )
				{
					void *data;
					size_t size;

					ospath resName = "scripts/";
					resName.Append( *cch );

					if( q3::ResourceLoader::GetData( resName, data, size ) )
					{
						//as a convenience the ResourceLoader has already appended a null
						//to the buffer data (note this is not counted in size).
						const char *shaderPos = Shader::FindShaderInFile( (const char*)data, shaderName );
						if( shaderPos )
						{
							m_shader = new Shader( shaderName );

							try
							{
								m_shader->Parse( shaderPos );
							}
							catch( ... )
							{
								delete m_shader;
								m_shader = null;
							}
						}

						q3::ResourceLoader::FreeData( data );

						if( m_shader )
							break;
					}
				}

				q3::ResourceLoader::FindClose( shaderFiles );

				if( !m_shader )
				{
					m_shader = new Shader( shaderName );
					try
					{
						m_shader->SetToTexture( shaderName );
					}
					catch( ... )
					{
						//should never happen since a missing texture
						//will map to the default "can't find it" texture
						delete m_shader;
						m_shader = null;
					}
				}
			}
		}

		if( m_shader && (m_attribsDirty || ctx != m_shaderCtx ) )
		{
			m_shader->UnloadGLResources();

			try
			{
				m_shader->LoadGLResources();
			} catch( ... )
			{
				//shader can't be displayed, fall back to solid color
				delete m_shader;
				m_shader = null;
			}
		}

		m_shaderCtx = ctx;

		m_attribsDirty = false;

		return MS::kSuccess;
	}

	MStatus DoBind( void )
	{
		MStatus stat;

		stat = EnsureState();
		if( stat != MS::kSuccess )
			return stat;

		return MS::kSuccess;
	}
	virtual MStatus /* override */ glBind( const MDagPath& /* shapePath */ ) { return DoBind(); }

	MStatus DoUnbind( void )
	{	
		return MS::kSuccess;
	}

	virtual MStatus /* override */ glUnbind(const MDagPath& /* shapePath */) { return DoUnbind(); }

	virtual MStatus /* override */ glGeometry( const MDagPath &shapePath,
		int prim, unsigned int writable, int indexCount, const unsigned int * indexArray,
		int vertexCount, const int * vertexIDs, const float * vertexArray,
		int normalCount, const float ** normalArrays, int colorCount,
		const float ** colorArrays, int texCoordCount, const float ** texCoordArrays )
	{
		return Draw( prim, writable, indexCount, indexArray, vertexCount,
			vertexIDs, vertexArray, normalCount, normalArrays, colorCount, 
			colorArrays, texCoordCount, texCoordArrays, &shapePath );
	}

	MStatus Draw( int prim, unsigned int /* writable */, int indexCount, const unsigned int * indexArray,
		int /* vertexCount */, const int * /* vertexIDs */, const float * vertexArray, int /* normalCount */,
		const float ** normalArrays, int /* colorCount */, const float ** /* colorArrays */, int texCoordCount,
		const float ** texCoordArrays, const MDagPath * /* objPath */ )
	{
		glPushAttrib( GL_ENABLE_BIT | GL_LIGHTING_BIT | GL_CURRENT_BIT );

		glDisable( GL_TEXTURE_1D );

		glPushClientAttrib( GL_CLIENT_VERTEX_ARRAY_BIT );
		
		glEnableClientState( GL_VERTEX_ARRAY );
		glEnableClientState( GL_NORMAL_ARRAY );

		glVertexPointer( 3, GL_FLOAT, 0, vertexArray );
		glNormalPointer( GL_FLOAT, 0, normalArrays[0] );

		if( texCoordCount )
		{
			glEnableClientState( GL_TEXTURE_COORD_ARRAY );
			glTexCoordPointer( 2, GL_FLOAT, 0, texCoordArrays[0] );
		}

		bool didDraw = false;
		if( m_shader )
		{
			ShaderContext ctx;

			float timeVal = 0.0F;

			MPlug timePlug( thisMObject(), attrAnimTime );
			if( timePlug.getValue( timeVal ) != MS::kSuccess )
				//not sure if Maya clobbered anything or not...
				timeVal = 0.0F;

			ctx.SetTime( (unsigned int)fabsf(timeVal * 1000.0F ) );

			size_t passes = m_shader->BeginShader( ctx );
			if( passes != (size_t)~0 )
			{
				//the ~0 case is a special warning flag which indicates that Q3 rendering cannot
				//continue for some reason (usually a missing GL extension), fall back to shaded
				for( size_t i = 0; i < passes; i++ )
					if( m_shader->BeginStage( ctx, i ) )
					{
						glDrawElements( prim, indexCount, GL_UNSIGNED_INT, indexArray );

						m_shader->EndStage( ctx );
					}

				m_shader->EndShader( ctx );
				didDraw = true;
			}
		}
		
		if( !didDraw )
		{
			glDisable( GL_TEXTURE_2D );

			float black[4] = { 0.0F, 0.0F, 0.0F, 1.0F };
			float diffuse[4] = { m_dispColor.x, m_dispColor.y, m_dispColor.z, 1.0F };

			if( glIsEnabled( GL_LIGHTING ) )
			{
				glDisable( GL_COLOR_MATERIAL );
				glMaterialfv( GL_FRONT, GL_AMBIENT_AND_DIFFUSE , diffuse );
				glMaterialfv( GL_FRONT, GL_SPECULAR, black );
				glMaterialfv( GL_FRONT, GL_EMISSION, black );
				glMateriali( GL_FRONT, GL_SHININESS, 0 ); 
			}
			else
				glColor4fv( diffuse );

			glDrawElements ( prim, indexCount, GL_UNSIGNED_INT, indexArray );
		}

		glPopClientAttrib();

		glPopAttrib();
		
		return MS::kSuccess;
	}

	virtual MStatus /* override */ getAvailableImages( const MString & /* uvSetName */,
		MStringArray &imageNames )
	{	
		if( m_shader )
			for( size_t is = 0; is < m_shader->GetStageCount(); is++ )
			{
				const Shader::StageInfo si = m_shader->GetStageInfo( is );

				for( size_t it = 0; it < si.numImages; it++ )
				{
					sstr< q3::qname::MaxLength + 10 > texName;

					texName.Append( "s%i t%i %s", is + 1, it + 1, (const char*)si.images[it].name );

					imageNames.append( MString( texName ) );
				}
			}

		return MS::kSuccess;
	}

private:
	q3::qname m_uvTexName;
	GLuint m_uvTex;
	gl::Context m_uvCtx;
public:

	virtual MStatus /* override */ renderImage( const MString &imageName,
		const float region[2][2], int &imageWidth, int &imageHeight )
	{
		if( !m_shader )
			return MS::kFailure;

		unsigned int is, it;
		
		if( imageName != "" )
		{
			if( sscanf( imageName.asChar(), "s%i t%i ", &is, &it ) != 2 )
				return MS::kUnknownParameter;

			is--;
			it--;
		}
		else
		{
			is = 0;
			it = 0;
		}

		if( is >= m_shader->GetStageCount() )
			return MS::kUnknownParameter;

		const Shader::StageInfo si = m_shader->GetStageInfo( is );

		if( it >= si.numImages )
			return MS::kUnknownParameter;

		const Shader::StageInfo::ImageInfo &ii = si.images[it];
		
		gl::Context ctx = gl::Context::GetCurrent();
		if( ctx != m_uvCtx )
		{
			if( m_uvCtx.MakeCurrent() )
			{
				glDeleteTextures( 1, &m_uvTex );
				m_uvTex = null;

				ctx.MakeCurrent();
			}

			m_uvCtx = ctx;
		}

		if( m_uvTex && m_uvTexName != ii.name )
		{
			glDeleteTextures( 1, &m_uvTex );
			m_uvTex = null;
		}

		if( !m_uvTex )
		{
			m_uvTex = si.LoadImage( it );
			m_uvTexName = ii.name;
			m_uvCtx = ctx;
		}

		if( !m_uvTex )
			return MS::kUnknownParameter;

		imageWidth = ii.width;
		imageHeight = ii.height;

		glPushAttrib( GL_ENABLE_BIT | GL_CURRENT_BIT | GL_TEXTURE_BIT );

		glEnable( GL_TEXTURE_2D );
		glBindTexture( GL_TEXTURE_2D, m_uvTex );
		
		glBegin( GL_QUADS );
		{
			//these images are loaded upside-down to the standard Maya way
			//of doing things...flip the UVs to correct for this

			glTexCoord2f( region[0][0], 1.0F - region[0][1] );
			glVertex2f( region[0][0], region[0][1] );

			glTexCoord2f( region[0][0], 1.0F - region[1][1] );
			glVertex2f( region[0][0], region[1][1] );
			
			glTexCoord2f( region[1][0], 1.0F - region[1][1] );
			glVertex2f( region[1][0], region[1][1] );
			
			glTexCoord2f( region[1][0], 1.0F - region[0][1] );
			glVertex2f( region[1][0], region[0][1] );
		}
		glEnd();

		glPopAttrib();

		return MS::kSuccess;
	}

	virtual MStatus /* override */ renderSwatchImage( MImage &image )
	{
		//black out the image since otherwise it's just random bits
		//don't waste time drawing an image into it since that would force
		//us to flip GL context and reload a bunch of data for what is
		//essentially not very much

		unsigned int w, h, c;
		image.getSize( w, h );
		c = image.depth();

		byte *pixBuff = image.pixels();

		while( h-- )
			for( int i = w; i--; )
				for( unsigned int i = 0; i < c; i++ )
					*(pixBuff++) = 0;

		return MS::kSuccess;
	}
};

/*
Note: this is an internal ID, it is (probably) not going to be unique
across maya binary (.mb) files from different workshops. We're releasing
this plugin - should we get Alias to give us our own range of unique ids?
*/
MTypeId Q3ShaderNode::id( 0x0007FF01 );

MObject Q3ShaderNode::attrShaderName;
MObject Q3ShaderNode::attrDispColor;
MObject Q3ShaderNode::attrAnimTime;

#if MAYA_API_VERSION >= 700
size_t Q3ShaderNode::cbScriptCommandRefCount = 0;
MCallbackId Q3ShaderNode::cbScriptCommand;

std::vector< Q3ShaderNode* > Q3ShaderNode::allNodes;
#endif

//
//  exports through plugin.h
//
MTypeId &g_shaderTypeID = Q3ShaderNode::id;
const char *g_shaderName = "Q3Shader";
void* (CDECL *g_shaderCreator)( void ) = &Q3ShaderNode::Creator;
MStatus (CDECL *g_shaderInit)( void ) = &Q3ShaderNode::Initialize;

class Q3ShaderDnDBehavior : public MPxDragAndDropBehavior
{
private:
static MStatus connectAttr( const MPlug &srcPlug, const MPlug &destPlug, bool force )
{
	if( srcPlug.isNull() || destPlug.isNull() )
		return MS::kFailure;

	MString cmd = "connectAttr ";

	if( force )
		cmd += "-f ";
	
	cmd += srcPlug.name() + " " + destPlug.name();
	
	return MGlobal::executeCommand( cmd );
}

public:
	static void* CDECL Creator() { return new Q3ShaderDnDBehavior; }

	Q3ShaderDnDBehavior( void ) { }
	virtual /* override */ ~Q3ShaderDnDBehavior( void ) { }

	virtual bool /* override */ shouldBeUsedFor( MObject &sourceNode,
		MObject & destinationNode, MPlug & /* sourcePlug */, MPlug & /* destinationPlug */ )
	{
		return MFnDependencyNode( sourceNode ).typeName() == g_shaderName
			&& destinationNode.hasFn( MFn::kLambert );
	}

	virtual MStatus /* override */ connectNodeToNode( MObject &sourceNode,
		MObject &destinationNode, bool force )
	{
		MFnDependencyNode src( sourceNode );

		if( src.typeName() != g_shaderName || !destinationNode.hasFn( MFn::kLambert ) )
			return MS::kFailure;
	
		MFnDependencyNode dest( destinationNode );
		return connectAttr( src.findPlug( "outColor", true ), dest.findPlug( "hardwareShader", true ), force );
	}

	virtual MStatus /* override */ connectNodeToAttr( MObject &sourceNode,
		MPlug &destinationPlug, bool force )
	{
		MFnDependencyNode src( sourceNode );

		if( src.typeName() != g_shaderName || !destinationPlug.node().hasFn( MFn::kLambert ) )
			return MS::kFailure;
		
		return connectAttr( src.findPlug( "outColor", true ), destinationPlug, force );
	}

	virtual MStatus /* override */ connectAttrToNode( MPlug &sourcePlug, 
		MObject &destinationNode, bool force )
	{
		MFnDependencyNode src( sourcePlug.node() );

		if( src.typeName() != g_shaderName || destinationNode.hasFn( MFn::kLambert ) )
			return MS::kFailure;

		MFnDependencyNode dest( destinationNode );

		return connectAttr( sourcePlug, dest.findPlug( "hardwareShader", true ), force );
	}

	virtual MStatus /* override */ connectAttrToAttr( MPlug &sourcePlug,
		MPlug &destinationPlug, bool force )
	{
		return connectAttr( sourcePlug, destinationPlug, force );
	}
};

//
//  exports through plugin.h
//
const char *g_behaviorName = "Q3ShaderBehavior";
void* (CDECL *g_behaviorCreator)( void ) = &Q3ShaderDnDBehavior::Creator;

};