#include "ShaderCompiler.hpp"

#include "Base/Logger.hpp"

#include <map>
#include <filesystem>

namespace OpenVulkano
{
#if defined(HAS_SHADERC)
	class ShaderIncluder : public shaderc::CompileOptions::IncluderInterface
	{
		struct IncludeData
		{
			shaderc_include_result result = {};
			std::string m_fullPath;
			Array<char> m_content;
		};

	public:
		ShaderIncluder(const std::string& path);
		~ShaderIncluder() override = default;

		shaderc_include_result* GetInclude(const char* requestedSource, shaderc_include_type type,
		                                   const char* requestingSource, uint64_t includeDepth) override;
		void ReleaseInclude(shaderc_include_result* data) override;

	private:
		std::string ResolveInclude(const std::string& requestedSource) const;

	private:
		std::string m_includePath;
	};

	shaderc_shader_kind CheckStage(const std::string& extensionName)
	{
		std::map<std::string, shaderc_shader_kind> stageMap = {
			{ "vert", shaderc_glsl_vertex_shader },       { "frag", shaderc_glsl_fragment_shader },
			{ "comp", shaderc_glsl_compute_shader },      { "geom", shaderc_glsl_geometry_shader },
			{ "tesc", shaderc_glsl_tess_control_shader }, { "tese", shaderc_glsl_tess_evaluation_shader },
			{ "rgen", shaderc_glsl_raygen_shader },       { "rint", shaderc_glsl_intersection_shader },
			{ "rahit", shaderc_glsl_anyhit_shader },      { "rchit", shaderc_glsl_closesthit_shader },
			{ "rmiss", shaderc_glsl_miss_shader },        { "rcall", shaderc_glsl_callable_shader },
			{ "task", shaderc_glsl_task_shader },         { "mesh", shaderc_glsl_mesh_shader }
		};

		auto it = stageMap.find(extensionName);
		if (it != stageMap.end()) return it->second;
		else throw std::runtime_error("Failed to find shader stage for extension '" + extensionName + "'!");
	}

	Array<uint32_t> ShaderCompiler::CompileGLSLToSpirv(const std::string& absPath, bool hasIncludes,
	                                                 const std::string& incPath, const std::string& entryPoint)
	{
		Array<char> file = Utils::ReadFile(absPath, false, true);

		shaderc::Compiler shaderCompiler;
		shaderc::CompileOptions options;

		if (hasIncludes && incPath.size() > 0) { options.SetIncluder(std::make_unique<ShaderIncluder>(incPath)); }
		options.SetSourceLanguage(shaderc_source_language_glsl);
		options.SetSuppressWarnings();

		auto extensionSplit = Utils::SplitAtLastOccurrence(absPath, '.');

		shaderc::PreprocessedSourceCompilationResult preResult =
		shaderCompiler.PreprocessGlsl(file.Data(), CheckStage(extensionSplit.second), entryPoint.c_str(), options);

		#if defined(_DEBUG)
		std::string test = static_cast<const char*>(preResult.begin(), preResult.end());
		// printf("Preprocessed shader: %s\n\n\n", test.c_str()); // Works fine
		// Logger::APP->info("Preprocessed shader: {0}", test); gives error
		#endif

		if (preResult.GetCompilationStatus() != shaderc_compilation_status_success)
		{
			throw std::runtime_error("Failed preprocessing shader. Reason: " + preResult.GetErrorMessage());
		}

		shaderc::CompilationResult<uint32_t> result = shaderCompiler.CompileGlslToSpv(
		static_cast<const char*>(preResult.begin()), CheckStage(extensionSplit.second), "", options);

		if (result.GetCompilationStatus() != shaderc_compilation_status_success)
		{
			throw std::runtime_error("Failed compiling shader. Reason: " + result.GetErrorMessage());
		}

		Array<uint32_t> returnResult(result.cend() - result.cbegin());
		// Copy the data to the return array
		std::copy(result.begin(), result.end(), returnResult.Data());
		return returnResult;
	}

	ShaderIncluder::ShaderIncluder(const std::string& path) : m_includePath(path) {}

	shaderc_include_result* ShaderIncluder::GetInclude(const char* requestedSource, shaderc_include_type type,
	                                                   const char* requestingSource, uint64_t includeDepth)
	{
		IncludeData* includeData = new IncludeData();
		includeData->m_fullPath = ResolveInclude(requestedSource);

		// Do not null terminate the string
		includeData->m_content = Utils::ReadFile(includeData->m_fullPath, false, false);

		shaderc_include_result* result = &includeData->result;
		result->content = includeData->m_content.Data();
		result->content_length = includeData->m_content.Size();
		result->source_name = includeData->m_fullPath.data();
		result->source_name_length = includeData->m_fullPath.size();
		result->user_data = includeData;

		return result;
	}

	void ShaderIncluder::ReleaseInclude(shaderc_include_result* data)
	{
		delete static_cast<IncludeData*>(data->user_data);
	}

	std::string ShaderIncluder::ResolveInclude(const std::string& requestedSource) const
	{
		// Check if the file exists in the include path
		std::string path = m_includePath + requestedSource;
		if (std::filesystem::exists(path)) return path;
		else throw std::runtime_error("Failed to resolve include '" + requestedSource + "'!");
	}
#endif
}