#include "externals.hh"

#include "c-opast.hh"
#include "c-sync.hh"

#include <ebcl/Algorithms.hh>

using namespace ebcl;
using namespace opast;


/*= A_Node ===================================================================*/

A_Node::A_Node( const E_Type type ,
		A_Node* const parent ) noexcept
	: type_( type ) , parent_( parent )
{
	assert( ( type == ROOT && !parent ) || ( type != ROOT && parent ) );
}

A_Node::~A_Node( ) { }

T_RootNode& A_Node::root( ) const noexcept
{
	A_Node const* node( this );
	while ( node->parent_ ) {
		node = node->parent_;
	}
	assert( node );
	assert( node->type_ == ROOT );
	return *dynamic_cast< T_RootNode* >(
			const_cast< A_Node* >( node ) );
}

bool A_Node::replace(
		A_Node const&		node ,
		T_OwnPtr< A_Node >&	replacement
		) noexcept
{
	const auto nc{ size( ) };
	for ( auto i = 0u ; i < nc ; i ++ ) {
		if ( child( i ).get( ) == &node ) {
			child( i ) = std::move( replacement );
			return true;
		}
	}
	for ( auto i = 0u ; i < nc ; i ++ ) {
		if ( child( i ) && child( i )->replace( node , replacement ) ) {
			return true;
		}
	}
	return false;
}

/*----------------------------------------------------------------------------*/

A_Node* opast::ASTVisitorBrowser(
		A_Node& node ,
		const uint32_t child ) noexcept
{
	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::EXPR_INPUT:
	    //
	    case A_Node::OP_PROGRAM: case A_Node::OP_PIPELINE:
	    case A_Node::OP_INPUT: case A_Node::OP_FULLSCREEN:
	    case A_Node::OP_ODBG: case A_Node::OP_LOCALS:
	    case A_Node::OP_MAINOUT: case A_Node::OP_OVERRIDES:
	    //
	    case A_Node::OP_USE_FRAMEBUFFER: case A_Node::OP_USE_PIPELINE:
	    case A_Node::OP_USE_PROGRAM: case A_Node::OP_USE_TEXTURE:
		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 );
		const auto nArgs( n.size( ) );
		if ( child < nArgs ) {
			return &n.argument( nArgs - child - 1 );
		}
		break;
	    }

	    // Argument node
	    case A_Node::TN_ARG:
	    {
		if ( child == 0 ) {
			return &( (T_ArgumentNode&) node ).expression( );
		}
		break;
	    }

	    // Instructions that match the common case, in which chidren are
	    // listed directly
	    case A_Node::OP_CLEAR: case A_Node::OP_COMPUTE:
	    {
		const auto nArgs( node.size( ) );
		if ( child < nArgs ) {
			return node.child( nArgs - child - 1 ).get( );
		}
		break;
	    }

	    // Instructions that may have undefined children at the end of
	    // their lists; we need to find the first non-NULL child in the
	    // list and start from there.
	    case A_Node::OP_IMAGE:
	    {
		auto nArgs( node.size( ) );
		while ( nArgs > 0 && !node.child( nArgs - 1 ) ) {
			nArgs --;
		}
		if ( child < nArgs ) {
			return node.child( nArgs - child - 1 ).get( );
		}
		break;
	    }

	    // Conditional instruction & associated technical nodes
	    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;
	    }

	    case A_Node::TN_CONDITION:
		if ( child == 0 ) {
			return &( (T_CondInstrNode::T_Expression&) node ).expression( );
		}
		break;
	    case A_Node::TN_CASE:
		if ( child == 0 ) {
			return &( (T_CondInstrNode::T_ValuedCase&) node ).instructions( );
		}
		break;
	    case A_Node::TN_DEFAULT:
		if ( child == 0 ) {
			return &( (T_CondInstrNode::T_DefaultCase&) node ).instructions( );
		}
		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.lods( ) ) {
			if ( c == 0 ) {
				return n.lods( );
			}
			c --;
		}
		if ( n.hasHeight( ) ) {
			if ( c == 0 ) {
				return &n.height( );
			}
			c --;
		}
		if ( n.hasWidth( ) && c == 0 ) {
			return &n.width( );
		}
		return nullptr;
	    }

	    // Framebuffer definition; may have LOD expressions
	    case A_Node::OP_FRAMEBUFFER:
	    {
		auto& n( (T_FramebufferInstrNode&) node );
		if ( child < n.colorAttachments( ) ) {
			return &n.colorAttachment( child );
		}
		if ( child == n.colorAttachments( ) && n.depthAttachment( ) ) {
			return n.depthAttachment( );
		}
		break;
	    }
	    case A_Node::TN_FBATT:
	    {
		    auto& n( (T_FramebufferInstrNode::T_Attachment&) node );
		    if ( n.lod( ) && child == 0 ) {
			    return n.lod( );
		    }
		    break;
	    }

	    // Sampler definition may have LOD expressions
	    case A_Node::OP_SAMPLER:
	    {
		auto& n( (T_SamplerInstrNode&) node );
		auto c = child;
		if ( n.minLod( ) ) {
			if ( c == 0 ) {
				return n.minLod( );
			}
			c--;
		}
		if ( n.maxLod( ) && c == 0 ) {
			return n.maxLod( );
		}
		break;
	    }

	    // Uniforms instruction
	    case A_Node::OP_UNIFORMS:
	    {
		    auto& n( (T_UniformsInstrNode&) node );
		    const auto nv( n.size( ) );
		    if ( child < nv ) {
			    return &n.value( nv - child - 1 );
		    }
		    break;
	    }

	    // Viewport instruction
	    case A_Node::OP_VIEWPORT:
	    {
		auto& n( (T_ViewportInstrNode&) node );
		auto c = child;
		for ( int i = 0 ; i < 4 ; i ++ ) {
			T_ViewportInstrNode::E_Parameter p{
				T_ViewportInstrNode::E_Parameter( 3 - i ) };
			if ( n.hasParameter( p ) ) {
				if ( c == 0 ) {
					return &n.parameter( p );
				}
				c --;
			}
		}
		break;
	    }

	}
	return nullptr;
}


/*= E_DataType ===============================================================*/

T_StringBuilder& opast::operator<<(
		T_StringBuilder& sb ,
		const E_DataType dt )
{
	switch ( dt ) {
		case E_DataType::UNKNOWN:
			sb << "unknown";
			return sb;
		case E_DataType::BUILTIN:
			sb << "built-in variable";
			return sb;
		case E_DataType::VARIABLE:
			sb << "variable";
			return sb;
		case E_DataType::FRAMEBUFFER:
			sb << "framebuffer";
			return sb;
		case E_DataType::INPUT:
			sb << "input";
			return sb;
		case E_DataType::PIPELINE:
			sb << "pipeline";
			return sb;
		case E_DataType::PROGRAM:
			sb << "program";
			return sb;
		case E_DataType::SAMPLER:
			sb << "sampler";
			return sb;
		case E_DataType::TEXTURE:
			sb << "texture";
			return sb;
	}

	sb << "*unknown data type*";
	return sb;
}

/*= A_FuncNode ===============================================================*/

A_FuncNode::A_FuncNode(
		const bool isInit ,
		T_RootNode* const root ) noexcept
	: A_Node( isInit ? DECL_INIT : DECL_FRAME , root ) ,
		name_( isInit ? "*init*" : "*frame*" )
{ }

A_FuncNode::A_FuncNode(
		T_String const& name ,
		T_RootNode* const root ) noexcept
	: A_Node( DECL_FN , root ) , name_( name )
{ }

T_Optional< T_SRDLocation > A_FuncNode::addLocalVariable(
		T_String const& name ,
		T_SRDLocation const& location ) noexcept
{
	auto const* const pnp( locals_.get( name ) );
	if ( pnp ) {
		return pnp->location;
	}
	T_Local_ lv{ name , location , false };
	lv.type = E_DataType::VARIABLE;
	locals_.add( std::move( lv ) );
	return {};
}


/*= T_RootNode ===============================================================*/

T_RootNode::T_RootNode( ) noexcept
	: A_Node( ROOT , nullptr )
{ }

T_RootNode::T_AddFunctionResult T_RootNode::addFunction(
		T_OwnPtr< A_FuncNode >& function ) noexcept
{
	T_String const& fn( function->name( ) );
	auto const* const pf( functions_.get( fn ) );
	if ( !pf ) {
		auto* const rv( function.get( ) );
		functions_.add( fn , children_.size( ) );
		children_.add( std::move( function ) );
		return *rv;
	}

	T_String dfn;
	uint32_t dupCtr( 0 );
	do {
		T_StringBuilder fnsb;
		fnsb << fn << " dup " << dupCtr ++;
		dfn = std::move( fnsb );
	} while ( functions_.contains( dfn ) );

	T_OwnPtr< A_FuncNode > df( NewOwned< T_FuncNode >( dfn , *this ) );
	auto* const rv( df.get( ) );
	functions_.add( fn , children_.size( ) );
	children_.add( std::move( df ) );
	return T_AddFunctionResult{ *rv , child( *pf )->location( ) };
}


/*= T_FuncNode ===============================================================*/

T_Optional< T_SRDLocation > T_FuncNode::addArgument(
		T_SRDToken const& token ) noexcept
{
	assert( token.type( ) == E_SRDTokenType::WORD );
	assert( token.hasLocation( ) );

	const auto pnp( locals_.indexOf( token.stringValue( ) ) );
	if ( pnp != T_HashIndex::INVALID_INDEX ) {
		return locals_[ pnp ].location;
	}
	locals_.add( T_Local_{ token.stringValue( ) , token.location( ) , true } );
	return {};
}

uint32_t T_FuncNode::arguments( ) const noexcept
{
	const auto n( locals_.size( ) );
	uint32_t na{ 0u };

	for ( auto i = 0u ; i < n ; i ++ ) {
		if ( locals_[ i ].argument ) {
			na ++;
		}
	}

	return na;
}

/*= T_InstrListNode ==========================================================*/

void T_InstrListNode::replaceMultiple(
		A_InstructionNode&	instruction ,
		T_InstrListNode*	replacement
		) noexcept
{
	T_OwnPtr< A_InstructionNode > oldInstr;
	uint32_t index{ 0xffffffff };
	for ( auto i = 0u ; i < children_.size( ) ; i ++ ) {
		if ( children_[ i ].get( ) == &instruction ) {
			oldInstr = std::move( children_[ i ] );
			index = i;
			break;
		}
	}
	assert( oldInstr );

	if ( replacement && replacement->size( ) ) {
		auto& rep{ *replacement };

		// Move existing children to the end of the list
		const auto oldSize{ children_.size( ) };
		for ( auto i = 1u ; i < rep.size( ) ; i ++ ) {
			children_.addNew( );
		}
		const auto newSize{ children_.size( ) };
		const auto nMove{ rep.size( ) > 1 ? ( oldSize - index ) : 0 };
		for ( auto i = 1u ; i < nMove ; i ++ ) {
			children_[ newSize - i ] = std::move(
					children_[ oldSize - i ] );
		}

		// Move replacement instructions
		for ( auto i = 0u ; i < rep.size( ) ; i ++ ) {
			children_[ index + i ] = std::move(
					rep.children_[ i ] );
			children_[ index + i ]->setParent( *this );
		}

		rep.children_.clear( );
	} else {
		children_.remove( index );
	}
}


/*= T_PipelineInstrNode ======================================================*/

T_Optional< T_SRDLocation > T_PipelineInstrNode::addProgram(
		T_SRDToken const& pidToken ) noexcept
{
	T_String const& name( pidToken.stringValue( ) );
	const auto pIndex( pids_.indexOf( name ) );
	if ( pIndex != -1 ) {
		return pidLocations_[ pIndex ];
	}
	pids_.add( name );
	pidLocations_.add( pidToken.location( ) );
	return {};
}


/*= T_UnaryOperatorNode ======================================================*/

T_UnaryOperatorNode::T_UnaryOperatorNode(
		A_Node& parent ,
		const E_Operator op ) noexcept
	: A_ExpressionNode( ([op]() {
		switch ( op ) {
			case NEG: return EXPR_NEG;
			case INV: return EXPR_INV;
			case NOT: return EXPR_NOT;
			case SIN: return EXPR_SIN;
			case COS: return EXPR_COS;
			case TAN: return EXPR_TAN;
			case SQRT: return EXPR_SQRT;
			case EXP: return EXPR_EXP;
			case LN: return EXPR_LN;
		}
		std::abort( );
	})( ) , parent ) , op_( op )
{
	children_.add( P_ExpressionNode{} );
}


/*= T_BinaryOperatorNode =====================================================*/

T_BinaryOperatorNode::T_BinaryOperatorNode(
		A_Node& parent ,
		const E_Operator op ) noexcept
	: A_ExpressionNode( ([op]() {
		switch ( op ) {
			case ADD: return EXPR_ADD;
			case SUB: return EXPR_SUB;
			case MUL: return EXPR_MUL;
			case DIV: return EXPR_DIV;
			case POW: return EXPR_POW;
			case CMP_EQ: return EXPR_CMP_EQ;
			case CMP_NE: return EXPR_CMP_NE;
			case CMP_GT: return EXPR_CMP_GT;
			case CMP_GE: return EXPR_CMP_GE;
			case CMP_LT: return EXPR_CMP_LT;
			case CMP_LE: return EXPR_CMP_LE;
		}
		std::abort( );
	})( ) , parent ) , op_( op )
{
	children_.add( P_ExpressionNode{} );
	children_.add( P_ExpressionNode{} );
}


/*= T_FramebufferInstrNode ===================================================*/

bool T_FramebufferInstrNode::hasAttachment(
		T_String const& name ) const noexcept
{
	const auto nc{ size( ) };
	for ( auto i = 0u ; i < nc ; i ++ ) {
		auto const& att{ (T_Attachment const&) *child( i ) };
		if ( att.id( ) == name ) {
			return true;
		}
	}

	return false;
}

bool T_FramebufferInstrNode::addColorAttachment(
		T_SRDToken const& id ,
		P_ExpressionNode lod ) noexcept
{
	if ( hasAttachment( id.stringValue( ) ) ) {
		return false;
	}
	colorAttachments_.add( children_.add(
		NewOwned< T_Attachment >( *this ,
			false , id , std::move( lod ) ) ) );
	return true;
}

bool T_FramebufferInstrNode::setDepthAttachment(
		T_SRDToken const& id ,
		P_ExpressionNode lod ) noexcept
{
	if ( depthAttachment_ || hasAttachment( id.stringValue( ) ) ) {
		return false;
	}
	depthAttachment_ = children_.add( NewOwned< T_Attachment >(
			*this , true , id , std::move( lod ) ) );
	return true;
}


/*= T_OutputDebugInstrNode ===================================================*/

T_OutputDebugInstrNode::T_OutputDebugInstrNode(
		T_InstrListNode& parent ,
		T_SRDToken const& texture ,
		const E_ODbgMode mode ,
		T_SRDLocation const& modeLocation ,
		T_SRDToken const& description ) noexcept
	: A_InstructionNode( OP_ODBG , parent ) ,
		idTexture_( texture.stringValue( ) ) ,
		locTexture_( texture.location( ) ) ,
		mode_( mode ) , locMode_( modeLocation ) ,
		description_( description.stringValue( ) ) ,
		locDescription_( description.location( ) )
{ }


/*= T_SamplerInstrNode =========================================================*/

T_Optional< T_SRDLocation > T_SamplerInstrNode::setSampling(
		E_TexSampling mode ,
		T_SRDLocation const& location ) noexcept
{
	if ( sampling_ ) {
		return sampling_->location;
	}
	sampling_ = T_Sampling_{ location , mode };
	return {};
}

T_Optional< T_SRDLocation > T_SamplerInstrNode::setMipmapSampling(
		E_TexSampling mode ,
		T_SRDLocation const& location ) noexcept
{
	if ( mipmaps_ ) {
		return mipmaps_->location;
	}
	mipmaps_ = T_Mipmaps_{ location , mode };
	return {};
}

T_Optional< T_SRDLocation > T_SamplerInstrNode::setNoMipmap(
		T_SRDLocation const& location ) noexcept
{
	if ( mipmaps_ ) {
		return mipmaps_->location;
	}
	mipmaps_ = T_Mipmaps_{ location , {} };
	return {};
}

T_Optional< T_SRDLocation > T_SamplerInstrNode::setWrapping(
		E_TexWrap mode ,
		T_SRDLocation const& location ) noexcept
{
	if ( wrapping_ ) {
		return wrapping_->location;
	}
	wrapping_ = T_Wrapping_{ location , mode };
	return {};
}

T_Optional< T_SRDLocation > T_SamplerInstrNode::setLOD(
		T_SRDLocation const& location ,
		P_ExpressionNode min ,
		P_ExpressionNode max ) noexcept
{
	if ( lod_ ) {
		return lod_->location;
	}

	T_Optional< uint32_t > pMin , pMax;
	if ( min ) {
		pMin = children_.add( NewOwned< T_ArgumentNode >(
				*this , std::move( min ) ) );
	}
	if ( max ) {
		pMax = children_.add( NewOwned< T_ArgumentNode >(
				*this , std::move( max ) ) );
	}

	lod_ = T_LOD_{ location , pMin , pMax };
	return {};
}


/*= T_LocalsInstrNode ===========================================================*/

T_Optional< T_SRDLocation > T_LocalsInstrNode::addVariable(
		T_SRDToken const& token ) noexcept
{
	const auto pnp( vars_.indexOf( token.stringValue( ) ) );
	if ( pnp != -1 ) {
		return varLocs_[ pnp ];
	}
	vars_.add( token.stringValue( ) );
	varLocs_.add( token.location( ) );
	return {};
}


/*= T_CondInstrNode ==========================================================*/

void T_CondInstrNode::handleRemoval(
		const uint32_t idx ) noexcept
{
	children_.removeSwap( idx );
	if ( idx >= children_.size( ) ) {
		return;
	}
	if ( defaultCase_ && idx == *defaultCase_ ) {
		defaultCase_ = idx;
	} else if ( expression_ && idx == *expression_ ) {
		expression_ = idx;
	} else {
		auto& vc{ (T_ValuedCase&) *child( idx ) };
		cases_.update( vc.value( ) , idx );
	}
}

void T_CondInstrNode::setExpression(
		P_ExpressionNode expression ) noexcept
{
	if ( !expression ) {
		if ( !expression_ ) {
			return;
		}
		const auto idx{ *expression_ };
		expression_.clear( );
		handleRemoval( idx );

	} else if ( expression_ ) {
		child( *expression_ ) = NewOwned< T_Expression >(
				*this , std::move( expression ) );
	} else {
		expression_ = children_.size( );
		children_.add( NewOwned< T_Expression >(
				*this , std::move( expression ) ) );
	}
}

void T_CondInstrNode::setCase(
		const int64_t value ,
		P_InstrListNode instrList ) noexcept
{
	auto const* const pp{ cases_.get( value ) };
	if ( instrList && pp ) {
		child( *pp ) = NewOwned< T_ValuedCase >( *this , value ,
				std::move( instrList ) );
	} else if ( instrList ) {
		cases_.add( value , children_.size( ) );
		children_.add( NewOwned< T_ValuedCase >( *this , value ,
				std::move( instrList ) ) );
	} else if ( pp ) {
		const auto idx{ *pp };
		cases_.remove( value );
		handleRemoval( idx );
	}
}

void T_CondInstrNode::setDefaultCase(
		P_InstrListNode defaultCase ) noexcept
{
	if ( !defaultCase ) {
		if ( !defaultCase_ ) {
			return;
		}
		const auto idx{ *defaultCase_ };
		defaultCase_.clear( );
		handleRemoval( idx );

	} else if ( defaultCase_ ) {
		child( *defaultCase_ ) = NewOwned< T_DefaultCase >(
				*this , std::move( defaultCase ) );
	} else {
		defaultCase_ = children_.size( );
		children_.add( NewOwned< T_DefaultCase >(
				*this , std::move( defaultCase ) ) );
	}
}