#include "IUI/ImGuiFilamentBridge.h"

#include <utils/EntityManager.h>
#include <tl/vector.h>
#include "IUI/ImGui/imgui.h"
#include <filamat/MaterialBuilder.h>
#include <filament/IndexBuffer.h>
#include <filament/Material.h>
#include <filament/MaterialInstance.h>
#include <filament/RenderableManager.h>
#include <filament/Scene.h>
#include <filament/Texture.h>
#include <filament/TextureSampler.h>
#include <filament/VertexBuffer.h>
#include <Kit/EngineManager.h>

#include "tl/narrow_cast.h"

using namespace filament::math;
using namespace filament;
using namespace utils;

namespace kit
{
static constexpr uint8_t IMGUI_PACKAGE[] =
{
// ReSharper disable All
#include "Materials/ImGuiMaterial.h"
// ReSharper restore All
};

ImGuiFilamentBridge::ImGuiFilamentBridge(tl::lent_ref<EngineManager> engineManager, tl::lent_ref<filament::Scene> flmScene)
	: m_engineManager(std::move(engineManager))
    , m_scene(std::move(flmScene))
{
	// Create a simple alpha-blended 2D blitting material.
	auto builder = Material::Builder()
		.package(IMGUI_PACKAGE, sizeof(IMGUI_PACKAGE));

	m_material = m_engineManager->adopt(builder.build(*m_engineManager->getEngine()));

	EntityManager& em = utils::EntityManager::get();
	m_renderable = em.create();
	m_scene->addEntity(m_renderable);
}

ImTextureID ImGuiFilamentBridge::createAtlasTexture()
{
	const ImGuiIO& io = ImGui::GetIO();
	// Create the grayscale texture that ImGui uses for its glyph atlas.
	unsigned char* pixels;
	int width, height;
	io.Fonts->GetTexDataAsRGBA32(&pixels, &width, &height);
	const size_t size = tl::narrow<size_t>(width * height * 4);
	uint8_t* buffer = m_engineManager->acquireUploadMemory(size);
	memcpy(buffer, pixels, size);

	auto builder = Texture::Builder()
	               .width(tl::narrow<uint32_t>(width))
	               .height(tl::narrow<uint32_t>(height))
	               .levels(tl::narrow<uint8_t>(1))
	               .format(Texture::InternalFormat::RGBA8)
	               .sampler(Texture::Sampler::SAMPLER_2D);

	m_texture = m_engineManager->adopt(builder.build(*m_engineManager->getEngine()));
	m_texture->setImage(*m_engineManager->getEngine(), 0, Texture::PixelBufferDescriptor(buffer, size, Texture::Format::RGBA, Texture::Type::UBYTE));

	const TextureSampler sampler(TextureSampler::MinFilter::LINEAR_MIPMAP_LINEAR, TextureSampler::MagFilter::LINEAR);
	m_material->setDefaultParameter("albedo", m_texture.get(), sampler);
	for (const auto& mi : m_materialInstances)
		mi->setParameter("albedo", m_texture.get(), sampler);

	return m_texture.get();
}

ImGuiFilamentBridge::~ImGuiFilamentBridge()
{
	m_engineManager->destroy(m_renderable);
}

//void ImGuiFilamentBridge::setDisplaySize(int width, int height, float scaleX, float scaleY)
//{
//	setContext();
//	ImGuiIO& io = ImGui::GetIO();
//	io.DisplaySize = ImVec2((float)width, (float)height);
//	io.DisplayFramebufferScale.x = scaleX;
//	io.DisplayFramebufferScale.y = scaleY;
//}

void ImGuiFilamentBridge::processImGuiCommands(ImDrawData& commands, const ImGuiIO& io)
{
	auto& rcm = m_engineManager->getEngine()->getRenderableManager();

	// Avoid rendering when minimized and scale coordinates for retina displays.
	const int fbwidth = static_cast<int>(commands.DisplaySize.x * commands.FramebufferScale.x);
	const int fbheight = static_cast<int>(commands.DisplaySize.y * commands.FramebufferScale.y);
	if (fbwidth == 0 || fbheight == 0)
		return;
	//commands.ScaleClipRects(commands.FramebufferScale);

	// Ensure that we have enough vertex buffers and index buffers.
	createBuffers(commands.CmdListsCount);

	// Count how many primitives we'll need, then create a Renderable builder.
	size_t primitiveCount = 0;
	//tl::unordered_map<uint64_t, filament::MaterialInstance*> scissorRects;
	for (int cmdListIndex = 0; cmdListIndex < commands.CmdListsCount; cmdListIndex++)
	{
		const ImDrawList* cmds = commands.CmdLists[cmdListIndex];
		primitiveCount += cmds->CmdBuffer.size();
	}
	auto rbuilder = RenderableManager::Builder(primitiveCount);
	rbuilder.boundingBox({{0, 0, 0}, {10000, 10000, 10000}}).culling(false);

	// Ensure that we have a material instance for each primitive.
	const size_t previousSize = m_materialInstances.size();
	if (primitiveCount > m_materialInstances.size())
	{
		//m_materialInstances.resize(primitiveCount);
		for (size_t i = previousSize; i < primitiveCount; i++)
			m_materialInstances.push_back(m_engineManager->adopt(m_material->createInstance()));
	}

	const ImVec2 clipOff = commands.DisplayPos; // (0,0) unless using multi-viewports
	const ImVec2 clipScale = commands.FramebufferScale; // (1,1) unless using retina display which are often (2,2)

	// Recreate the Renderable component and point it to the vertex buffers.
	rcm.destroy(m_renderable);
	size_t bufferIndex = 0;
	size_t primIndex = 0;
	for (int cmdListIndex = 0; cmdListIndex < commands.CmdListsCount; cmdListIndex++)
	{
		const ImDrawList* cmds = commands.CmdLists[cmdListIndex];
		size_t indexOffset = 0;
		populateVertexData(bufferIndex,
		                   cmds->VtxBuffer.Size * sizeof(ImDrawVert),
		                   cmds->VtxBuffer.Data,
		                   cmds->IdxBuffer.Size * sizeof(ImDrawIdx),
		                   cmds->IdxBuffer.Data);
		for (const auto& pcmd : cmds->CmdBuffer)
		{
			if (pcmd.UserCallback)
				pcmd.UserCallback(cmds, &pcmd);
			else
			{
				MaterialInstance& materialInstance = *m_materialInstances[primIndex];

				// Project scissor/clipping rectangles into framebuffer space
				const ImVec2 clipMin((pcmd.ClipRect.x - clipOff.x) * clipScale.x, (pcmd.ClipRect.y - clipOff.y) * clipScale.y);
				const ImVec2 clipMax((pcmd.ClipRect.z - clipOff.x) * clipScale.x, (pcmd.ClipRect.w - clipOff.y) * clipScale.y);

				materialInstance.setScissor(uint32_t(clipMin.x),
				                             uint32_t(fbheight - clipMax.y),
				                             (uint16_t)(clipMax.x - clipMin.x),
				                             (uint16_t)(clipMax.y - clipMin.y));
				TextureSampler sampler(TextureSampler::MinFilter::LINEAR_MIPMAP_LINEAR, TextureSampler::MagFilter::LINEAR);
				if (pcmd.TextureId)
					materialInstance.setParameter("albedo", (Texture const*)pcmd.TextureId, sampler);
				else
					materialInstance.setParameter("albedo", m_texture.get(), sampler);

				rbuilder
					.geometry(primIndex,
					          RenderableManager::PrimitiveType::TRIANGLES,
					          m_vertexBuffers[bufferIndex].get(),
					          m_indexBuffers[bufferIndex].get(),
					          indexOffset,
					          pcmd.ElemCount)
					.blendOrder(primIndex, tl::narrow<uint16_t>(primIndex))
					.material(primIndex, &materialInstance);
				primIndex++;
			}
			indexOffset += pcmd.ElemCount;
		}
		bufferIndex++;
	}
	if (commands.CmdListsCount > 0)
		rbuilder.build(*m_engineManager->getEngine(), m_renderable);
}

tl::unique_ref<filament::VertexBuffer> ImGuiFilamentBridge::createVertexBuffer(size_t capacity) const
{
	auto builder = VertexBuffer::Builder()
	               .vertexCount(uint32_t(capacity))
	               .bufferCount(1)
	               .attribute(VertexAttribute::POSITION,
	                          0,
	                          VertexBuffer::AttributeType::FLOAT2,
	                          0,
	                          sizeof(ImDrawVert))
	               .attribute(VertexAttribute::UV0,
	                          0,
	                          VertexBuffer::AttributeType::FLOAT2,
	                          sizeof(filament::math::float2),
	                          sizeof(ImDrawVert))
	               .attribute(VertexAttribute::COLOR,
	                          0,
	                          VertexBuffer::AttributeType::UBYTE4,
	                          2 * sizeof(filament::math::float2),
	                          sizeof(ImDrawVert))
	               .normalized(COLOR);

	return m_engineManager->adopt(builder.build(*m_engineManager->getEngine()));
}

tl::unique_ref<filament::IndexBuffer> ImGuiFilamentBridge::createIndexBuffer(size_t capacity) const
{
	auto builder = IndexBuffer::Builder()
	               .indexCount(uint32_t(capacity))
	               .bufferType(IndexBuffer::IndexType::UINT);
	return m_engineManager->adopt(builder.build(*m_engineManager->getEngine()));
}

void ImGuiFilamentBridge::createBuffers(int numRequiredBuffers)
{
	if (numRequiredBuffers > m_vertexBuffers.size())
	{
		const size_t previousSize = m_vertexBuffers.size();
		for (size_t i = previousSize; i < numRequiredBuffers; i++)
		{
			// Pick a reasonable starting capacity; it will grow if needed.
			m_vertexBuffers.push_back(createVertexBuffer(1000));
		}
	}
	if (numRequiredBuffers > m_indexBuffers.size())
	{
		const size_t previousSize = m_indexBuffers.size();
		for (size_t i = previousSize; i < numRequiredBuffers; i++)
		{
			// Pick a reasonable starting capacity; it will grow if needed.
			m_indexBuffers.push_back(createIndexBuffer(1000));
		}
	}
}

void ImGuiFilamentBridge::populateVertexData(size_t bufferIndex, size_t vbSizeInBytes, void* vbImguiData, size_t ibSizeInBytes, void* ibImguiData)
{
	// Create a new vertex buffer if the size isn't large enough, then copy the ImGui data into
	// a staging area since Filament's render thread might consume the data at any time.
	const size_t requiredVertCount = vbSizeInBytes / sizeof(ImDrawVert);
	const size_t capacityVertCount = m_vertexBuffers[bufferIndex]->getVertexCount();
	if (requiredVertCount > capacityVertCount)
		m_vertexBuffers[bufferIndex] = createVertexBuffer(requiredVertCount);

	{
		const size_t nVbBytes = requiredVertCount * sizeof(ImDrawVert);
		uint8_t* vbFilamentData = m_engineManager->acquireUploadMemory(nVbBytes);
		memcpy(vbFilamentData, vbImguiData, nVbBytes);
		m_vertexBuffers[bufferIndex]->setBufferAt(*m_engineManager->getEngine(), 0, VertexBuffer::BufferDescriptor(vbFilamentData, nVbBytes));
	}

	// Create a new index buffer if the size isn't large enough, then copy the ImGui data into
	// a staging area since Filament's render thread might consume the data at any time.
	const size_t requiredIndexCount = ibSizeInBytes / sizeof(ImDrawIdx);
	const size_t capacityIndexCount = m_indexBuffers[bufferIndex]->getIndexCount();
	if (requiredIndexCount > capacityIndexCount)
		m_indexBuffers[bufferIndex] = createIndexBuffer(requiredIndexCount);

	{
		const size_t nIbBytes = requiredIndexCount * sizeof(ImDrawIdx);
		uint8_t* ibFilamentData = m_engineManager->acquireUploadMemory(nIbBytes);
		memcpy(ibFilamentData, ibImguiData, nIbBytes);
		m_indexBuffers[bufferIndex]->setBuffer(*m_engineManager->getEngine(), IndexBuffer::BufferDescriptor(ibFilamentData, nIbBytes));
	}
}
} // namespace filagui