#include "Material.hpp"

#include "AppConfig.hpp"
#include "EditorCore.hpp"
#include "EditorUtils.hpp"
#include "RapidJsonHelper.hpp"
#include "ScopeGuard.hpp"
#include "Utils.hpp"

#include <imgui.h>
#include <rapidjson/document.h>
#include <cstdlib>
#include <cstring>
#include <utility>

using namespace std::literals;

Material::Material() {
}

namespace ProjectBrussel_UNITY_ID {
bool TryFindShaderId(Shader* shader, std::string_view name, int& out) {
	auto& info = shader->GetInfo();
	auto iter = info.things.find(name);
	if (iter == info.things.end()) return false;
	auto& id = iter->second;

	if (id.kind != ShaderThingId::KD_Uniform) return false;

	out = id.index;
	return true;
}

template <class TUniform>
TUniform& ObtainUniform(Shader* shader, const char* name, std::vector<TUniform>& uniforms, GLint location) {
	for (auto& uniform : uniforms) {
		if (uniform.location == location) {
			return uniform;
		}
	}

	auto& uniform = uniforms.emplace_back();
	uniform.location = location;
	if (!TryFindShaderId(shader, name, uniform.infoUniformIndex)) {
		uniform.infoUniformIndex = -1;
	}

	return uniform;
}

rapidjson::Value MakeVectorJson(const Material::VectorUniform& vector, rapidjson::Document& root) {
	int len = vector.actualLength;

	rapidjson::Value result(rapidjson::kArrayType);
	result.Reserve(len, root.GetAllocator());

	for (int i = 0; i < len; ++i) {
		result.PushBack(vector.value[i], root.GetAllocator());
	}

	return result;
}

Material::VectorUniform ReadVectorFromJson(const rapidjson::Value& rv) {
	assert(rv.IsArray());
	Material::VectorUniform result;
	int len = result.actualLength = rv.Size();
	for (int i = 0; i < len; ++i) {
		result.value[i] = rv[i].GetFloat();
	}
	return result;
}

rapidjson::Value MakeMatrixJson(const Material::MatrixUniform& matrix, rapidjson::Document& root) {
	int w = matrix.actualWidth;
	int h = matrix.actualHeight;

	rapidjson::Value result(rapidjson::kArrayType);
	result.Reserve(h, root.GetAllocator());

	for (int y = 0; y < h; ++y) {
		rapidjson::Value row(rapidjson::kArrayType);
		row.Reserve(w, root.GetAllocator());

		for (int x = 0; x < w; ++x) {
			// Each item in a column is consecutive in memory in glm::mat<> structs
			row.PushBack(matrix.value[x * h + y], root.GetAllocator());
		}

		result.PushBack(row, root.GetAllocator());
	}

	return result;
}

Material::MatrixUniform ReadMatrixFromjson(const rapidjson::Value& rv) {
	assert(rv.IsArray());
	assert(rv.Size() > 0);
	assert(rv[0].IsArray());
	Material::MatrixUniform result;
	int w = result.actualWidth = rv[0].Size();
	int h = result.actualHeight = rv.Size();
	for (int y = 0; y < h; ++y) {
		auto& row = rv[y];
		assert(row.IsArray());
		assert(row.Size() == w);
		for (int x = 0; x < w; ++x) {
			auto& val = row[x];
			assert(val.IsNumber());
			result.value[x * h + y] = val.GetFloat();
		}
	}
	return result;
}
} // namespace ProjectBrussel_UNITY_ID

void Material::SetFloat(const char* name, float value) {
	assert(IsValid());

	GLint location = glGetUniformLocation(mShader->GetProgram(), name);
	auto& uniform = ProjectBrussel_UNITY_ID::ObtainUniform(mShader.Get(), name, mBoundScalars, location);
	uniform.floatValue = value;
	uniform.actualType = GL_FLOAT;
}

void Material::SetInt(const char* name, int32_t value) {
	assert(IsValid());

	GLint location = glGetUniformLocation(mShader->GetProgram(), name);
	auto& uniform = ProjectBrussel_UNITY_ID::ObtainUniform(mShader.Get(), name, mBoundScalars, location);
	uniform.intValue = value;
	uniform.actualType = GL_INT;
}

void Material::SetUInt(const char* name, uint32_t value) {
	assert(IsValid());

	GLint location = glGetUniformLocation(mShader->GetProgram(), name);
	auto& uniform = ProjectBrussel_UNITY_ID::ObtainUniform(mShader.Get(), name, mBoundScalars, location);
	uniform.uintValue = value;
	uniform.actualType = GL_UNSIGNED_INT;
}

template <int length>
void Material::SetVector(const char* name, const glm::vec<length, float>& vec) {
	assert(IsValid());

	static_assert(length >= 1 && length <= 4);

	GLint location = glGetUniformLocation(mShader->GetProgram(), name);
	auto& uniform = ProjectBrussel_UNITY_ID::ObtainUniform(mShader.Get(), name, mBoundVectors, location);
	uniform.actualLength = length;
	std::memset(uniform.value, 0, sizeof(uniform.value));
	std::memcpy(uniform.value, &vec[0], length * sizeof(float));
}

template void Material::SetVector<1>(const char*, const glm::vec<1, float>&);
template void Material::SetVector<2>(const char*, const glm::vec<2, float>&);
template void Material::SetVector<3>(const char*, const glm::vec<3, float>&);
template void Material::SetVector<4>(const char*, const glm::vec<4, float>&);

template <int width, int height>
void Material::SetMatrix(const char* name, const glm::mat<width, height, float>& mat) {
	static_assert(width >= 1 && width <= 4);
	static_assert(height >= 1 && height <= 4);

	GLint location = glGetUniformLocation(mShader->GetProgram(), name);
	auto& uniform = ProjectBrussel_UNITY_ID::ObtainUniform(mShader.Get(), name, mBoundMatrices, location);
	uniform.actualWidth = width;
	uniform.actualHeight = height;
	std::memset(uniform.value, 0, sizeof(uniform.value));
	std::memcpy(uniform.value, &mat[0][0], width * height * sizeof(float));
}

template void Material::SetMatrix<2, 2>(const char*, const glm::mat<2, 2, float>&);
template void Material::SetMatrix<3, 3>(const char*, const glm::mat<3, 3, float>&);
template void Material::SetMatrix<4, 4>(const char*, const glm::mat<4, 4, float>&);

template void Material::SetMatrix<2, 3>(const char*, const glm::mat<2, 3, float>&);
template void Material::SetMatrix<3, 2>(const char*, const glm::mat<3, 2, float>&);

template void Material::SetMatrix<2, 4>(const char*, const glm::mat<2, 4, float>&);
template void Material::SetMatrix<4, 2>(const char*, const glm::mat<4, 2, float>&);

template void Material::SetMatrix<3, 4>(const char*, const glm::mat<3, 4, float>&);
template void Material::SetMatrix<4, 3>(const char*, const glm::mat<4, 3, float>&);

void Material::SetTexture(const char* name, Texture* texture) {
	assert(IsValid());

	GLint location = glGetUniformLocation(mShader->GetProgram(), name);

	for (auto& uniform : mBoundTextures) {
		if (uniform.location == location) {
			uniform.value.Attach(texture);
			return;
		}
	}

	auto& uniform = mBoundTextures.emplace_back();
	uniform.value.Attach(texture);
	uniform.location = location;
}

std::span<const Material::VectorUniform> Material::GetVectors() const {
	return mBoundVectors;
}

std::span<const Material::MatrixUniform> Material::GetMatrices() const {
	return mBoundMatrices;
}

std::span<const Material::TextureUniform> Material::GetTextures() const {
	return mBoundTextures;
}

Shader* Material::GetShader() const {
	return mShader.Get();
}

void Material::SetShader(Shader* shader) {
	mShader.Attach(shader);
	auto& info = shader->GetInfo();

	mBoundScalars.clear();
	mBoundVectors.clear();
	mBoundMatrices.clear();
	mBoundTextures.clear();
	for (int i = 0; i < info.uniforms.size(); ++i) {
		auto& decl = info.uniforms[i];
		switch (decl->kind) {
			case ShaderVariable::KD_Math: {
				auto& mathDecl = static_cast<ShaderMathVariable&>(*decl);
				if (mathDecl.width == 1) {
					if (mathDecl.height == 1) {
						// Scalar
						auto& scalar = mBoundScalars.emplace_back();
						scalar.location = decl->location;
						scalar.infoUniformIndex = i;
					} else {
						// Vector
						auto& vec = mBoundVectors.emplace_back();
						vec.location = decl->location;
						vec.infoUniformIndex = i;
						vec.actualLength = mathDecl.height;
					}
				} else {
					// Matrix
					auto& mat = mBoundMatrices.emplace_back();
					mat.location = decl->location;
					mat.infoUniformIndex = i;
					mat.actualWidth = mathDecl.width;
					mat.actualHeight = mathDecl.height;
				}
			} break;

			case ShaderVariable::KD_Sampler: {
				auto& uniform = mBoundTextures.emplace_back();
				uniform.location = decl->location;
				uniform.infoUniformIndex = i;
			} break;
		}
	}
}

bool Material::IsValid() const {
	return mShader != nullptr;
}

static constexpr int IdentifyMatrixSize(int width, int height) {
	return width * 10 + height;
}

void Material::UseUniforms() const {
	for (auto& uniform : mBoundScalars) {
		switch (uniform.actualType) {
			case GL_FLOAT: glUniform1f(uniform.location, uniform.intValue); break;
			case GL_INT: glUniform1i(uniform.location, uniform.intValue); break;
			case GL_UNSIGNED_INT: glUniform1ui(uniform.location, uniform.intValue); break;
			default: break;
		}
	}

	for (auto& uniform : mBoundVectors) {
		switch (uniform.actualLength) {
			case 1: glUniform1fv(uniform.location, 1, &uniform.value[0]); break;
			case 2: glUniform2fv(uniform.location, 1, &uniform.value[0]); break;
			case 3: glUniform3fv(uniform.location, 1, &uniform.value[0]); break;
			case 4: glUniform4fv(uniform.location, 1, &uniform.value[0]); break;
			default: break;
		}
	}

	for (auto& uniform : mBoundMatrices) {
		switch (IdentifyMatrixSize(uniform.actualWidth, uniform.actualHeight)) {
			case IdentifyMatrixSize(2, 2): glUniformMatrix2fv(uniform.location, 1, GL_FALSE, uniform.value); break;
			case IdentifyMatrixSize(3, 3): glUniformMatrix3fv(uniform.location, 1, GL_FALSE, uniform.value); break;
			case IdentifyMatrixSize(4, 4): glUniformMatrix4fv(uniform.location, 1, GL_FALSE, uniform.value); break;

			case IdentifyMatrixSize(2, 3): glUniformMatrix2x3fv(uniform.location, 1, GL_FALSE, uniform.value); break;
			case IdentifyMatrixSize(3, 2): glUniformMatrix3x2fv(uniform.location, 1, GL_FALSE, uniform.value); break;

			case IdentifyMatrixSize(2, 4): glUniformMatrix2x4fv(uniform.location, 1, GL_FALSE, uniform.value); break;
			case IdentifyMatrixSize(4, 2): glUniformMatrix4x2fv(uniform.location, 1, GL_FALSE, uniform.value); break;

			case IdentifyMatrixSize(3, 4): glUniformMatrix3x4fv(uniform.location, 1, GL_FALSE, uniform.value); break;
			case IdentifyMatrixSize(4, 3): glUniformMatrix4x3fv(uniform.location, 1, GL_FALSE, uniform.value); break;

			default: break;
		}
	}

	int i = 0;
	for (auto& uniform : mBoundTextures) {
		glActiveTexture(GL_TEXTURE0 + i);
		glBindTexture(GL_TEXTURE_2D, uniform.value->GetHandle());
		glUniform1i(uniform.location, i);
		++i;
	}
}

bool Material::Write(rapidjson::Value& value, rapidjson::Document& root) const {
	using namespace ProjectBrussel_UNITY_ID;

	if (!IsValid()) return false;

	auto& shaderInfo = mShader->GetInfo();

	value.AddMember("ShaderName", mShader->GetName(), root.GetAllocator());

	rapidjson::Value fields(rapidjson::kArrayType);
	for (auto& scalar : mBoundScalars) {
		rapidjson::Value rvField(rapidjson::kObjectType);
		rvField.AddMember("Name", shaderInfo.uniforms[scalar.infoUniformIndex]->name, root.GetAllocator());
		rvField.AddMember("Type", "Scalar", root.GetAllocator());
		switch (scalar.actualType) {
			case GL_FLOAT: rvField.AddMember("Value", scalar.floatValue, root.GetAllocator()); break;
			case GL_INT: rvField.AddMember("Value", scalar.intValue, root.GetAllocator()); break;
			case GL_UNSIGNED_INT: rvField.AddMember("Value", scalar.uintValue, root.GetAllocator()); break;
		}
		fields.PushBack(rvField, root.GetAllocator());
	}
	for (auto& vector : mBoundVectors) {
		rapidjson::Value rvField(rapidjson::kObjectType);
		rvField.AddMember("Name", shaderInfo.uniforms[vector.infoUniformIndex]->name, root.GetAllocator());
		rvField.AddMember("Type", "Vector", root.GetAllocator());
		rvField.AddMember("Value", MakeVectorJson(vector, root).Move(), root.GetAllocator());
		fields.PushBack(rvField, root.GetAllocator());
	}
	for (auto& matrix : mBoundMatrices) {
		rapidjson::Value rvField(rapidjson::kObjectType);
		rvField.AddMember("Name", shaderInfo.uniforms[matrix.infoUniformIndex]->name, root.GetAllocator());
		rvField.AddMember("Type", "Matrix", root.GetAllocator());
		rvField.AddMember("Value", MakeMatrixJson(matrix, root).Move(), root.GetAllocator());
		fields.PushBack(rvField, root.GetAllocator());
	}
	for (auto& texture : mBoundTextures) {
		// TODO
	}
	value.AddMember("Fields", fields, root.GetAllocator());

	return true;
}

bool Material::Read(const rapidjson::Value& value) {
	using namespace ProjectBrussel_UNITY_ID;

	{
		auto rvShaderName = rapidjson::GetProperty(value, rapidjson::kStringType, "ShaderName"sv);
		if (!rvShaderName) return false;

		auto shader = ShaderManager::instance->FindShader(rapidjson::AsStringView(*rvShaderName));
		if (!shader) return false;

		mShader.Attach(shader);
	}
	auto& shaderInfo = mShader->GetInfo();

	auto fields = rapidjson::GetProperty(value, rapidjson::kArrayType, "Fields"sv);
	if (!fields) return false;

	for (auto& rvField : fields->GetArray()) {
		if (!rvField.IsObject()) continue;

		auto rvName = rapidjson::GetProperty(rvField, rapidjson::kStringType, "Name"sv);

		auto rvType = rapidjson::GetProperty(rvField, rapidjson::kStringType, "Type"sv);
		if (!rvType) continue;
		auto type = rapidjson::AsStringView(*rvType);

		auto rvValue = rapidjson::GetProperty(rvField, "Value"sv);

		if (type == "Scalar"sv) {
			ScalarUniform uniform;
			if (rvName) {
				TryFindShaderId(mShader.Get(), rapidjson::AsStringView(*rvName), uniform.infoUniformIndex);
				uniform.location = shaderInfo.uniforms[uniform.infoUniformIndex]->location;
			}
			if (rvValue->IsFloat()) {
				uniform.actualType = GL_FLOAT;
				uniform.floatValue = rvValue->GetFloat();
			} else if (rvValue->IsInt()) {
				uniform.actualType = GL_INT;
				uniform.intValue = rvValue->GetInt();
			} else if (rvValue->IsUint()) {
				uniform.actualType = GL_UNSIGNED_INT;
				uniform.uintValue = rvValue->GetUint();
			}
			mBoundScalars.push_back(std::move(uniform));
		} else if (type == "Vector"sv) {
			auto uniform = ReadVectorFromJson(*rvValue);
			if (rvName) {
				TryFindShaderId(mShader.Get(), rapidjson::AsStringView(*rvName), uniform.infoUniformIndex);
				uniform.location = shaderInfo.uniforms[uniform.infoUniformIndex]->location;
			}
			mBoundVectors.push_back(std::move(uniform));
		} else if (type == "Matrix"sv) {
			auto uniform = ReadMatrixFromjson(*rvValue);
			if (rvName) {
				TryFindShaderId(mShader.Get(), rapidjson::AsStringView(*rvName), uniform.infoUniformIndex);
				uniform.location = shaderInfo.uniforms[uniform.infoUniformIndex]->location;
			}
			mBoundMatrices.push_back(uniform);
		} else if (type == "Texture"sv) {
			// TODO
		}
	}

	return true;
}

IresMaterial::IresMaterial()
	: IresObject(KD_Material)
	, mInstance(new Material()) {
}

Material* IresMaterial::GetInstance() const {
	return mInstance.Get();
}

void IresMaterial::InvalidateInstance() {
	mInstance.Attach(new Material());
}

void IresMaterial::ShowEditor(EditorInstance& editor) {
	using namespace Tags;

	IresObject::ShowEditor(editor);

	auto shader = mInstance->GetShader();
	if (shader) {
		auto& name = shader->GetName();
		bool isAnnoymous = name.empty();
		if (isAnnoymous) {
			ImGui::Text("Shader <annoymous at %p>", (void*)shader);
		} else {
			ImGui::Text("Shader: %s", name.c_str());
		}
	} else {
		ImGui::TextUnformatted("Shader: <null>");
	}
	if (ImGui::BeginDragDropTarget()) {
		if (auto payload = ImGui::AcceptDragDropPayload(BRUSSEL_TAG_Shader)) {
			auto shader = *static_cast<Shader* const*>(payload->Data);
			mInstance->SetShader(shader);
		}
		ImGui::EndDragDropTarget();
	}
	ImGui::SameLine();
	if (ImGui::Button("GoTo", shader == nullptr)) {
		editor.GetInspector().SelectTarget(EditorInspector::ITT_Shader, shader);
	}

	if (!shader) return;
	auto& info = shader->GetInfo();

	for (auto& field : mInstance->mBoundScalars) {
		auto& decl = static_cast<ShaderMathVariable&>(*info.uniforms[field.infoUniformIndex]);
		decl.ShowInfo();

		ImGui::Indent();
		switch (decl.scalarType) {
			case GL_FLOAT: ImGui::InputFloat("##", &field.floatValue); break;
			case GL_INT: ImGui::InputInt("##", &field.intValue); break;
			// TODO proper uint edit?
			case GL_UNSIGNED_INT: ImGui::InputInt("##", (int32_t*)(&field.uintValue), 0, std::numeric_limits<int32_t>::max()); break;
			default: ImGui::TextUnformatted("Unsupported scalar type"); break;
		}
		ImGui::Unindent();
	}
	for (auto& field : mInstance->mBoundVectors) {
		auto& decl = static_cast<ShaderMathVariable&>(*info.uniforms[field.infoUniformIndex]);
		decl.ShowInfo();

		ImGui::Indent();
		switch (decl.semantic) {
			case VES_Color1:
			case VES_Color2: {
				ImGui::ColorEdit4("##", field.value);
			} break;

			default: {
				ImGui::InputFloat4("##", field.value);
			} break;
		}
		ImGui::Unindent();
	}
	for (auto& field : mInstance->mBoundMatrices) {
		auto& decl = static_cast<ShaderMathVariable&>(*info.uniforms[field.infoUniformIndex]);
		decl.ShowInfo();

		// TODO
	}
	for (auto& field : mInstance->mBoundTextures) {
		auto& decl = static_cast<ShaderSamplerVariable&>(*info.uniforms[field.infoUniformIndex]);
		decl.ShowInfo();

		// TODO
	}
}

void IresMaterial::Write(rapidjson::Value& value, rapidjson::Document& root) const {
	IresObject::Write(value, root);
	mInstance->Write(value, root);
}

void IresMaterial::Read(const rapidjson::Value& value) {
	InvalidateInstance();

	IresObject::Read(value);
	mInstance->Read(value);
}