#include <PreCompiled.h>
#include <Renderer/Image.h>
#include <Maths/Vec4.h>
#include <Maths/Vec2.h>
#include <Renderer/RenderCore.h>
#include <Utils/CallOnExit.h>
#include <Core/Memory.h>

#define STB_IMAGE_IMPLEMENTATION

#define STBI_ASSERT(x) (mlx::Assert(x, "internal stb assertion " #x))
#define STBI_MALLOC(x) (mlx::MemManager::Malloc(x))
#define STBI_REALLOC(p, x) (mlx::MemManager::Realloc(p, x))
#define STBI_FREE(x) (mlx::MemManager::Free(x))

#ifdef MLX_COMPILER_GCC
	#pragma GCC diagnostic push
	#pragma GCC diagnostic ignored "-Wstringop-overflow"
		#include <stb_image.h>
	#pragma GCC diagnostic pop
#else
	#include <stb_image.h>
#endif

namespace mlx
{
	MLX_FORCEINLINE mlx_color ReverseColor(mlx_color color)
	{
		mlx_color reversed_color;
		reversed_color.r = color.a;
		reversed_color.g = color.b;
		reversed_color.b = color.g;
		reversed_color.a = color.r;
		return reversed_color;
	}

	void Image::Init(ImageType type, std::uint32_t width, std::uint32_t height, VkFormat format, VkImageTiling tiling, VkImageUsageFlags usage, bool is_multisampled, [[maybe_unused]] std::string_view debug_name)
	{
		MLX_PROFILE_FUNCTION();
		Verify(width > 0 && height > 0, "width or height cannot be null");

		m_type = type;
		m_width = width;
		m_height = height;
		m_format = format;
		m_tiling = tiling;
		m_is_multisampled = is_multisampled;
		#ifdef DEBUG
			m_debug_name = std::move(debug_name);
		#endif

		VmaAllocationCreateInfo alloc_info{};
		alloc_info.usage = VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE;

		VkImageCreateInfo image_info{};
		image_info.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
		image_info.imageType = VK_IMAGE_TYPE_2D;
		image_info.extent.width = width;
		image_info.extent.height = height;
		image_info.extent.depth = 1;
		image_info.mipLevels = 1;
		image_info.arrayLayers = 1;
		image_info.format = format;
		image_info.tiling = tiling;
		image_info.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
		image_info.usage = usage;
		image_info.samples = (m_is_multisampled ? VK_SAMPLE_COUNT_4_BIT : VK_SAMPLE_COUNT_1_BIT);
		image_info.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
		#ifdef DEBUG
			m_allocation = RenderCore::Get().GetAllocator().CreateImage(&image_info, &alloc_info, m_image, m_debug_name.c_str());
		#else
			m_allocation = RenderCore::Get().GetAllocator().CreateImage(&image_info, &alloc_info, m_image, nullptr);
		#endif
	}

	void Image::CreateImageView(VkImageViewType type, VkImageAspectFlags aspect_flags, int layer_count) noexcept
	{
		MLX_PROFILE_FUNCTION();
		m_image_view = kvfCreateImageView(RenderCore::Get().GetDevice(), m_image, m_format, type, aspect_flags, layer_count);
		#ifdef MLX_HAS_DEBUG_UTILS_FUNCTIONS
			VkDebugUtilsObjectNameInfoEXT name_info{};
			name_info.sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT;
			name_info.objectType = VK_OBJECT_TYPE_IMAGE_VIEW;
			name_info.objectHandle = reinterpret_cast<std::uint64_t>(m_image_view);
			name_info.pObjectName = m_debug_name.c_str();
			RenderCore::Get().vkSetDebugUtilsObjectNameEXT(RenderCore::Get().GetDevice(), &name_info);
		#endif
	}

	void Image::CreateSampler() noexcept
	{
		MLX_PROFILE_FUNCTION();
		m_sampler = kvfCreateSampler(RenderCore::Get().GetDevice(), VK_FILTER_NEAREST, VK_SAMPLER_ADDRESS_MODE_REPEAT, VK_SAMPLER_MIPMAP_MODE_NEAREST);
	}

	void Image::TransitionLayout(VkImageLayout new_layout, VkCommandBuffer cmd)
	{
		MLX_PROFILE_FUNCTION();
		if(new_layout == m_layout)
			return;
		bool is_single_time_cmd_buffer = (cmd == VK_NULL_HANDLE);
		if(is_single_time_cmd_buffer)
			cmd = kvfCreateCommandBuffer(RenderCore::Get().GetDevice());
		kvfTransitionImageLayout(RenderCore::Get().GetDevice(), m_image, KVF_IMAGE_COLOR, cmd, m_format, m_layout, new_layout, is_single_time_cmd_buffer);
		if(is_single_time_cmd_buffer)
			kvfDestroyCommandBuffer(RenderCore::Get().GetDevice(), cmd);
		m_layout = new_layout;
	}

	void Image::Clear(VkCommandBuffer cmd, Vec4f color)
	{
		MLX_PROFILE_FUNCTION();
		VkImageSubresourceRange subresource_range{};
		subresource_range.baseMipLevel = 0;
		subresource_range.layerCount = 1;
		subresource_range.levelCount = 1;
		subresource_range.baseArrayLayer = 0;

		bool is_single_time_cmd_buffer = (cmd == VK_NULL_HANDLE);
		if(is_single_time_cmd_buffer)
		{
			cmd = kvfCreateCommandBuffer(RenderCore::Get().GetDevice());
			kvfBeginCommandBuffer(cmd, VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT);
		}

		VkImageLayout old_layout = m_layout;
		TransitionLayout(VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, cmd);
		subresource_range.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
		VkClearColorValue clear_color = VkClearColorValue({ { color.x, color.y, color.z, color.w } });
		RenderCore::Get().vkCmdClearColorImage(cmd, m_image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, &clear_color, 1, &subresource_range);
		TransitionLayout(old_layout, cmd);

		if(is_single_time_cmd_buffer)
		{
			kvfEndCommandBuffer(cmd);
			VkFence fence = kvfCreateFence(RenderCore::Get().GetDevice());
			kvfSubmitSingleTimeCommandBuffer(RenderCore::Get().GetDevice(), cmd, KVF_GRAPHICS_QUEUE, fence);
			kvfDestroyFence(RenderCore::Get().GetDevice(), fence);
			kvfDestroyCommandBuffer(RenderCore::Get().GetDevice(), cmd);
		}
	}

	void Image::DestroySampler() noexcept
	{
		MLX_PROFILE_FUNCTION();
		if(m_sampler != VK_NULL_HANDLE)
			kvfDestroySampler(RenderCore::Get().GetDevice(), m_sampler);
		m_sampler = VK_NULL_HANDLE;
	}

	void Image::DestroyImageView() noexcept
	{
		MLX_PROFILE_FUNCTION();
		if(m_image_view != VK_NULL_HANDLE)
			kvfDestroyImageView(RenderCore::Get().GetDevice(), m_image_view);
		m_image_view = VK_NULL_HANDLE;
	}

	void Image::Destroy() noexcept
	{
		MLX_PROFILE_FUNCTION();
		DestroySampler();
		DestroyImageView();

		if(m_image != VK_NULL_HANDLE)
		{
			#ifdef DEBUG
				RenderCore::Get().GetAllocator().DestroyImage(m_allocation, m_image, m_debug_name.c_str());
			#else
				RenderCore::Get().GetAllocator().DestroyImage(m_allocation, m_image, nullptr);
			#endif
		}
		m_image = VK_NULL_HANDLE;
		m_layout = VK_IMAGE_LAYOUT_UNDEFINED;
		m_width = 0;
		m_height = 0;
		m_is_multisampled = false;
	}

	void Texture::Init(CPUBuffer pixels, std::uint32_t width, std::uint32_t height, VkFormat format, bool is_multisampled, [[maybe_unused]] std::string_view debug_name)
	{
		MLX_PROFILE_FUNCTION();
		Image::Init(ImageType::Color, width, height, format, VK_IMAGE_TILING_OPTIMAL, VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT, is_multisampled, std::move(debug_name));
		Image::CreateImageView(VK_IMAGE_VIEW_TYPE_2D, VK_IMAGE_ASPECT_COLOR_BIT);
		Image::CreateSampler();
		if(pixels)
		{
			GPUBuffer staging_buffer;
			std::size_t size = width * height * kvfFormatSize(format);
			staging_buffer.Init(BufferType::Staging, size, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, std::move(pixels), debug_name);
			VkCommandBuffer cmd = kvfCreateCommandBuffer(RenderCore::Get().GetDevice());
			kvfBeginCommandBuffer(cmd, VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT);
			TransitionLayout(VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, cmd);
			kvfCopyBufferToImage(cmd, Image::Get(), staging_buffer.Get(), staging_buffer.GetOffset(), VK_IMAGE_ASPECT_COLOR_BIT, { width, height, 1 });
			kvfEndCommandBuffer(cmd);
			VkFence fence = kvfCreateFence(RenderCore::Get().GetDevice());
			kvfSubmitSingleTimeCommandBuffer(RenderCore::Get().GetDevice(), cmd, KVF_GRAPHICS_QUEUE, fence);
			kvfDestroyFence(RenderCore::Get().GetDevice(), fence);
			kvfDestroyCommandBuffer(RenderCore::Get().GetDevice(), cmd);
			staging_buffer.Destroy();
		}
		TransitionLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
	}

	void Texture::Destroy() noexcept
	{
		if(m_staging_buffer.has_value())
			m_staging_buffer->Destroy();
		Image::Destroy();
	}

	void Texture::SetPixel(int x, int y, mlx_color color) noexcept
	{
		MLX_PROFILE_FUNCTION();
		if(x < 0 || y < 0 || static_cast<std::uint32_t>(x) >= m_width || static_cast<std::uint32_t>(y) >= m_height)
			return;
		if(!m_staging_buffer.has_value())
			OpenCPUBuffer();
		if constexpr(std::endian::native == std::endian::little)
			m_staging_buffer->GetMap<mlx_color*>()[(y * m_width) + x] = ReverseColor(color);
		else
			m_staging_buffer->GetMap<mlx_color*>()[(y * m_width) + x] = color;
		m_has_been_modified = true;
	}

	void Texture::SetRegion(int x, int y, int w, int h, mlx_color* pixels) noexcept
	{
		MLX_PROFILE_FUNCTION();
		if(w < 0 || h < 0 || x < -w || y < -h
			|| x >= static_cast<int>(m_width) || y >= static_cast<int>(m_height))
			return;
		if(!m_staging_buffer.has_value())
			OpenCPUBuffer();
		const int
			start_x = std::max<int>(x, 0),
			start_y = std::max<int>(y, 0),
			start_row = start_y * m_width,
			start_i = (start_y - y) * w + (start_x - x),
			end_x = std::min<int>(x + w, m_width),
			end_y = std::min<int>(y + h, m_height),
			end_row = end_y * m_width,
			incr = (x + w) - end_x + (start_x - x);
		for(int i = start_i, dx = start_x, row = start_row;; i++, dx++)
		{
			if(dx >= end_x)
			{
				i += incr;
				dx = start_x;
				row += m_width;
				if(row >= end_row)
					break;
			}
			if constexpr(std::endian::native == std::endian::little)
				m_staging_buffer->GetMap<mlx_color*>()[row + dx] = ReverseColor(pixels[i]);
			else
				m_staging_buffer->GetMap<mlx_color*>()[row + dx] = pixels[i];
		}
		m_has_been_modified = true;
	}

	void Texture::SetLinearRegion(int x, int y, std::size_t len, mlx_color* pixels) noexcept
	{
		MLX_PROFILE_FUNCTION();
		if(x >= static_cast<int>(m_width) || y >= static_cast<int>(m_height))
			return;
		if(!m_staging_buffer.has_value())
			OpenCPUBuffer();
		int
			start = y * m_width + x,
			dest_start = std::max<int>(start, 0),
			src_start = dest_start - start,
			dest_end = std::min<int>(start + len, m_width * m_height);
		if constexpr(std::endian::native == std::endian::little)
		{
			for(int i = dest_start, j = src_start; i < dest_end; i++, j++)
				m_staging_buffer->GetMap<mlx_color*>()[i] = ReverseColor(pixels[j]);
		}
		else
		{
			std::memcpy(
				&m_staging_buffer->GetMap<mlx_color*>()[dest_start],
				&pixels[src_start], dest_end - dest_start);
		}
		m_has_been_modified = true;
	}

	mlx_color Texture::GetPixel(int x, int y) noexcept
	{
		MLX_PROFILE_FUNCTION();
		if(x < 0 || y < 0 || static_cast<std::uint32_t>(x) >= m_width || static_cast<std::uint32_t>(y) >= m_height)
			return { .rgba = 0x00000000 };
		if(!m_staging_buffer.has_value())
			OpenCPUBuffer();
		if constexpr(std::endian::native == std::endian::little)
			return ReverseColor(m_staging_buffer->GetMap<mlx_color*>()[(y * m_width) + x]);
		else
			return m_staging_buffer->GetMap<mlx_color*>()[(y * m_width) + x];
	}

	void Texture::GetRegion(int x, int y, int w, int h, mlx_color* dst) noexcept
	{
		MLX_PROFILE_FUNCTION();
		if(w < 0 || h < 0 || x < 0 || y < 0 || static_cast<std::uint32_t>(x + w) >= m_width || static_cast<std::uint32_t>(y + h) >= m_height)
			return;
		if(!m_staging_buffer.has_value())
			OpenCPUBuffer();
		for(std::uint32_t i = 0, dx = x, row = y * m_width;; i++, dx++)
		{
			if(dx >= m_width)
			{
				dx = x;
				row += m_width;
				if(row >= m_height * m_width)
					break;
			}
			if constexpr(std::endian::native == std::endian::little)
				dst[i] = ReverseColor(m_staging_buffer->GetMap<mlx_color*>()[row + dx]);
			else
				dst[i] = m_staging_buffer->GetMap<mlx_color*>()[row + dx];
		}
	}

	void Texture::Clear(VkCommandBuffer cmd, Vec4f color)
	{
		MLX_PROFILE_FUNCTION();
		Image::Clear(cmd, color);
		if(m_staging_buffer.has_value())
		{
			mlx_color processed_color;
			processed_color.r = static_cast<std::uint8_t>(color.r * 255.f);
			processed_color.g = static_cast<std::uint8_t>(color.g * 255.f);
			processed_color.b = static_cast<std::uint8_t>(color.b * 255.f);
			processed_color.a = static_cast<std::uint8_t>(color.a * 255.f);
			for(std::size_t i = 0; i < m_width * m_height; i++)
				m_staging_buffer->GetMap<mlx_color*>()[i] = processed_color;
		}
	}

	void Texture::Update(VkCommandBuffer cmd)
	{
		MLX_PROFILE_FUNCTION();
		if(!m_has_been_modified)
			return;
		VkImageLayout old_layout = m_layout;
		TransitionLayout(VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, cmd);
		kvfCopyBufferToImage(cmd, Image::Get(), m_staging_buffer->Get(), m_staging_buffer->GetOffset(), VK_IMAGE_ASPECT_COLOR_BIT, { m_width, m_height, 1 });
		TransitionLayout(old_layout, cmd);

		m_has_been_modified = false;
	}

	void Texture::OpenCPUBuffer()
	{
		MLX_PROFILE_FUNCTION();
		if(m_staging_buffer.has_value())
			return;
		#ifdef DEBUG
			DebugLog("Texture: enabling CPU mapping for '%'", m_debug_name);
		#endif
		m_staging_buffer.emplace();
		std::size_t size = m_width * m_height * kvfFormatSize(m_format);
		#ifdef DEBUG
			m_staging_buffer->Init(BufferType::Staging, size, VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, {}, m_debug_name);
		#else
			m_staging_buffer->Init(BufferType::Staging, size, VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, {}, {});
		#endif

		SyncCPUBuffer();
	}

	void Texture::SyncCPUBuffer(VkCommandBuffer cmd)
	{
		if(!m_staging_buffer.has_value())
			return;
		VkImageLayout old_layout = m_layout;
		if(cmd == VK_NULL_HANDLE)
			cmd = kvfCreateCommandBuffer(RenderCore::Get().GetDevice());
		kvfBeginCommandBuffer(cmd, VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT);
		TransitionLayout(VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, cmd);
		kvfCopyImageToBuffer(cmd, m_staging_buffer->Get(), m_image, m_staging_buffer->GetOffset(), VK_IMAGE_ASPECT_COLOR_BIT, { m_width, m_height, 1 });
		TransitionLayout(old_layout, cmd);
		kvfEndCommandBuffer(cmd);
		VkFence fence = kvfCreateFence(RenderCore::Get().GetDevice());
		kvfSubmitSingleTimeCommandBuffer(RenderCore::Get().GetDevice(), cmd, KVF_GRAPHICS_QUEUE, fence);
		kvfDestroyFence(RenderCore::Get().GetDevice(), fence);
		kvfDestroyCommandBuffer(RenderCore::Get().GetDevice(), cmd);
	}

	void Texture::CopyTo(Texture& other)
	{
		VkImageLayout old_layout = m_layout;
		VkImageLayout other_old_layout = other.GetLayout();

		VkImageSubresourceLayers subresource{};
		subresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
		subresource.mipLevel = 0;
		subresource.layerCount = 1;
		subresource.baseArrayLayer = 0;

		VkExtent3D extent{};
		extent.width = m_width;
		extent.height = m_height;
		extent.depth = 1;

		VkOffset3D offset{};
		offset.x = 0;
		offset.y = 0;
		offset.z = 0;

		VkImageCopy region{};
		region.srcSubresource = subresource;
		region.dstSubresource = subresource;
		region.extent = extent;
		region.srcOffset = offset;
		region.dstOffset = offset;

		VkCommandBuffer cmd = kvfCreateCommandBuffer(RenderCore::Get().GetDevice());
		kvfBeginCommandBuffer(cmd, VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT);

		TransitionLayout(VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, cmd);
		other.TransitionLayout(VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, cmd);

		kvfCopyImageToImage(cmd, m_image, m_layout, other.Get(), other.GetLayout(), 1, &region);

		TransitionLayout(old_layout, cmd);
		other.TransitionLayout(other_old_layout, cmd);

		kvfEndCommandBuffer(cmd);

		SyncCPUBuffer(cmd);
	}

	void Texture::Resize(std::uint32_t width, std::uint32_t height)
	{
		#ifdef DEBUG
			Texture new_texture = Texture(CPUBuffer{}, width, height, m_format, m_is_multisampled, m_debug_name);
		#else
			Texture new_texture = Texture(CPUBuffer{}, width, height, m_format, m_is_multisampled, std::string_view{});
		#endif

		if(m_staging_buffer.has_value())
		{
			new_texture.OpenCPUBuffer();
			new_texture.m_staging_buffer->CopyFrom(*m_staging_buffer);
		}

		// Suboptimal operations, should bake all of them in a single command buffer
		new_texture.Clear(VK_NULL_HANDLE, Vec4f{ 0.f });
		CopyTo(new_texture);

		Swap(new_texture);

		#ifdef DEBUG
			DebugLog("Texture: resized '%'", m_debug_name);
		#endif
	}

	void Texture::Swap(Texture& texture) noexcept
	{
		MLX_PROFILE_FUNCTION();

		#ifdef DEBUG
			std::swap(m_debug_name, texture.m_debug_name);
		#endif
		std::swap(m_allocation, texture.m_allocation);
		std::swap(m_image, texture.m_image);
		std::swap(m_image_view, texture.m_image_view);
		std::swap(m_sampler, texture.m_sampler);
		std::swap(m_format, texture.m_format);
		std::swap(m_tiling, texture.m_tiling);
		std::swap(m_layout, texture.m_layout);
		std::swap(m_type, texture.m_type);
		std::swap(m_width, texture.m_width);
		std::swap(m_height, texture.m_height);
		std::swap(m_is_multisampled, texture.m_is_multisampled);

		if(m_staging_buffer.has_value() && texture.m_staging_buffer.has_value())
			m_staging_buffer->Swap(*texture.m_staging_buffer);
		else if(m_staging_buffer.has_value())
			m_staging_buffer.reset();
		else if(texture.m_staging_buffer.has_value())
			texture.m_staging_buffer.reset();

		m_has_been_modified = true;
		texture.m_has_been_modified = true;
	}

	Texture* StbTextureLoad(const std::filesystem::path& file, int* w, int* h)
	{
		using namespace std::literals;
		MLX_PROFILE_FUNCTION();
		std::string filename = file.string();

		if(file.stem() == "terracotta.pie")
			Message("banana, banana, banana, banana, terracotta banana terracotta, terracotta pie");

		if(!std::filesystem::exists(file))
		{
			Error("Image loader: file not found %", file);
			return nullptr;
		}
		if(stbi_is_hdr(filename.c_str()))
		{
			Error("Image loader: unsupported image format from % (HDR image)", file);
			return nullptr;
		}

		Vec2i size;
		int channels;

		std::uint8_t* data = stbi_load(filename.c_str(), &size.x, &size.y, &channels, 4);
		if(data == nullptr)
		{
			Error("Image loader: could not load % due to %", file, stbi_failure_reason());
			return nullptr;
		}
		CallOnExit defer([&]() { stbi_image_free(data); });

		CPUBuffer buffer(size.x * size.y * 4);
		std::memcpy(buffer.GetData(), data, buffer.GetSize());

		if(w != nullptr)
			*w = size.x;
		if(h != nullptr)
			*h = size.y;
		return new Texture(std::move(buffer), size.x, size.y, VK_FORMAT_R8G8B8A8_SRGB, false, std::move(filename));
	}
}
