#include "Texture.hpp"

#include "Macros.hpp"
#include "PodVector.hpp"
#include "ScopeGuard.hpp"

#include <stb_image.h>
#include <stb_rect_pack.h>
#include <bit>
#include <cstring>
#include <utility>

Texture::~Texture() {
	glDeleteTextures(1, &mHandle);
}

namespace ProjectBrussel_UNITY_ID {
GLint MapTextureFilteringToGL(Tags::TexFilter option) {
	using namespace Tags;
	switch (option) {
		case TF_Linear: return GL_LINEAR;
		case TF_Nearest: return GL_NEAREST;
	}
	return 0;
}
} // namespace ProjectBrussel_UNITY_ID

Texture::ErrorCode Texture::InitFromFile(const char* filePath, const TextureProperties& props) {
	using namespace ProjectBrussel_UNITY_ID;

	if (IsValid()) {
		return EC_AlreadyInitialized;
	}

	int width, height;
	int channels;

	auto result = (uint8_t*)stbi_load(filePath, &width, &height, &channels, 4);
	if (!result) {
		return EC_FileIoFailed;
	}
	DEFER { stbi_image_free(result); };

	glGenTextures(1, &mHandle);
	glBindTexture(GL_TEXTURE_2D, mHandle);

	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, MapTextureFilteringToGL(props.minifyingFilter));
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, MapTextureFilteringToGL(props.magnifyingFilter));
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, width, height, 0, GL_RGBA, GL_UNSIGNED_BYTE, result);

	mProps.size = { width, height };

	return EC_Success;
}

Texture::ErrorCode Texture::InitFromImage(const Image& image, const TextureProperties& props) {
	using namespace ProjectBrussel_UNITY_ID;

	if (IsValid()) {
		return EC_AlreadyInitialized;
	}

	GLint sourceFormat;
	switch (image.GetChannels()) {
		case 1: sourceFormat = GL_RED; break;
		case 2: sourceFormat = GL_RG; break;
		case 3: sourceFormat = GL_RGB; break;
		case 4: sourceFormat = GL_RGBA; break;
		default: return EC_InvalidImage;
	}

	auto size = image.GetSize();
	uint8_t* dataPtr = image.GetDataPtr();

	glGenTextures(1, &mHandle);
	glBindTexture(GL_TEXTURE_2D, mHandle);

	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, MapTextureFilteringToGL(props.minifyingFilter));
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, MapTextureFilteringToGL(props.magnifyingFilter));
	glTexImage2D(GL_TEXTURE_2D, 0, sourceFormat, size.x, size.y, 0, sourceFormat, GL_UNSIGNED_BYTE, dataPtr);

	// Copy all [In] properties
	mProps = props;
	mProps.size = size;

	return EC_Success;
}

Texture::ErrorCode Texture::InitAtlas(const AtlasInput& in, AtlasOutput* out) {
	// Force RGBA for easier time uploading to GL texture
	constexpr int kDesiredChannels = 4;

	PodVector<stbrp_rect> rects;
	rects.resize(in.sources.size());

	for (size_t i = 0; i < in.sources.size(); ++i) {
		auto size = in.sources[i].image.GetSize();
		auto& rect = rects[i];
		rect.w = static_cast<stbrp_coord>(size.x);
		rect.h = static_cast<stbrp_coord>(size.y);
	}

	int atlasWidth;
	int atlasHeight;

	// 1. Pack the candidate rectanges onto the (not yet allocated) atlas
	// Note that the coordinates here are top-left origin
	switch (in.packingMode) {
		case PM_KeepSquare: {
			atlasWidth = 512;
			atlasHeight = 512;

			PodVector<stbrp_node> nodes;
			while (true) {
				// No need to zero initialize stbrp_node, library will take care of that
				nodes.resize(atlasWidth);

				stbrp_context ctx;
				stbrp_init_target(&ctx, atlasWidth, atlasHeight, &nodes[0], (int)nodes.size());
				int result = stbrp_pack_rects(&ctx, rects.data(), (int)rects.size());

				if (result != 1) {
					atlasWidth *= 2;
					atlasHeight *= 2;
				} else {
					// Break out of the while loop
					break;
				}
			}
		} break;

		case PM_VerticalExtension:
		case PM_HorizontalExtension: {
			constexpr int kMaxHeight = 1024 * 32;
			atlasWidth = 0;
			atlasHeight = 0;

			PodVector<stbrp_node> nodes;
			stbrp_context ctx;
			stbrp_init_target(&ctx, atlasWidth, atlasHeight, &nodes[0], nodes.size());
			stbrp_pack_rects(&ctx, rects.data(), rects.size());

			// Calculate width/height needed for atlas
			auto& limiter = in.packingMode == PM_VerticalExtension ? atlasHeight : atlasWidth;
			for (auto& rect : rects) {
				int bottom = rect.y + rect.h;
				limiter = std::max(limiter, bottom);
			}
			limiter = std::bit_ceil<uint32_t>(limiter);
		} break;
	}

	// 2. Allocate atlas bitmap

	// Number of bytes in *bitmap*
	auto bytes = atlasWidth * atlasHeight * kDesiredChannels * sizeof(uint8_t);
	// Note that the origin (first pixel) is the bottom-left corner, to be consistent with OpenGL
	auto bitmap = std::make_unique<uint8_t[]>(bytes);
	std::memset(bitmap.get(), 0, bytes * sizeof(uint8_t));

	// 3. Put all candidate images to the atlas bitmap
	// TODO don't flip
	// We essentially flip the candidate images vertically when putting into the atlas bitmap, so that when OpenGL reads
	// these bytes, it sees the "bottom row" (if talking in top-left origin) first
	// (empty spots are set with 0, "flipping" doesn't apply to them)
	//
	// Conceptually, we flip the atlas bitmap vertically so that the origin is at bottom-left
	// i.e. all the coordinates we talk (e.g. rect.x/y) are still in top-left origin

	// Unit: bytes
	size_t bitmapRowStride = atlasWidth * kDesiredChannels * sizeof(uint8_t);
	for (size_t i = 0; i < in.sources.size(); ++i) {
		auto& rect = rects[i];
		// Data is assumed to be stored in top-left origin
		auto data = in.sources[i].image.GetDataPtr();

		// We need to copy row by row, because the candidate image bytes won't land in a continuous chunk in our atlas bitmap
		// Unit: bytes
		size_t incomingRowStride = rect.w * kDesiredChannels * sizeof(uint8_t);
		// Unit: bytes
		size_t bitmapX = rect.x * kDesiredChannels * sizeof(uint8_t);
		for (int y = 0; y < rect.h; ++y) {
			auto src = data + y * incomingRowStride;

			int bitmapY = y;
			auto dst = bitmap.get() + bitmapY * bitmapRowStride + bitmapX;

			std::memcpy(dst, src, incomingRowStride);
		}
	}

	// 4. Upload to VRAM
	GLuint atlasTexture;
	glGenTextures(1, &atlasTexture);
	glBindTexture(GL_TEXTURE_2D, atlasTexture);

	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, atlasWidth, atlasHeight, 0, GL_RGBA, GL_UNSIGNED_BYTE, bitmap.get());

	// 5. Generate atlas texture info
	mHandle = atlasTexture;
	mProps.size = { atlasWidth, atlasHeight };

	// 6. Generate output information
	if (out) {
		out->elements.reserve(in.sources.size());
		for (size_t i = 0; i < in.sources.size(); ++i) {
			auto& rect = rects[i];
			auto& source = in.sources[i];
			out->elements.push_back(AltasElement{
				.name = source.name,
				.subregion = Subregion{
					.u0 = (float)(rect.x) / atlasWidth,
					.v0 = (float)(rect.y + rect.h) / atlasHeight,
					.u1 = (float)(rect.x + rect.w) / atlasWidth,
					.v1 = (float)(rect.y) / atlasHeight,
				},
				.subregionSize = glm::ivec2(rect.w, rect.h),
			});
		}
	}

	return EC_Success;
}

const TextureProperties& Texture::GetInfo() const {
	return mProps;
}

GLuint Texture::GetHandle() const {
	return mHandle;
}

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

Texture* IresTexture::CreateInstance() const {
	return new Texture();
}

Texture* IresTexture::GetInstance() {
	if (mInstance == nullptr) {
		mInstance.Attach(CreateInstance());
	}
	return mInstance.Get();
}

void IresTexture::Write(IresWritingContext& ctx, rapidjson::Value& value, rapidjson::Document& root) const {
	IresObject::Write(ctx, value, root);
	// TODO
}

void IresTexture::Read(IresLoadingContext& ctx, const rapidjson::Value& value) {
	IresObject::Read(ctx, value);
	// TODO
}