#include "generate_atlas.hpp"
#include "../texture.hpp"
#include "../shader.hpp"
#include "atlas.hpp"
#include "image.hpp"
#include "uv.hpp"
#include <iostream>
#include <iterator>
#include <stb/stb_rect_pack.h>
#include <GL/glew.h>

void Graphics::Texture::generate_atlas(GeneratedAtlas& ga) {
	int total_area = 0;
	int min_size = 1;
	int padding = 0;
	int offset = 0;

	for(const Image* i : IMAGES) {
		int w = i->m_size.x + padding;
		int h = i->m_size.y + padding;
		if(w > min_size) min_size = w;
		if(h > min_size) min_size = h;
		total_area += w * h;
	}

	// find an approximate size for the atlas that can hold all the textures
	int size = std::pow(2, std::ceil(std::log2(std::max((double)min_size, std::pow(total_area, 1.0/3.0)))));

	std::vector<stbrp_rect> rects;
	std::vector<UV> uvs(std::size(IMAGES));
	rects.reserve(std::size(IMAGES));

	for(int i = 0; i < std::size(IMAGES); i++) {
		const Image* a = IMAGES[i];
		stbrp_rect rect;
		rect.id = i;
		rect.w = a->m_size.x + padding;
		rect.h = a->m_size.y + padding;
		rects.push_back(rect);
	}

	stbrp_context context;
	std::vector<stbrp_node> nodes(size);
	Atlas atlas({size, size});
	int zpos;

	for(zpos = 0; rects.size() > 0; zpos++) {
		stbrp_init_target(&context, size, size, nodes.data(), nodes.size());
		stbrp_pack_rects(&context, rects.data(), rects.size());

		for(auto it = rects.begin(); it != rects.end();) {
			if(!it->was_packed) {
				it++;
				continue;
			}
			
			Image* src = IMAGES[it->id];
			atlas.draw(*src, {it->x + offset, it->y + offset, zpos});
			src->m_id = it->id;
			src->clear_host();
			
			uvs[it->id] = {
				.uv0={float(it->x + offset) / size, float(it->y + offset) / size},
				.uv1={float(it->x + offset + src->m_size.x) / size, float(it->y + offset + src->m_size.y) / size},
				.zpos=zpos,
				.edges=src->m_edges,
			};

			it = rects.erase(it);
		}
	}

	std::cout << "Finished stitching " << size << "x" << size << "x" << atlas.m_size.z << " atlas" << std::endl;

	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, Graphics::Shader::SSBO_ATLAS_BUFFER, ga.m_ssbo_id);
	glNamedBufferData(ga.m_ssbo_id, sizeof(UV) * uvs.size(), uvs.data(), GL_STATIC_DRAW);

	glActiveTexture(GL_TEXTURE0 + Graphics::Shader::TEX_ATLAS);
	glBindTexture(GL_TEXTURE_2D_ARRAY, ga.m_tex_id);
	glTexImage3D(GL_TEXTURE_2D_ARRAY, 0, GL_RGBA8, atlas.m_size.x, atlas.m_size.y, atlas.m_size.z, 0, GL_RGBA, GL_UNSIGNED_BYTE, atlas.m_data.data());
	glTexParameteri(GL_TEXTURE_2D_ARRAY, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
	glTexParameteri(GL_TEXTURE_2D_ARRAY, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
	glTexParameteri(GL_TEXTURE_2D_ARRAY, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
	glTexParameteri(GL_TEXTURE_2D_ARRAY, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
	glGenerateMipmap(GL_TEXTURE_2D_ARRAY);
	glActiveTexture(GL_TEXTURE0);
}