#include "Shader.hpp"

#include "ScopeGuard.hpp"

#include <cstddef>
#include <fstream>
#include <iostream>
#include <sstream>
#include <utility>

using namespace std::literals::string_literals;
using namespace std::literals::string_view_literals;

Shader::~Shader() {
	glDeleteProgram(mHandle);
}

namespace ProjectBrussel_UNITY_ID {
Shader::ErrorCode LoadFile(std::string& out, const char* filePath) {
	std::ifstream ifs(filePath);
	if (!ifs) {
		return Shader::FileIOFailed;
	}

	std::stringstream buf;
	buf << ifs.rdbuf();
	out = buf.str();

	return Shader::Success;
}

// Grabs section [begin, end)
Shader::ErrorCode CreateShader(GLuint& out, const char* src, int beginIdx, int endIdx, GLenum type) {
	out = glCreateShader(type);

	const GLchar* begin = &src[beginIdx];
	const GLint len = endIdx - beginIdx;
	glShaderSource(out, 1, &begin, &len);

	glCompileShader(out);
	GLint compileStatus;
	glGetShaderiv(out, GL_COMPILE_STATUS, &compileStatus);
	if (compileStatus == GL_FALSE) {
		GLint len;
		glGetShaderiv(out, GL_INFO_LOG_LENGTH, &len);

		std::string log(len, '\0');
		glGetShaderInfoLog(out, len, nullptr, log.data());

		return Shader ::CompilationFailed;
	}

	return Shader::Success;
}

Shader::ErrorCode CreateShader(GLuint& out, std::string_view str, GLenum type) {
	return CreateShader(out, str.data(), 0, str.size(), type);
}

Shader::ErrorCode LinkShaderProgram(GLuint program) {
	glLinkProgram(program);
	GLint linkStatus;
	glGetProgramiv(program, GL_LINK_STATUS, &linkStatus);
	if (linkStatus == GL_FALSE) {
		GLint len;
		glGetProgramiv(program, GL_INFO_LOG_LENGTH, &len);

		std::string log(len, '\0');
		glGetProgramInfoLog(program, len, nullptr, log.data());

		return Shader::LinkingFailed;
	}

	return Shader::Success;
}
} // namespace ProjectBrussel_UNITY_ID

#define CATCH_ERROR_IMPL(x, name)  \
	auto name = x;                 \
	if (name != Shader::Success) { \
		return name;               \
	}
#define CATCH_ERROR(x) CATCH_ERROR_IMPL(x, UNIQUE_NAME(result))

Shader::ErrorCode Shader::InitFromSources(const ShaderSources& sources) {
	using namespace ProjectBrussel_UNITY_ID;

	if (IsValid()) {
		return ShaderAlreadyCreated;
	}

	GLuint program = glCreateProgram();
	ScopeGuard sg = [&]() { glDeleteProgram(program); };

	GLuint vertex = 0;
	DEFER { glDeleteShader(vertex); };
	if (!sources.vertex.empty()) {
		CATCH_ERROR(CreateShader(vertex, sources.vertex, GL_VERTEX_SHADER));
		glAttachShader(program, vertex);
	}

	GLuint geometry = 0;
	DEFER { glDeleteShader(geometry); };
	if (!sources.geometry.empty()) {
		CATCH_ERROR(CreateShader(geometry, sources.geometry, GL_GEOMETRY_SHADER));
		glAttachShader(program, geometry);
	}

	GLuint tessControl = 0;
	DEFER { glDeleteShader(tessControl); };
	if (!sources.tessControl.empty()) {
		CATCH_ERROR(CreateShader(tessControl, sources.tessControl, GL_TESS_CONTROL_SHADER));
		glAttachShader(program, tessControl);
	}

	GLuint tessEval = 0;
	DEFER { glDeleteShader(tessEval); };
	if (!sources.tessEval.empty()) {
		CATCH_ERROR(CreateShader(tessEval, sources.tessEval, GL_TESS_EVALUATION_SHADER));
		glAttachShader(program, tessEval);
	}

	GLuint fragment = 0;
	DEFER { glDeleteShader(fragment); };
	if (!sources.fragment.empty()) {
		CATCH_ERROR(CreateShader(fragment, sources.fragment, GL_FRAGMENT_SHADER));
		glAttachShader(program, fragment);
	}

	CATCH_ERROR(LinkShaderProgram(program));

	sg.Dismiss();
	mHandle = program;

	return Success;
}

Shader::ErrorCode Shader::InitFromFiles(const ShaderFilePaths& files) {
	using namespace ProjectBrussel_UNITY_ID;

	if (IsValid()) {
		return ShaderAlreadyCreated;
	}

	GLuint program = glCreateProgram();
	ScopeGuard sg = [&]() { glDeleteProgram(program); };

	GLuint vertex = 0;
	DEFER { glDeleteShader(vertex); };
	if (files.vertex) {
		std::string src;
		CATCH_ERROR(LoadFile(src, files.vertex));
		CATCH_ERROR(CreateShader(vertex, src, GL_VERTEX_SHADER));
		glAttachShader(program, vertex);
	}

	GLuint geometry = 0;
	DEFER { glDeleteShader(geometry); };
	if (files.geometry) {
		std::string src;
		CATCH_ERROR(LoadFile(src, files.geometry));
		CATCH_ERROR(CreateShader(geometry, src, GL_GEOMETRY_SHADER));
		glAttachShader(program, geometry);
	}

	GLuint tessControl = 0;
	DEFER { glDeleteShader(tessControl); };
	if (files.tessControl) {
		std::string src;
		CATCH_ERROR(LoadFile(src, files.tessControl));
		CATCH_ERROR(CreateShader(tessControl, src, GL_TESS_CONTROL_SHADER));
		glAttachShader(program, tessControl);
	}

	GLuint tessEval = 0;
	DEFER { glDeleteShader(tessEval); };
	if (files.tessEval) {
		std::string src;
		CATCH_ERROR(LoadFile(src, files.tessEval));
		CATCH_ERROR(CreateShader(tessEval, src, GL_TESS_EVALUATION_SHADER));
		glAttachShader(program, tessEval);
	}

	GLuint fragment = 0;
	DEFER { glDeleteShader(fragment); };
	if (files.fragment) {
		std::string src;
		CATCH_ERROR(LoadFile(src, files.fragment));
		CATCH_ERROR(CreateShader(fragment, src, GL_FRAGMENT_SHADER));
		glAttachShader(program, fragment);
	}

	CATCH_ERROR(LinkShaderProgram(program));

	sg.Dismiss();
	mHandle = program;

	return Success;
}

Shader::ErrorCode Shader::InitFromSource(std::string_view source) {
	using namespace ProjectBrussel_UNITY_ID;

	GLuint vertex = 0;
	DEFER { glDeleteShader(vertex); };

	GLuint geometry = 0;
	DEFER { glDeleteShader(geometry); };

	GLuint tessControl = 0;
	DEFER { glDeleteShader(tessControl); };

	GLuint tessEval = 0;
	DEFER { glDeleteShader(tessEval); };

	GLuint fragment = 0;
	DEFER { glDeleteShader(fragment); };

	int prevBegin = -1; // Excluding #type marker
	int prevEnd = -1; // [begin, end)
	std::string prevShaderVariant;

	auto CommitSection = [&]() -> ErrorCode {
		if (prevBegin == -1 || prevEnd == -1) {
			// Not actually "succeeding" here, but we just want to skip this call and continue
			return Success;
		}

		if (prevShaderVariant == "vertex" && !vertex) {
			CATCH_ERROR(CreateShader(vertex, source.data(), prevBegin, prevEnd, GL_VERTEX_SHADER));
		} else if (prevShaderVariant == "geometry" && !geometry) {
			CATCH_ERROR(CreateShader(geometry, source.data(), prevBegin, prevEnd, GL_GEOMETRY_SHADER));
		} else if (prevShaderVariant == "tessellation_control" && !tessControl) {
			CATCH_ERROR(CreateShader(tessControl, source.data(), prevBegin, prevEnd, GL_TESS_CONTROL_SHADER));
		} else if (prevShaderVariant == "tessellation_evaluation" && !tessEval) {
			CATCH_ERROR(CreateShader(tessEval, source.data(), prevBegin, prevEnd, GL_TESS_EVALUATION_SHADER));
		} else if (prevShaderVariant == "fragment" && !fragment) {
			CATCH_ERROR(CreateShader(fragment, source.data(), prevBegin, prevEnd, GL_FRAGMENT_SHADER));
		} else {
			return InvalidShaderVariant;
		}

		prevBegin = -1;
		prevEnd = -1;
		prevShaderVariant.clear();

		return Success;
	};

	constexpr const char* kMarker = "#type ";
	bool matchingDirective = true; // If true, we are matching marker pattern; if false, we are accumulating shader variant identifier
	int matchIndex = 0; // Current index of the pattern trying to match
	std::string shaderVariant;

	// Don't use utf8 iterator, shader sources are expected to be ASCII only
	for (size_t i = 0; i < source.size(); ++i) {
		char c = source[i];

		if (matchingDirective) {
			if (c == kMarker[matchIndex]) {
				// Matched the expected character, go to next char in pattern
				matchIndex++;

				// If we are at the end of the marker pattern...
				if (kMarker[matchIndex] == '\0') {
					matchingDirective = false;
					matchIndex = 0;
					continue;
				}

				// This might be a shader variant directive -> might be end of a section
				if (c == '#') {
					prevEnd = i;
					continue;
				}
			} else {
				// Unexpected character, rollback to beginning
				matchIndex = 0;
			}
		} else {
			if (c == '\n') {
				// Found complete shader variant directive

				CATCH_ERROR(CommitSection()); // Try commit section, for the first apparent of #type this should do nothing, as `prevEnd` will still be -1
				prevBegin = i + 1; // +1 to skip new line (technically not needed)
				prevShaderVariant = std::move(shaderVariant);

				matchingDirective = true;
				shaderVariant.clear();
			} else {
				// Simply accumulate to shader variant buffer
				shaderVariant += c;
			}
		}
	}

	// Commit the last section
	prevEnd = static_cast<int>(source.size());
	CATCH_ERROR(CommitSection());

	GLuint program = glCreateProgram();
	ScopeGuard sg = [&]() { glDeleteProgram(program); };

	if (vertex) glAttachShader(program, vertex);
	if (geometry) glAttachShader(program, geometry);
	if (tessControl) glAttachShader(program, tessControl);
	if (tessEval) glAttachShader(program, tessEval);
	if (fragment) glAttachShader(program, fragment);

	CATCH_ERROR(LinkShaderProgram(program));

	sg.Dismiss();
	mHandle = program;

	return Success;
}

Shader::ErrorCode Shader::InitFromFile(const char* filePath) {
	std::string src;
	CATCH_ERROR(ProjectBrussel_UNITY_ID::LoadFile(src, filePath));

	return InitFromSource(src);
}

#undef CATCH_ERROR

GLuint Shader::GetProgram() const {
	return mHandle;
}

bool Shader::IsValid() const {
	return mHandle != 0;
}