#include "IUI/Icons.h"

#include "Kit/EngineManager.h"
#include "IUI/ImGui/imgui.h"
#include "IUI/ImGui/imgui_internal.h"
#include "fs/ReadStream.h"

#include "Kit/Systems/RenderSystem/TextureCPU.h"
#include "Kit/util/Atlaser.h"

#include "Kit/util/stb/stb_image.h"

#include <filament/Engine.h>
#include <filament/Texture.h>
#include "tl/chrono.h"
#include "Logger.h"
#include "Kit/util/Image.h"
#include "tl/algorithm/string.h"

#include "fs/NativeFilesystem.h"
#include "tl/narrow_cast.h"


namespace kit
{
namespace icons
{
constexpr uint32_t k_atlasSize = 1024;

tl::lent_ptr<kit::EngineManager> s_engineManager;
tl::lent_ptr<fs::IFilesystem> s_filesystem;
fs::AbsPath s_folderPath;

using Image = kit::Image<kit::pf::R8G8B8A8>;

struct TextureData
{
	tl::unique_ref<filament::Texture> texture;
	Image image;
	kit::Atlaser atlaser;
};

tl::vector<TextureData> s_textures;
tl::unordered_map<tl::string, Icon> s_icons;

void init(tl::lent_ref<kit::EngineManager> engineManager, tl::lent_ref<fs::IFilesystem> filesystem, fs::AbsPathView folderPath)
{
	s_engineManager = std::move(engineManager);
	s_filesystem = std::move(filesystem);
	s_folderPath = std::move(folderPath);
}

void shutdown()
{
	s_icons.clear();
	s_textures.clear();
	s_engineManager = nullptr;
	s_filesystem = nullptr;
}

Icon loadIcon(const tl::string& name)
{
	TL_ASSERT(s_engineManager);
	fs::AbsPath filePath = s_folderPath + (name + ".png");

	tl::memory_buffer rawData;
	rawData.reserve(32768);
	auto readResult = fs::readFile(rawData, *s_filesystem, filePath);
	if (readResult.has_error())
	{
		LOGE("{}", readResult.error());
		return {};
	}

	int w = 0, h = 0, comp = 0;

	//expand anything to RGBA
	constexpr int reqComp = 4;
	tl::shared_ptr<uint8_t> imageData(kit_stbi_load_from_memory(rawData.data(), (int)rawData.size(), &w, &h, &comp, reqComp), tl::aligned_free16);
	if (!imageData)
	{
		LOGE("Unknown image format for icon '{}'", filePath);
		return {};
	}
	if ((w < 1) || (h < 1))
	{
		LOGE("Invalid size for icon '{}': {}x{}", filePath, w, h);
		return {};
	}
	if (comp != 4)
	{
		LOGE("Invalid number of components for icon '{}': {}", filePath, comp);
		return {};
	}

	tl::lent_ref<filament::Engine> engine = s_engineManager->getEngine();

	const math::vec2u size(tl::narrow<uint32_t>(w), tl::narrow<uint32_t>(h));
	math::vec2u offset;

	TextureData* textureData = nullptr;

	for (TextureData& td : s_textures)
	{
		kit::Atlaser& atlaser = td.atlaser;
		const tl::optional<math::vec2u> region = atlaser.reserveRegion(size);
		if (!region)
			continue;

		textureData = &td;
		offset = region.value();
		break;
	}

	if (!textureData)
	{
		constexpr bool isSRGB = false; //we're using icons for imgui, which doesn't use a linear FB
		constexpr auto format = isSRGB ? filament::Texture::InternalFormat::SRGB8_A8 : filament::Texture::InternalFormat::RGBA8;
		constexpr filament::Texture::Sampler sampler = filament::Texture::Sampler::SAMPLER_2D;

		const uint8_t targetMipmapCount = 1 + (uint8_t)std::log2(k_atlasSize);

		const auto texture = filament::Texture::Builder()
		                     .width(k_atlasSize)
		                     .height(k_atlasSize)
		                     .depth(1)
		                     .levels(targetMipmapCount)
		                     .sampler(sampler)
		                     .format(format)
		                     .build(*engine);

		TextureData td
		{
			.texture = s_engineManager->adopt(texture),
			.image = Image(math::vec2u(k_atlasSize)),
			.atlaser = kit::Atlaser(math::vec2u(k_atlasSize), 4)
		};
		const tl::optional<math::vec2u> region = td.atlaser.reserveRegion(size);
		TL_ASSERT(region);
		offset = region.value();

		s_textures.push_back(std::move(td));

		textureData = &s_textures.back();
	}

	constexpr auto pixelDataFormat = filament::Texture::PixelBufferDescriptor::PixelDataFormat::RGBA;
	constexpr auto pixelDataType = filament::Texture::PixelBufferDescriptor::PixelDataType::UBYTE;

	{
		Image img;
		img.allocateNonOwning(math::vec2u(w, h), (kit::pf::R8G8B8A8*)imageData.get());
		textureData->image.copyFrom(img, {0, 0}, img.getSize(), offset);
	}

	//for debugging purposes
	//textureData->image.saveAsTga(*fs::native.openSink(fs::AbsPath::from_string("icons.tga", fs::native.getCurrentFolder()), fs::Mode::CreateOrOpenAndClear).value());

	textureData->texture->setImage(*engine,
	                               0,
	                               filament::Texture::PixelBufferDescriptor(textureData->image.getData(), k_atlasSize * k_atlasSize * 4, pixelDataFormat, pixelDataType));

	textureData->texture->generateMipmaps(*engine);

	Icon icon
	{
		.texture = textureData->texture,
		.uv0 = math::to_vec2<float>(offset) / float(k_atlasSize),
		.uv1 = math::to_vec2<float>(offset + size) / float(k_atlasSize)
	};

	return icon;
}

const Icon& getIcon(const tl::string& name)
{
	auto it = s_icons.find(name);
	if (it == s_icons.end())
	{
		Icon data = loadIcon(name);
		it = s_icons.emplace(name, std::move(data)).first;
	}
	return it->second;
}

static const Icon k_empty;

struct c_str_hasher
{
	using hash_code_t = tl::unordered_map<tl::string, Icon>::hash_code_t;

	hash_code_t operator()(const char* c_str) const noexcept
	{
		return tl::string::raw_hash(c_str);
	}
};

const Icon& getIcon(const char* name)
{
	if (!name)
		return k_empty;

	auto it = s_icons.find_as(name, c_str_hasher(), tl::equal_to_2<tl::string, const char*>());
	if (it == s_icons.end())
	{
		Icon data = loadIcon(name);
		it = s_icons.emplace(name, std::move(data)).first;
	}
	return it->second;
}
}
}
