#include "externals.hh"
#include "opast.hh"
#include <ebcl/Files.hh>
#include <ebcl/SRDText.hh>

using namespace ebcl;
using namespace opast;

namespace {

/*============================================================================*/


void PrintStreamError(
		char const* const prefix ,
		T_String const& name ,
		X_StreamError const& error )
{
	T_StringBuilder sb;
	sb << prefix << " '" << name << "': " << error.what( );
	if ( error.code( ) == E_StreamError::SYSTEM_ERROR ) {
		sb << " (error code " << error.systemError( ) << ")";
	}
	sb << '\n' << '\0';
	fprintf( stderr , "%s" , sb.data( ) );
}

void WriteSRDError(
		T_StringBuilder& sb ,
		T_SRDError const& error )
{
	sb << error.location( ) << " - " << error.error( ) << "\n";
}


/*============================================================================*/
// FIXME TESTING, MOVE THIS LATER

template<
	typename Enum ,
	typename Storage = uint32_t
> class T_Flags
{
    private:
	Storage flags_;

    public:
	constexpr T_Flags( ) noexcept
		: flags_( 0 ) {}

	constexpr T_Flags( T_Flags const& other ) noexcept
		: flags_( other.flags_ ) { }

	constexpr T_Flags( const Enum flag ) noexcept
		: flags_( 1 << int( flag ) ) {}

	constexpr T_Flags( std::initializer_list< Enum > flags ) noexcept
		: flags_( 0 )
	{
		for ( auto f : flags ) {
			flags_ |= ( 1 << int( f ) );
		}
	}

	explicit constexpr T_Flags( const Storage flags ) noexcept
		: flags_( flags ) { }

	constexpr T_Flags operator |=( T_Flags other ) noexcept
		{ flags_ |= other.flags_; return *this; }
	constexpr T_Flags operator &=( T_Flags other ) noexcept
		{ flags_ &= other.flags_; return *this; }
	constexpr T_Flags operator ^=( T_Flags other ) noexcept
		{ flags_ ^= other.flags_; return *this; }

	constexpr T_Flags operator ~( ) const noexcept
		{ return T_Flags( ~flags_ ); }
	constexpr T_Flags operator &( T_Flags other ) const noexcept
		{ return T_Flags( flags_ & other.flags_ ); }
	constexpr T_Flags operator |( T_Flags other ) const noexcept
		{ return T_Flags( flags_ & other.flags_ ); }
	constexpr T_Flags operator ^( T_Flags other ) const noexcept
		{ return T_Flags( flags_ ^ other.flags_ ); }

	constexpr operator bool( ) const noexcept
		{ return flags_ != 0; }
	constexpr bool operator!( ) const noexcept
		{ return flags_ == 0; }
	explicit constexpr operator Storage( ) const noexcept
		{ return flags_; }

	constexpr bool operator ==( const T_Flags other ) const noexcept
		{ return flags_ == other.flags_; }
	constexpr bool operator !=( const T_Flags other ) const noexcept
		{ return flags_ != other.flags_; }

	constexpr bool isSet( const T_Flags value ) const noexcept
		{ return ( *this & value ) == value; }
	constexpr bool isClear( const T_Flags value ) const noexcept
		{ return !( *this & value ); }
};


template< typename NodeType >
class T_Visitor
{
    public:
	using T_Node = NodeType;

	// Node browser. Returns the Nth child of the specified node, or null
	// if there are no children left.
	using F_NodeBrowser = std::function< T_Node*( T_Node& , uint32_t ) >;

	// Node action. Second parameter indicates whether the action is
	// being called before (false) or after (true) visiting the children.
	using F_NodeAction = std::function< bool( T_Node& , bool ) >;

    private:
	enum E_State_ {
		BEFORE , CHILDREN , AFTER
	};
	struct T_NodeRef_ {
		T_Node* node;
		uint32_t child;
		E_State_ state{ BEFORE };

		explicit T_NodeRef_( T_Node* node )
			: node( node ) , child( 0 ) {}
		explicit T_NodeRef_( T_Node* node , uint32_t child )
			: node( node ) , child( child ) {}
	};

	F_NodeBrowser nodeBrowser_;
	T_Array< T_NodeRef_ > stack_;

    public:
	T_Visitor( ) = delete;
	T_Visitor( T_Visitor const& ) = default;
	T_Visitor( T_Visitor&& ) noexcept = default;

	explicit T_Visitor( F_NodeBrowser browser ) noexcept;

	void visit( T_Node& root , F_NodeAction action );
};

template< typename T >
inline T_Visitor< T >::T_Visitor(
		F_NodeBrowser browser ) noexcept
	: nodeBrowser_( std::move( browser ) )
{ }

template< typename T >
inline void T_Visitor< T >::visit(
		T_Node& root ,
		F_NodeAction action )
{
	stack_.addNew( &root , 0 );

	while ( !stack_.empty( ) ) {
		auto& n( stack_.last( ) );
		switch ( n.state ) {
		    case BEFORE:
			n.state = action( *n.node , false ) ? CHILDREN : AFTER;
			break;
		    case CHILDREN: {
			T_Node* child( nodeBrowser_( *n.node , n.child ++ ) );
			if ( child ) {
				stack_.addNew( child , 0 );
			} else {
				n.state = AFTER;
			}
			break;
		    }
		    case AFTER:
			action( *n.node , true );
			stack_.removeLast( );
			break;
		}
	}
}

A_Node* OpASTBrowser(
		A_Node& node ,
		const uint32_t child )
{
	switch ( node.type( ) ) {

	    // Root node
	    case A_Node::ROOT: {
		auto& n( (T_RootNode&) node );
		if ( child < n.nFunctions( ) ) {
			return &n.function( child );
		}
		break;
	    }

	    // Functions / special blocks
	    case A_Node::DECL_FN: case A_Node::DECL_INIT: case A_Node::DECL_FRAME:
		if ( child == 0 ) {
			auto& n( (A_FuncNode&) node );
			return &n.instructions( );
		}
		break;

	    // Instruction list
	    case A_Node::ILIST: {
		auto& n( (T_InstrListNode&) node );
		if ( child < n.size( ) ) {
			return &n.node( child );
		}
		break;
	    }

	    // Unary operators
	    case A_Node::EXPR_NEG: case A_Node::EXPR_INV:
	    case A_Node::EXPR_NOT: case A_Node::EXPR_SIN:
	    case A_Node::EXPR_COS: case A_Node::EXPR_TAN:
	    case A_Node::EXPR_SQRT: case A_Node::EXPR_EXP:
	    case A_Node::EXPR_LN:
	    {
		auto& n( (T_UnaryOperatorNode&) node );
		if ( child == 0 && n.hasArgument( ) ) {
			return &n.argument( );
		}
		break;
	    }

	    // Binary operators
	    case A_Node::EXPR_ADD: case A_Node::EXPR_SUB:
	    case A_Node::EXPR_MUL: case A_Node::EXPR_DIV:
	    case A_Node::EXPR_POW:
	    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:
	    {
		auto& n( (T_BinaryOperatorNode&) node );
		if ( child == 0 && n.hasLeft( ) ) {
			return &n.left( );
		}
		if ( child < 2 && n.hasRight( ) ) {
			return &n.right( );
		}
		break;
	    }

	    // Nodes that do not have children
	    case A_Node::EXPR_ID: case A_Node::EXPR_CONST:
	    case A_Node::OP_PROGRAM: case A_Node::OP_PIPELINE:
	    case A_Node::OP_INPUT:
		break;

	    // Profile instruction
	    case A_Node::OP_PROFILE:
		if ( child == 0 ) {
			return &( ((T_ProfileInstrNode&) node).instructions( ) );
		}
		break;

	    // Call instruction
	    case A_Node::OP_CALL:
	    {
		auto& n( (T_CallInstrNode&) node );
		if ( child < n.arguments( ) ) {
			return &n.argument( child );
		}
		break;
	    }

	    // Conditional instruction
	    case A_Node::OP_COND:
	    {
		auto& n( (T_CondInstrNode&) node );
		auto c = child;
		if ( n.hasExpression( ) ) {
			if ( c == 0 ) {
				return &n.expression( );
			}
			c --;
		}
		if ( !n.cases( ).empty( ) ) {
			if ( c < n.cases( ).size( ) ) {
				return &n.getCase( n.cases( )[ c ] );
			}
			c -= n.cases( ).size( );
		}
		if ( n.hasDefaultCase( ) && c == 0 ) {
			return &n.defaultCase( );
		}
		break;
	    }

	    // Set instruction
	    case A_Node::OP_SET:
		if ( child == 0 ) {
			auto& n( (T_SetInstrNode&) node );
			if ( n.hasExpression( ) ) {
				return &n.expression( );
			}
		}
		break;

	    // Texture instruction
	    case A_Node::OP_TEXTURE:
	    {
		auto& n( (T_TextureInstrNode&) node );
		auto c = child;
		if ( n.hasWidth( ) ) {
			if ( c == 0 ) {
				return &n.width( );
			}
			c --;
		}
		if ( n.hasHeight( ) && c == 0 ) {
			return &n.height( );
		}
		break;
	    }

	}
	return nullptr;
}


bool checkCalls( T_RootNode& root )
{
	T_Visitor< A_Node > visitor( OpASTBrowser );
	T_Array< T_SRDError > errors;

	T_MultiArray< uint32_t > calls;
	uint32_t cfi;
	for ( cfi = 0 ; cfi < root.nFunctions( ) ; cfi ++ ) {
		calls.next( );
		visitor.visit( root.function( cfi ) , [&]( A_Node& node , bool exit ) -> bool {
			if ( exit || dynamic_cast< A_ExpressionNode* >( &node ) ) {
				return false;
			}
			if ( node.type( ) != A_Node::OP_CALL ) {
				return true;
			}

			auto& call( (T_CallInstrNode&) node );
			const auto callee( root.functionIndex( call.id( ) ) );
			if ( callee >= 0 ) {
				if ( !calls.contains( cfi , callee ) ) {
					calls.add( (uint32_t) callee );
				}

				// Check argument count while we're at it
				auto& fn( (T_FuncNode&) root.function( callee ) );
				if ( fn.arguments( ) != call.arguments( ) ) {
					T_StringBuilder sb;
					sb << "function expects " << fn.arguments( )
						<< " argument" << ( fn.arguments( ) == 1 ? "" : "s" )
						<< ", " << call.arguments( )
						<< " argument" << ( call.arguments( ) == 1 ? "" : "s" )
						<< " provided";
					errors.addNew( std::move( sb ) , call.location( ) );
				}
			} else {
				errors.addNew( "unknown function" , call.idLocation( ) );
			}
			return false;
		} );
	}

	if ( !errors.empty( ) ) {
		T_StringBuilder sb;
		for ( auto const& err : errors ) {
			WriteSRDError( sb , err );
		}
		sb << "Parser failed\n" << '\0';
		fprintf( stderr , "%s" , sb.data( ) );
		return false;
	}

	T_Visitor< uint32_t > callGraphVisitor(
			[&]( uint32_t& v , uint32_t child ) -> uint32_t* {
				const uint32_t nc( calls.sizeOf( v ) );
				if ( child < nc ) {
					return &calls.get( v , child );
				}
				return nullptr;
			} );
	enum class E_CallInfo_ {
		INIT_CHECKED , FRAME_CHECKED ,
		INIT_CALLED , FRAME_CALLED
	};
	using T_CallInfo_ = T_Flags< E_CallInfo_ , uint8_t >;
	T_CallInfo_ callInfo[ calls.size( ) ];

	uint32_t initId( root.functionIndex( "*init*" ) );
	callGraphVisitor.visit( initId , [&]( uint32_t& id , const bool exit ) -> bool {
		if ( exit || callInfo[ id ] & E_CallInfo_::INIT_CALLED ) {
			return false;
		}
		callInfo[ id ] |= E_CallInfo_::INIT_CALLED;
		return true;
	} );
	uint32_t frameId( root.functionIndex( "*frame*" ) );
	callGraphVisitor.visit( frameId , [&]( uint32_t& id , const bool exit ) -> bool {
		if ( exit || callInfo[ id ] & E_CallInfo_::FRAME_CALLED ) {
			return false;
		}
		callInfo[ id ] |= E_CallInfo_::FRAME_CALLED;
		return true;
	} );

	T_StringBuilder sb;
	for ( auto callerId = 0u ; callerId < calls.size( ) ; callerId ++ ) {
		const auto nCallees( calls.sizeOf( callerId ) );
		if ( !nCallees ) {
			continue;
		}
		sb << root.function( callerId ).name( ) << " calls";
		for ( auto i = 0u ; i < nCallees ; i ++ ) {
			auto const& callee( root.function( calls.get( callerId , i ) ) );
			sb << ' ' << callee.name( );

		}
		sb << '\n';
	}
	sb << '\0';
	printf( "%s" , sb.data( ) );

	return true;
}

/*============================================================================*/

} // namespace

int main( int argc , char** argv )
{
	// Open file
	const T_String inputName( argc >= 2 ? argv[ 1 ] : "demo.srd" );
	T_File input( inputName , E_FileMode::READ_ONLY );
	try {
		input.open( );
	} catch ( X_StreamError const& e ) {
		PrintStreamError( "Could not open" , inputName , e );
		return 1;
	}

	// Load SRD data
	T_SRDMemoryTarget srdOut;
	srdOut.clearComments( true ).clearFlushToken( true );
	try {
		T_SRDTextReader srdReader{ srdOut };
		T_FileInputStream fis{ input };
		srdReader.read( inputName , fis );
	} catch ( X_StreamError const& e ) {
		PrintStreamError( "Could not open" , inputName , e );
		return 1;
	} catch ( X_SRDErrors const& e ) {
		T_StringBuilder sb;
		const auto nErrors( e.errors.size( ) );
		for ( auto i = 0u ; i < nErrors ; i ++ ) {
			WriteSRDError( sb , e.errors[ i ] );
		}
		sb << "No parsing happened due to format errors\n" << '\0';
		fprintf( stderr , "%s" , sb.data( ) );
		return 2;
	}

	// Parse the fuck
	T_Parser parser;
	if ( parser.parse( srdOut.list( ) ) ) {
		printf( "Success!\n" );

		auto result( parser.result( ) );
		checkCalls( *result );
		return 0;
	} else {
		T_StringBuilder sb;
		for ( auto const& err : parser.errors( ) ) {
			WriteSRDError( sb , err );
		}
		sb << "Parser failed\n" << '\0';
		fprintf( stderr , "%s" , sb.data( ) );
		return 3;
	}
}