#include "externals.hh"
#include "utilities.hh"
#include "uniforms.hh"


/*= Uniform types ============================================================*/

E_UniformCategory GetUniformCategory(
	__rd__ const E_UniformType type )
{
	switch ( type ) {
		case E_UniformType::F1:
		case E_UniformType::F2:
		case E_UniformType::F3:
		case E_UniformType::F4:
			return E_UniformCategory::FLOAT;
		case E_UniformType::I1:
		case E_UniformType::I2:
		case E_UniformType::I3:
		case E_UniformType::I4:
			return E_UniformCategory::INT;
		case E_UniformType::SAMPLER2D:
			return E_UniformCategory::SAMPLER2D;
	}
	std::abort( );
}


uint32_t GetElementCount(
	__rd__ const E_UniformType type )
{
	switch ( type ) {
		case E_UniformType::F1:
		case E_UniformType::I1:
			return 1;
		case E_UniformType::F2:
		case E_UniformType::I2:
			return 2;
		case E_UniformType::F3:
		case E_UniformType::I3:
			return 3;
		case E_UniformType::F4:
		case E_UniformType::I4:
			return 4;

		case E_UniformType::SAMPLER2D:
			return 0;
	}
	std::abort( );
}

bool IsVectorCategory(
	__rd__ const E_UniformCategory category )
{
	assert( category != E_UniformCategory::__COUNT__ );
	return category == E_UniformCategory::FLOAT
		|| category == E_UniformCategory::INT;
}


/*= PackUniforms =============================================================*/

namespace {

T_PackedUniforms PackVectorUniforms_(
		__rd__ T_UniformDeclarations const& declarations ,
		__rd__ const E_UniformCategory category )
{
	std::vector< std::string > packOrder;
	packOrder.reserve( declarations.size( ) );
	for ( auto entry : declarations ) {
		auto const& u( entry.second );
		if ( GetUniformCategory( u.type ) == category ) {
			packOrder.push_back( entry.first );
		}
	}
	std::sort( packOrder.begin( ) , packOrder.end( ) ,
		[&declarations]( auto const& a , auto const& b )
		{
			auto const& ua( declarations.find( a )->second );
			auto const& ub( declarations.find( b )->second );
			return GetElementCount( ub.type ) < GetElementCount( ua.type );
		} );

	const auto n( packOrder.size( ) );
	auto index = 0u;
	T_PackedUniforms pu;

	// Place vec4's
	while ( index < n && GetElementCount( declarations.find( packOrder[ index ] )->second.type ) == 4 ) {
		pu.positions[ packOrder[ index ] ] = pu.elements;
		pu.elements += 4;
		index ++;
	}

	// Place vec3's
	uint32_t s3left( 0 );
	while ( index < n && GetElementCount( declarations.find( packOrder[ index ] )->second.type ) == 3 ) {
		pu.positions[ packOrder[ index ] ] = pu.elements;
		pu.elements += 4;
		s3left ++;
		index ++;
	}

	// Place vec2's
	const uint32_t s2s( pu.elements );
	bool s2spare( false );
	while ( index < n && GetElementCount( declarations.find( packOrder[ index ] )->second.type ) == 2 ) {
		if ( !s2spare ) {
			pu.elements += 4;
		}
		const uint32_t pos( pu.elements - ( s2spare ? 2 : 4 ) );
		pu.positions[ packOrder[ index ] ] = pos;
		s2spare = !s2spare;
		index ++;
	}

	// Place single elements
	uint32_t s1pos( s2spare ? ( pu.elements - 2 ) : pu.elements );
	while ( index < n ) {
		uint32_t pos;
		if ( s3left ) {
			pos = s2s - s3left * 4 + 3;
			s3left --;
		} else {
			if ( s1pos == pu.elements ) {
				pu.elements += 4;
			}
			pos = s1pos;
			s1pos ++;
		}
		pu.positions[ packOrder[ index ] ] = pos;
		index ++;
	}

	return pu;
}

T_PackedUniforms PackOpaqueUniforms_(
		__rd__ T_UniformDeclarations const& declarations ,
		__rd__ const E_UniformCategory category )
{
	T_PackedUniforms pu;
	for ( auto entry : declarations ) {
		auto const& u( entry.second );
		if ( GetUniformCategory( u.type ) == category ) {
			pu.positions[ entry.first ] = pu.elements;
			pu.elements ++;
		}
	}
	return pu;
}

} // namespace

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

T_PackedUniforms PackUniforms(
		__rd__ T_UniformDeclarations const& declarations ,
		__rd__ const E_UniformCategory category )
{
	assert( category != E_UniformCategory::__COUNT__ );
	if ( IsVectorCategory( category ) ) {
		return PackVectorUniforms_( declarations , category );
	}
	return PackOpaqueUniforms_( declarations , category );
}


/*= T_UniformSet =============================================================*/

constexpr uint32_t T_UniformSet::Invalid;

T_UniformSet::T_UniformSet( )
	: vectorElements_( 0 ) , opaqueElements_( 0 )
{ }

T_UniformSet::T_UniformSet(
		__rw__ T_UniformDeclarations declarations )
	: declarations_( std::move( declarations ) ) ,
		vectorElements_( 0 ) , opaqueElements_( 0 )
{
	for ( auto i = 0 ; i < int( E_UniformCategory::__COUNT__ ) ; i ++ ) {
		const auto cat{ E_UniformCategory( i ) };
		packed_[ i ] = PackUniforms( declarations_ , cat );
		if ( IsVectorCategory( cat ) ) {
			vectorElements_ += packed_[ i ].elements;
		} else {
			opaqueElements_ += packed_[ i ].elements;
		}
	}
}

uint32_t T_UniformSet::position(
	__rd__ std::string const& uniform ) const
{
	const auto pos( declarations_.find( uniform ) );
	if ( pos == declarations_.end( ) ) {
		return Invalid;
	}

	const int ci( int( GetUniformCategory( pos->second.type ) ) );
	const auto ePos( packed_[ ci ].positions.find( uniform ) );
	assert( ePos != packed_[ ci ].positions.end( ) );
	return ePos->second;
}


/*= T_UniformStorage =========================================================*/

T_UniformStorage::T_UniformStorage(
		__rw__ T_UniformSet&& uniforms )
	: uniforms_( std::move( uniforms ) )
{
	values_.resize( uniforms_.opaqueElements( )
			+ uniforms_.vectorElements( ) , 0 );

	uint32_t pos = 0;
	for ( auto i = 0 ; i < int( E_UniformCategory::__COUNT__ ) ; i ++ ) {
		starts_[ i ] = pos;
		pos += uniforms_.inCategory( E_UniformCategory( i ) );
	}
}

float* T_UniformStorage::getf(
	__rd__ std::string const& name )
{
	auto const& dPos( uniforms_.declarations( ).find( name ) );
	if ( dPos == uniforms_.declarations( ).end( ) ) {
		return nullptr;
	}
	auto const& decl( dPos->second );

	const auto cat( GetUniformCategory( decl.type ) );
	const auto pos( starts_[ int( cat  ) ] + uniforms_.position( name ) );
	return reinterpret_cast< float* >( &values_[ pos ] );
}

int* T_UniformStorage::geti(
	__rd__ std::string const& name )
{
	auto const& dPos( uniforms_.declarations( ).find( name ) );
	if ( dPos == uniforms_.declarations( ).end( ) ) {
		return nullptr;
	}
	auto const& decl( dPos->second );

	const auto cat( GetUniformCategory( decl.type ) );
	const auto pos( starts_[ int( cat  ) ] + uniforms_.position( name ) );
	return &values_[ pos ];
}


/*= T_ProgramUniforms ========================================================*/

T_ProgramUniforms::T_ProgramUniforms(
		__rw__ T_UniformSet uniforms )
	: T_UniformStorage( std::move( uniforms ) )
{ }

void T_ProgramUniforms::apply(
	__rd__ const GLuint id )
{
	GLuint uLoc( 0 );
	for ( auto i = 0 ; i < int( E_UniformCategory::__COUNT__ ) ; i ++ ) {
		const auto cat{ E_UniformCategory( i ) };
		const auto count{ uniforms_.inCategory( cat ) };
		if ( count != 0 ) {
			switch ( cat ) {
				case E_UniformCategory::FLOAT:
					glProgramUniform4fv( id , uLoc , count / 4 ,
							reinterpret_cast< float* >( &values_[ starts_[ i ] ] ) );
					uLoc += count / 4;
					break;
				case E_UniformCategory::INT:
					glProgramUniform4iv( id , uLoc , count / 4 ,
							 &values_[ starts_[ i ] ] );
					uLoc += count / 4;
					break;
				case E_UniformCategory::SAMPLER2D:
					glProgramUniform1iv( id , uLoc , count ,
							 &values_[ starts_[ i ] ] );
					uLoc += count;
					break;

				case E_UniformCategory::__COUNT__:
					std::abort( );
			}
		}
	}
}


/*= T_UniformBuffer ==========================================================*/

T_UniformBuffer::T_UniformBuffer(
		__rw__ T_UniformSet set )
	: T_UniformStorage( std::move( set ) ) , buffer_( 0 )
{
	assert( uniforms_.opaqueElements( ) == 0 );

	glGenBuffers( 1 , &buffer_ );
	glBindBuffer( GL_UNIFORM_BUFFER , buffer_ );
	glBufferData( GL_UNIFORM_BUFFER , values_.size( ) * 4 , &values_[ 0 ] , GL_STREAM_DRAW );
	GL_ASSERT( );
}

T_UniformBuffer::T_UniformBuffer(
		__rw__ T_UniformBuffer&& ori ) noexcept
	: T_UniformStorage( std::move( ori ) ) , buffer_( 0 )
{
	std::swap( buffer_ , ori.buffer_ );
}

T_UniformBuffer& T_UniformBuffer::operator=(
		__rw__ T_UniformBuffer&& other ) noexcept
{
	if ( buffer_ ) {
		glBindBuffer( GL_UNIFORM_BUFFER , 0 );
		glDeleteBuffers( 1 , &buffer_ );
		buffer_ = 0;
	}
	std::swap( buffer_ , other.buffer_ );
	T_UniformStorage::operator=( std::move( other ) );
	return *this;
}

void T_UniformBuffer::update( )
{
	glBindBuffer( GL_UNIFORM_BUFFER , buffer_ );
	glInvalidateBufferData( GL_UNIFORM_BUFFER );
	glBufferData( GL_UNIFORM_BUFFER , values_.size( ) * 4 , &values_[ 0 ] , GL_STREAM_DRAW );
}

void T_UniformBuffer::bindTo(
		__rd__ const GLuint binding )
{
	glBindBufferBase( GL_UNIFORM_BUFFER , binding , buffer_ );
}

T_UniformBuffer::~T_UniformBuffer( )
{
	if ( buffer_ ) {
		glBindBuffer( GL_UNIFORM_BUFFER , 0 );
		glDeleteBuffers( 1 , &buffer_ );
	}
}