#include "externals.hh"
#include "control.hh"
#include <ebcl/Algorithms.hh>

#define INVASIVE_TRACES

using namespace ebcl;
using namespace ops;
using namespace opast;

namespace {

struct T_CompilerImpl_
{
	P_OpProgram compile(
			T_ParserOutput const& input ) noexcept;

    private:
	T_Visitor< A_Node > astVisitor{ ASTVisitorBrowser };
	T_Set< uint32_t > constants{ UseTag< IndexBacked< > >( ) };
	T_KeyValueTable< T_String , uint32_t > locations;

	T_ParserOutput* input;
	P_OpProgram output;

	uint32_t fiVariables , fiFramebuffers ,
		 fiPipelines , fiPrograms ,
		 fiSamplers , fiTextures;

	uint32_t sdMain , sdFPU;
	T_AutoArray< uint32_t , 32 > condJumps;

	void gatherConstants( ) noexcept;
	void countAssets( ) noexcept;

	bool compileNode( uint32_t funcIndex ,
			A_Node& node ,
			bool exit ) noexcept;
	void processFunction(
			bool exit ,
			uint32_t args ,
			uint32_t lvars ,
			T_SRDLocation const& location ) noexcept;
	bool processIdentifier(
			uint32_t funcIndex ,
			T_IdentifierExprNode const& node ) noexcept;

	void addInstruction(
			E_OpType op ,
			T_SRDLocation const& location ) noexcept;
	void addInstruction(
			E_OpType op ,
			uint32_t arg0 ,
			T_SRDLocation const& location ) noexcept;
	void addInstruction(
			E_OpType op ,
			uint32_t arg0 , uint32_t arg1 ,
			T_SRDLocation const& location ) noexcept;
};


P_OpProgram T_CompilerImpl_::compile(
		T_ParserOutput const& in ) noexcept
{
	input = const_cast< T_ParserOutput* >( &in );
	output = NewOwned< T_OpProgram >( );

	// Gather all constants used in expressions, count resources
	gatherConstants( );
	countAssets( );

	// Get function indices
	// FIXME ideally we should remap functions so that init is always 0
	// and frame is always 1
	output->init = input->root.functionIndex( "*init*" );
	output->frame = input->root.functionIndex( "*frame*" );
#ifdef INVASIVE_TRACES
	printf( "function indices\n\t%d\tinit\n\t%d\tframe\n" ,
			output->init , output->frame );
#endif

	// Compile each function
#ifdef INVASIVE_TRACES
	uint32_t nInstr = 0;
#endif
	uint32_t cfi;
	for ( cfi = 0u ; cfi < input->root.nFunctions( ) ; cfi ++ ) {
		output->ops.next( );
		auto& func( input->root.function( cfi ) );
#ifdef INVASIVE_TRACES
		printf( "compiling function %s\n" ,
				func.name( ).toOSString( ).data( ) );
#endif
		sdMain = sdFPU = 0;
		astVisitor.visit( func ,
			[=]( A_Node& node , const bool exit ) -> bool {
				return compileNode( cfi , node , exit );
			} );
#ifdef INVASIVE_TRACES
		T_StringBuilder dump , temp;
		for ( auto i = 0u ; i < output->ops.sizeOf( cfi ) ; i ++ ) {
			temp << "(" << output->ops.get( cfi , i ) << ")";
			while ( temp.length( ) < 30 ) {
				temp << ' ';
			}
			dump << "\t\t" << temp << "{ " << output->ops.get( cfi , i ).location << " }\n";
			temp.clear( );
		}
		dump << '\t' << output->ops.sizeOf( cfi ) << " instructions\n"
			<< '\0';
		printf( "%s" , dump.data( ) );
		nInstr += output->ops.sizeOf( cfi );
#endif
	}
#ifdef INVASIVE_TRACES
	printf( "total %d instructions\n" , nInstr );
#endif

	return std::move( output );
}

void T_CompilerImpl_::gatherConstants( ) noexcept
{
	constants.clear( );
	astVisitor.visit( input->root , [&]( A_Node& node , const bool exit ) {
		if ( exit && node.type( ) == A_Node::EXPR_CONST ) {
			T_OpValue value;
			value.f = dynamic_cast< T_ConstantExprNode& >( node ).floatValue( );
			constants.add( value.u );
		}
		return true;
	} );

#ifdef INVASIVE_TRACES
	printf( "%d constants\n" , constants.size( ) );
	for ( auto i = 0u ; i < constants.size( ) ; i ++ ) {
		printf( "    %08x" , constants[ i ] );
		if ( i % 4 == 3 ) {
			printf( "\n" );
		}
	}
	if ( constants.size( ) % 4 ) {
		printf( "\n" );
	}
#endif
}

void T_CompilerImpl_::countAssets( ) noexcept
{
	locations.clear( );
	locations.add( T_String::Pooled( "time" ) , 0u );
	locations.add( T_String::Pooled( "width" ) , 1u );
	locations.add( T_String::Pooled( "height" ) , 2u );

	auto const nt{ input->types.size( ) };
	for ( auto i = 0u ; i < nt ; i ++ ) {
		const auto t{ input->types.values( )[ i ] };
		auto const& n{ input->types.keys( )[ i ] };
		switch ( t ) {
		    case E_DataType::FRAMEBUFFER:
			locations.add( n , output->nFramebuffers );
			output->nFramebuffers ++;
			break;
		    case E_DataType::PIPELINE:
			locations.add( n , output->nPipelines );
			output->nPipelines ++;
			break;
		    case E_DataType::PROGRAM:
			locations.add( n , output->nPrograms );
			output->nPrograms ++;
			break;
		    case E_DataType::SAMPLER:
			locations.add( n , output->nSamplers );
			output->nSamplers ++;
			break;
		    case E_DataType::TEXTURE:
			locations.add( n , output->nTextures );
			output->nTextures ++;
			break;
		    case E_DataType::VARIABLE:
			locations.add( n , output->nVariables + 3 );
			output->nVariables ++;
			break;

		    case E_DataType::INPUT:
			assert( !output->inputs.contains(
						input->types.keys( )[ i ] ) );
			output->inputs.add( input->types.keys( )[ i ] );
			break;

		    case E_DataType::BUILTIN:
		    case E_DataType::UNKNOWN:
			break;
		}
	}

	fiVariables = 3 + constants.size( );
	fiFramebuffers = fiVariables + output->nVariables;
	fiPipelines = fiFramebuffers + output->nFramebuffers;
	fiPrograms = fiPipelines + output->nPipelines;
	fiSamplers = fiPrograms + output->nPrograms;
	fiTextures = fiSamplers + output->nSamplers;

	for ( auto i = 0u ; i < nt ; i ++ ) {
		const auto li{ locations.indexOf( input->types.keys( )[ i ] ) };
		if ( li == T_HashIndex::INVALID_INDEX ) {
			continue;
		}

		const auto t{ input->types.values( )[ i ] };
		auto& pos{ locations[ li ] };
		switch ( t ) {
		    case E_DataType::FRAMEBUFFER: pos += fiFramebuffers; break;
		    case E_DataType::PIPELINE: pos += fiPipelines; break;
		    case E_DataType::PROGRAM: pos += fiPrograms; break;
		    case E_DataType::SAMPLER: pos += fiSamplers; break;
		    case E_DataType::TEXTURE: pos += fiTextures; break;

		    case E_DataType::VARIABLE:
		    case E_DataType::INPUT:
		    case E_DataType::BUILTIN:
		    case E_DataType::UNKNOWN:
			break;
		}
	}

#ifdef INVASIVE_TRACES
	printf( "assets\n\t%d framebuffers\n\t%d pipelines\n"
			"\t%d programs\n\t%d samplers\n\t%d textures\n"
			"\t%d variables\n\t%d inputs\n" ,
		output->nFramebuffers , output->nPipelines ,
		output->nPrograms , output->nSamplers ,
		output->nTextures , output->nVariables ,
		output->inputs.size( ) );
	printf( "table ranges\n\t0\t2\tBuilt-ins\n"
			"\t3\t%d\tConstants\n"
			"\t%d\t%d\tVariables\n"
			"\t%d\t%d\tFramebuffers\n"
			"\t%d\t%d\tPipelines\n"
			"\t%d\t%d\tPrograms\n"
			"\t%d\t%d\tSamplers\n"
			"\t%d\t%d\tTextures\n" ,
		fiVariables - 1 , fiVariables , fiFramebuffers - 1 , fiFramebuffers ,
		fiPipelines - 1 , fiPipelines , fiPrograms - 1 , fiPrograms ,
		fiSamplers - 1 , fiSamplers , fiTextures - 1 , fiTextures ,
		fiTextures + output->nTextures - 1 );

	T_Array< uint32_t > indices( locations.size( ) );
	indices.ensureCapacity( locations.size( ) );
	for ( auto i = 0u ; i < locations.size( ) ; i ++ ) {
		indices.add( i );
	}
	indices.sort( [this]( uint32_t a , uint32_t b ) {
		return T_Comparator< uint32_t >::compare(
				locations.values( )[ a ] ,
				locations.values( )[ b ] );
	} );
	T_StringBuilder lmap;
	lmap << "location map (constants not included)\n";
	for ( auto idx : indices ) {
		lmap << '\t' << locations.values( )[ idx ]
			<< '\t' << locations.keys( )[ idx ]
			<< '\n';
	}
	lmap << '\0';
	printf( "%s" , lmap.data( ) );
#endif
}

bool T_CompilerImpl_::compileNode(
		const uint32_t funcIndex ,
		A_Node& node ,
		const bool exit ) noexcept
{
	switch ( node.type( ) ) {

	    case A_Node::ROOT:
		fprintf( stderr , "Internal error: root node found during compilation\n" );
		std::abort( );
		break;

	    case A_Node::DECL_FN:
	    {
		T_FuncNode& fn( (T_FuncNode&) node );
		const auto args( fn.arguments( ) );
		processFunction( exit , args , fn.locals( ) - args , node.location( ) );
		break;
	    }

	    case A_Node::DECL_INIT:
	    case A_Node::DECL_FRAME:
	    {
		A_FuncNode& fn( (A_FuncNode&) node );
		processFunction( exit , 0 , fn.locals( ) , node.location( ) );
		break;
	    }

	    case A_Node::OP_CALL:
		if ( exit ) {
			auto& call( (T_CallInstrNode&) node );
			const auto fi( input->root.functionIndex( call.id( ) ) );
			assert( fi >= 0 );
			auto& callee( input->root.function( fi ) );
			assert( callee.type( ) == A_Node::DECL_FN );
			auto& fcallee( (T_FuncNode&) callee );
			const auto args( fcallee.arguments( ) );
			assert( sdMain > args );
			addInstruction( OP_CALL , fi , node.location( ) );
			sdMain -= args;
		}
		break;
	    case A_Node::TN_ARG:
	    {
		auto& n( (T_CallInstrNode::T_Argument&)node );
		if ( n.isIdentifier( ) && !exit ) {
			const bool main{ processIdentifier( funcIndex ,
					(T_IdentifierExprNode&)  n.expression( ) ) };
			if ( !main ) {
				addInstruction( OP_PUSH , node.location( ) );
				addInstruction( OP_FP_SSTORE , 0 , node.location( ) );
				sdMain ++;
				sdFPU --;
			}
			return false;
		} else if ( exit && !n.isIdentifier( ) ) {
			addInstruction( OP_PUSH , node.location( ) );
			addInstruction( OP_FP_SSTORE , 0 , node.location( ) );
			sdMain ++;
			sdFPU --;
		}
		break;
	    }

	    case A_Node::OP_SET:
		if ( exit ) {
			assert( sdFPU > 0 );
			auto& id( ((T_SetInstrNode&)node).id( ) );
			auto& func{ input->root.function( funcIndex ) };
			if ( func.hasLocal( id ) ) {
				const auto pos( func.getLocalIndex( id ) + 1 );
				addInstruction( OP_FP_SSTORE , sdMain - pos - 1 ,
						node.location( ) );
			} else {
				addInstruction( OP_FP_STORE ,
						*locations.get( ((T_SetInstrNode&)node).id( ) ) ,
						node.location( ) );
			}
			sdFPU --;
		}
		break;

	    case A_Node::OP_COND:
		if ( exit ) {
			assert( sdMain > 0 );
			addInstruction( OP_POP , 0 , node.location( ) );
			sdMain --;
		}
		break;

	    case A_Node::TN_CONDITION:
		if ( exit ) {
			addInstruction( OP_PUSH , node.location( ) );
			addInstruction( OP_FP_SSTORE_INT , 0 , node.location( ) );
			sdFPU --;
			sdMain ++;
		}
		break;

	    case A_Node::TN_CASE:
		if ( exit ) {
			const auto cpos( output->ops.sizeOf( funcIndex ) );
			const auto diff( cpos - condJumps.last( ) );
			output->ops.get( funcIndex , condJumps.last( ) ).arg1 = diff - 1;
			condJumps.removeLast( );
#ifdef INVASIVE_TRACES
			printf( "\tCOND JUMP UPDATED: %d\n" , diff - 1 );
#endif
		} else {
			auto& c( (T_CondInstrNode::T_ValuedCase&) node );
			condJumps.add( output->ops.sizeOf( funcIndex ) );
			addInstruction( OP_COND_JUMP , 0 , c.value( ) ,
					node.location( ) );
		}
		break;

	    case A_Node::EXPR_CMP_EQ: case A_Node::EXPR_CMP_NE:
	    case A_Node::EXPR_CMP_GT: case A_Node::EXPR_CMP_GE:
	    case A_Node::EXPR_CMP_LT: case A_Node::EXPR_CMP_LE:
		if ( exit ) {
			const uint32_t op( dynamic_cast< T_BinaryOperatorNode& >( node ).op( )
					- T_BinaryOperatorNode::CMP_EQ );
			addInstruction( OP_FP_CMP , op , node.location( ) );
		}
		break;

	    case A_Node::EXPR_ADD:
		if ( exit ) {
			assert( sdFPU > 2 );
			addInstruction( OP_FP_ADD , node.location( ) );
			sdFPU --;
		}
		break;
	    case A_Node::EXPR_SUB:
		if ( exit ) {
			assert( sdFPU > 2 );
			addInstruction( OP_FP_SUB , node.location( ) );
			sdFPU --;
		}
		break;
	    case A_Node::EXPR_MUL:
		if ( exit ) {
			assert( sdFPU > 2 );
			addInstruction( OP_FP_MUL , node.location( ) );
			sdFPU --;
		}
		break;
	    case A_Node::EXPR_DIV:
		if ( exit ) {
			assert( sdFPU > 2 );
			addInstruction( OP_FP_DIV , node.location( ) );
			sdFPU --;
		}
		break;

	    case A_Node::EXPR_CONST:
		if ( !exit ) {
			T_OpValue value;
			value.f = dynamic_cast< T_ConstantExprNode& >( node ).floatValue( );
			addInstruction( OP_FP_LOAD ,
					constants.indexOf( value.u ) + 3 ,
					node.location( ) );
			sdFPU ++;
		}
		break;

	    case A_Node::EXPR_ID:
		if ( !exit ) {
			processIdentifier( funcIndex ,
					dynamic_cast< T_IdentifierExprNode& >( node ) );
		}
		break;
	}

	return true;
}

void T_CompilerImpl_::processFunction(
		const bool exit ,
		const uint32_t args ,
		const uint32_t lvars ,
		T_SRDLocation const& location ) noexcept
{
	if ( exit ) {
		assert( sdMain == args + lvars + 1 );
		if ( lvars ) {
			addInstruction( OP_POP , lvars - 1 , location );
			sdMain -= lvars;
		}
		sdMain -= 1 + args;
		addInstruction( OP_RET , args , location );
		assert( sdMain == 0 );
	} else {
		if ( lvars ) {
			addInstruction( OP_RES_STACK , lvars - 1 , location );
			sdMain += lvars;
		}
		sdMain += 1 + args;
	}
}

// Returns true if the identifier caused a push to the main stack, false
// if it pushed to the FPU stack.
bool T_CompilerImpl_::processIdentifier(
		const uint32_t funcIndex ,
		T_IdentifierExprNode const& node ) noexcept
{
	auto& func{ input->root.function( funcIndex ) };
	if ( func.hasLocal( node.id( ) ) ) {
		const E_DataType dt{ func.getLocalType( node.id( ) ) };
		assert( dt != E_DataType::UNKNOWN );

		uint32_t stackPos;
		if ( func.isArgument( node.id( ) ) ) {
			auto const& fn( (T_FuncNode&) func );
			const auto nArgs( fn.arguments( ) );
			stackPos = nArgs - 1 - func.getLocalIndex( node.id( ) );
		} else {
			stackPos = func.getLocalIndex( node.id( ) ) + 1;
		}
		assert( stackPos < sdMain );

		const auto p( sdMain - ( stackPos + 1 ) );
		if ( dt == E_DataType::VARIABLE ) {
			addInstruction( OP_FP_SLOAD , p , node.location( ) );
			sdFPU ++;
			return false;
		}

		addInstruction( OP_SLOAD , p , node.location( ) );
		addInstruction( OP_PUSH , node.location( ) );
		sdMain ++;
		return true;
	}

	assert( input->types.contains( node.id( ) ) );
	const E_DataType dt{ *( input->types.get( node.id( ) ) ) };
	assert( dt != E_DataType::UNKNOWN );
	assert( locations.contains( node.id( ) ) );
	if ( dt == E_DataType::VARIABLE || dt == E_DataType::BUILTIN ) {
		addInstruction( OP_FP_LOAD , *locations.get( node.id( ) ) ,
				node.location( ) );
		sdFPU ++;
		return false;
	}

	addInstruction( OP_LOAD , *locations.get( node.id( ) ) ,
			node.location( ) );
	addInstruction( OP_PUSH , node.location( ) );
	sdMain ++;
	return true;
}


void T_CompilerImpl_::addInstruction(
		const E_OpType op ,
		T_SRDLocation const& location ) noexcept
{
	assert( ArgumentsFor( op ) == 0 );
	output->ops.addNew( op , location );
}

void T_CompilerImpl_::addInstruction(
		const E_OpType op ,
		const uint32_t arg0 ,
		T_SRDLocation const& location ) noexcept
{
	assert( ArgumentsFor( op ) == 1 );
	output->ops.addNew( op , location , arg0 );
}

void T_CompilerImpl_::addInstruction(
		const E_OpType op ,
		const uint32_t arg0 ,
		const uint32_t arg1 ,
		T_SRDLocation const& location ) noexcept
{
	assert( ArgumentsFor( op ) == 2 );
	output->ops.addNew( op , location , arg0 , arg1 );
}

}


/*= T_Compiler =================================================================*/

T_Compiler::T_Compiler( ) noexcept
	: A_PrivateImplementation( new T_CompilerImpl_( ) )
{ }

P_OpProgram T_Compiler::compile(
		T_ParserOutput const& input ) noexcept
{
	return p< T_CompilerImpl_ >( ).compile( input );
}