#include "externals.hh"
#include "shaders.hh"
#include "globals.hh"


namespace {

const std::regex PreprocDirective_( "^\\s*//!\\s*([a-z]+(\\s+([^\\s]+))*)\\s*$" );

const std::map< std::string , E_ShaderInput > InputTypes_( ([] {
	std::map< std::string , E_ShaderInput > t;
	t.emplace( "chunk" , E_ShaderInput::CHUNK );
	t.emplace( "library" , E_ShaderInput::LIBRARY );
	t.emplace( "lib" , E_ShaderInput::LIBRARY );
	t.emplace( "vertex" , E_ShaderInput::VERTEX );
	t.emplace( "fragment" , E_ShaderInput::FRAGMENT );
	t.emplace( "compute" , E_ShaderInput::COMPUTE );
	t.emplace( "geo" , E_ShaderInput::GEOMETRY );
	t.emplace( "geometry" , E_ShaderInput::GEOMETRY );
	return t;
})());


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

// Input reader state and functions, used when loading a source file
struct T_InputReader_
{
	using T_Tokens_ = std::vector< std::string >;

	FILE* const file;
	T_ShaderInput& input;
	uint32_t line{ 0 };

	char* buffer{ nullptr };
	ssize_t readCount;

	std::string accumulator{ };
	uint32_t accumLines{ 0 };

	T_InputReader_( __rd__ FILE* const file ,
			__rw__ T_ShaderInput& input )
		: file( file ) , input( input )
	{ }
	~T_InputReader_( );

	void read( );
	void nl( );
	void addAccumulated( );
	void handleDirective(
			__rd__ T_Tokens_ const& tokens );
	void error( __rd__ std::string const& err );
};

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

T_InputReader_::~T_InputReader_( )
{
	if ( buffer ) {
		free( buffer );
	}
	fclose( file );
}

void T_InputReader_::read( )
{
	size_t bsz( 0 );
	while ( ( readCount = getline( &buffer , &bsz , file ) ) != -1 ) {
		line ++;
		std::cmatch match;
		if ( std::regex_match( buffer , match , PreprocDirective_ ) ) {
			const T_Tokens_ tokens( ([](std::string const& a) {
				using stri = std::istream_iterator< std::string >;
				std::istringstream iss( a );
				return T_Tokens_{ stri( iss ) , stri( ) };
			})( match[ 1 ].str( ) ) );
			assert( tokens.size( ) >= 1 );
			handleDirective( tokens );
		} else {
			accumulator += buffer;
			accumLines ++;
		}
	}
	addAccumulated( );
}

void T_InputReader_::nl( )
{
	accumLines ++;
	accumulator += '\n';
}

void T_InputReader_::addAccumulated( )
{
	if ( accumLines ) {
		auto& ck( input.chunks );
		ck.emplace_back( E_ShaderInputChunk::CODE ,
				std::move( accumulator ) , accumLines );
		accumulator = {};
		accumLines = 0;
	}
}

void T_InputReader_::handleDirective(
		__rd__ T_Tokens_ const& tokens )
{
	auto const& directive( tokens[ 0 ] );

	if ( directive == "include" ) {
		if ( tokens.size( ) != 2 ) {
			nl( );
			error( "invalid arguments" );
			return;
		}
		addAccumulated( );
		auto& ck( input.chunks );
		ck.emplace_back( E_ShaderInputChunk::INCLUDE , tokens[ 1 ] , 1 );

	} else if ( directive == "type" ) {
		nl( );
		if ( tokens.size( ) != 2 ) {
			error( "invalid arguments" );
			return;
		}
		auto pos( InputTypes_.find( tokens[ 1 ] ) );
		if ( pos == InputTypes_.end( ) ) {
			error( "unknown type" );
		} else {
			input.type = pos->second;
		}

	} else {
		nl( );
		error( "unknown directive" );
	}
}

void T_InputReader_::error(
		__rd__ std::string const& err )
{
	input.errors.push_back( T_ShaderInputError{
			line , err } );
}

} // namespace


/*= T_ShaderInput ============================================================*/

bool T_ShaderInput::load(
		__rd__ std::string const& path )
{
	type = E_ShaderInput::CHUNK;
	chunks.clear( );
	errors.clear( );

	FILE* const file{ fopen( path.c_str( ) , "r" ) };
	if ( !file ) {
		return false;
	}

	T_InputReader_ reader( file , *this );
	reader.read( );

	return true;
}


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

namespace {

// Code builder, state and functions
struct T_CodeBuilder_
{
	struct T_StackEntry_ {
		std::string name;
		T_ShaderInput const* input;
		uint32_t pos;
	};

	T_ShaderCodeLoader& loader;
	const std::string name;
	T_Frankenshader& code;

	T_ShaderInput const* main;
	std::map< std::string , uint32_t > pos;
	std::vector< T_StackEntry_ > stack;
	std::set< std::string > libraries;
	T_ShaderInput const* current;
	uint32_t cpos{ 0 };
	std::string cname;

	T_CodeBuilder_( __rw__ T_ShaderCodeLoader& loader ,
			__rd__ std::string const& name ,
			__rw__ T_Frankenshader& code )
		: loader( loader ) , name( name ) , code( code ) ,
			main( loader.getInput( name ) )
	{ }

	bool buildCode( );
	void appendChunk( __rd__ T_ShaderInputChunk const& chunk );
	void include( __rd__ std::string const& name ,
			__rd__ const uint32_t lines );
	void next( );
};

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

bool T_CodeBuilder_::buildCode( )
{
	code = T_Frankenshader{ };
	code.files.emplace( name , main != nullptr );
	if ( !main ) {
		return false;
	}
	cname = name;
	current = main;

	while ( cpos < current->chunks.size( ) ) {
		auto& chunk( current->chunks[ cpos ] );
		if ( chunk.type == E_ShaderInputChunk::CODE ) {
			appendChunk( chunk );
		} else {
			include( chunk.text , chunk.lines );
		}

		next( );
	}

	return true;
}

void T_CodeBuilder_::appendChunk(
		__rd__ T_ShaderInputChunk const& chunk )
{
	code.sources.push_back( cname );
	code.counts.push_back( chunk.lines );
	code.starts.push_back( pos[ cname ] );
	code.code += chunk.text;
	pos[ cname ] += chunk.lines;
}

void T_CodeBuilder_::include(
		__rd__ std::string const& nname ,
		__rd__ const uint32_t lines )
{
	const auto prevPos( pos[ cname ] );
	pos[ cname ] += lines;

	// Avoid recursion
	if ( cname == nname || 0 != count_if( stack.begin( ) , stack.end( ) ,
				[nname] ( T_StackEntry_ const& e ) {
					return nname == e.name;
				} ) ) {
		code.errors.push_back( T_ShaderError{ cname , prevPos ,
				"recursive inclusion of '" + nname + "'" } );
		return;
	}

	// Avoid including libraries more than once
	if ( libraries.find( nname ) != libraries.end( ) ) {
		return;
	}

	T_ShaderInput const* const isi( loader.getInput( nname ) );
	code.files.emplace( nname , isi != nullptr );

	// Check for problems
	if ( !isi ) {
		// Not found
		code.errors.push_back( T_ShaderError{ cname , prevPos ,
				"file not found" } );
		return;
	}
	if ( isi->type != E_ShaderInput::CHUNK && isi->type != E_ShaderInput::LIBRARY ) {
		// Trying to load a top-level shader
		code.errors.push_back( T_ShaderError{ cname , prevPos ,
				"trying to include a top-level file" } );
		return;
	}

	// Add input loader errors
	if ( libraries.find( nname ) == libraries.end( ) ) {
		for ( auto const& errs : isi->errors ) {
			code.errors.push_back( T_ShaderError{
					nname , errs.line , errs.error } );
		}
	}
	libraries.insert( nname );

	// Enter the new file
	stack.push_back( T_StackEntry_{ cname , current , cpos } );
	cname = nname;
	current = isi;
	cpos = UINT32_MAX;
	pos[ cname ] = 1;
}

void T_CodeBuilder_::next( )
{
	cpos ++;
	while ( cpos == current->chunks.size( ) && !stack.empty( ) ) {
		T_StackEntry_ const& se( stack[ stack.size( ) - 1 ] );
		pos.erase( cname );
		cpos = se.pos + 1;
		current = se.input;
		cname = se.name;
		stack.pop_back( );
	}
}


} // namespace


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

T_ShaderCodeLoader::T_ShaderCodeLoader(
		__rw__ T_FilesWatcher& watcher )
	: watcher_( watcher )
{
}

bool T_ShaderCodeLoader::load(
		__rd__ std::string const& name ,
		__wr__ T_Frankenshader& code )
{
	T_CodeBuilder_ cb( *this , name , code );
	const bool rv( cb.buildCode( ) );

	// Update dependencies
	for ( auto const& dep : code.files ) {
		deps_[ name ].insert( dep.first );
	}

	return rv;
}

T_ShaderInput const* T_ShaderCodeLoader::getInput(
		__rd__ std::string const& name )
{
	auto pos( inputs_.find( name ) );
	if ( pos != inputs_.end( ) ) {
		return pos->second.get( );
	}
	T_ShaderInput ni;
	if ( !ni.load( "shaders/" + name ) ) {
		return nullptr;
	}
	inputs_.emplace( name , std::make_unique< T_ShaderInput >( std::move( ni ) ) );
	return inputs_.find( name )->second.get( );
}

void T_ShaderCodeLoader::removeInput(
		__rd__ std::string const& name )
{
	inputs_.erase( name );
}

void T_ShaderCodeLoader::onFileUpdated(
		__rd__ std::string const& name )
{
	inputs_.erase( name );
	auto& deps( deps_[ name ] );
	for ( auto const& dep : deps ) {
		pending_.insert( dep );
	}
}


bool T_ShaderPipeline::valid( ) const noexcept
{
	return Globals::Shaders( ).pipelines_[ smIndex_ ].id != 0;
}


#if 1
void testLoadShaderFile( )
{
	const std::string source( "test-loader.glsl" );
	T_ShaderCodeLoader loader( Globals::Watcher( ) );
	T_Frankenshader code;
	if ( loader.load( source , code ) ) {
		printf( "SUCCESS! TYPE = %d\n" , int( code.type ) );
		printf( "FILES:\n" );
		for ( auto const& e : code.files ) {
			printf( " -> %s (%s)\n" , e.first.c_str( ) ,
					e.second ? "found" : "missing" );
		}
		if ( code.errors.size( ) ) {
			printf( "ERRORS:\n" );
			for ( auto const& err : code.errors ) {
				printf( "%s:%d: %s\n" , err.source.c_str( ) , err.line , err.error.c_str( ) );
			}
		}
	} else {
		printf( "fail :'(\n" );
	}
}
#endif