// Copyright (C) 2009-2021, Panagiotis Christopoulos Charitos and contributors.
// All rights reserved.
// Code licensed under the BSD License.
// http://www.anki3d.org/LICENSE

#pragma once

#include <AnKi/Resource/Common.h>
#include <AnKi/Gr/ShaderProgram.h>
#include <AnKi/Util/HashMap.h>
#include <AnKi/Util/StringList.h>
#include <AnKi/ShaderCompiler/ShaderProgramBinary.h>

namespace anki {

/// @addtogroup resource
/// @{

/// This is a ray tracing library. Essentially a shader program with some functionality on how to get group indices.
class ShaderProgramRaytracingLibrary
{
	friend class ShaderProgramResourceSystem;

public:
	~ShaderProgramRaytracingLibrary()
	{
		m_libraryName.destroy(m_alloc);
		m_resourceHashToShaderGroupHandleIndex.destroy(m_alloc);
	}

	CString getLibraryName() const
	{
		return m_libraryName;
	}

	U32 getRayTypeCount() const
	{
		ANKI_ASSERT(m_rayTypeCount < MAX_U32);
		return m_rayTypeCount;
	}

	const ShaderProgramPtr& getShaderProgram() const
	{
		return m_program;
	}

	/// Given the filename of a program (that contains ray tracing shaders) and a specific mutation get the shader
	/// handle index.
	U32 getShaderGroupHandleIndex(CString resourceFilename, U64 mutationHash) const
	{
		return getIndex(generateShaderGroupGroupHash(resourceFilename, mutationHash, m_alloc));
	}

private:
	GenericMemoryPoolAllocator<U8> m_alloc;
	String m_libraryName;
	U32 m_rayTypeCount = MAX_U32;
	ShaderProgramPtr m_program;
	HashMap<U64, U32> m_resourceHashToShaderGroupHandleIndex;

	/// Given the filename of a program (that contains ray tracing shaders) and a specific mutation get a hash back.
	static U64 generateShaderGroupGroupHash(CString resourceFilename, U64 mutationHash,
											GenericMemoryPoolAllocator<U8> alloc);

	/// The hash generated by generateShaderGroupGroupHash() can be used to retrieve the group position in the
	/// m_program.
	U32 getIndex(U64 groupHash) const
	{
		auto it = m_resourceHashToShaderGroupHandleIndex.find(groupHash);
		ANKI_ASSERT(it != m_resourceHashToShaderGroupHandleIndex.getEnd());
		return *it;
	}
};

/// A system that does some work on shader programs before resources start loading.
class ShaderProgramResourceSystem
{
public:
	ShaderProgramResourceSystem(CString cacheDir, GrManager* gr, ResourceFilesystem* fs,
								const GenericMemoryPoolAllocator<U8>& alloc)
		: m_alloc(alloc)
		, m_gr(gr)
		, m_fs(fs)
	{
		m_cacheDir.create(alloc, cacheDir);
	}

	~ShaderProgramResourceSystem();

	ANKI_USE_RESULT Error init();

	ConstWeakArray<ShaderProgramRaytracingLibrary> getRayTracingLibraries() const
	{
		return m_rtLibraries;
	}

private:
	GenericMemoryPoolAllocator<U8> m_alloc;
	String m_cacheDir;
	GrManager* m_gr;
	ResourceFilesystem* m_fs;
	DynamicArray<ShaderProgramRaytracingLibrary> m_rtLibraries;

	/// Iterate all programs in the filesystem and compile them to AnKi's binary format.
	static Error compileAllShaders(CString cacheDir, GrManager& gr, ResourceFilesystem& fs,
								   GenericMemoryPoolAllocator<U8>& alloc, StringListAuto& rtProgramFilenames);

	static Error createRayTracingPrograms(CString cacheDir, const StringListAuto& rtProgramFilenames, GrManager& gr,
										  GenericMemoryPoolAllocator<U8>& alloc,
										  DynamicArray<ShaderProgramRaytracingLibrary>& libs);
};
/// @}

} // end namespace anki