#if UF_USE_VULKAN

#define VMA_VULKAN_VERSION 1002000
#define VMA_DEFAULT_LARGE_HEAP_BLOCK_SIZE (200ull * 1024 * 1024)
#define VMA_IMPLEMENTATION

#include <uf/ext/vulkan/buffer.h>
#include <uf/ext/vulkan/vulkan.h>
#include <uf/ext/vulkan/device.h>

void ext::vulkan::Buffer::swap( ext::vulkan::Buffer& buffer ) {
	std::swap(this->aliased, buffer.aliased);
	std::swap(this->device, buffer.device);
	std::swap(this->buffer, buffer.buffer);
	std::swap(this->memory, buffer.memory);
	std::swap(this->descriptor, buffer.descriptor);
	std::swap(this->alignment, buffer.alignment);
	std::swap(this->stride, buffer.stride);
	std::swap(this->address, buffer.address);
	std::swap(this->mapped, buffer.mapped);
	std::swap(this->usage, buffer.usage);
	std::swap(this->memoryProperties, buffer.memoryProperties);
	std::swap(this->allocation, buffer.allocation);
	std::swap(this->allocationInfo, buffer.allocationInfo);
	std::swap(this->count, buffer.count);
	std::swap(this->written, buffer.written);
	std::swap(this->staged, buffer.staged);
}
ext::vulkan::Buffer ext::vulkan::Buffer::alias() const {
	ext::vulkan::Buffer buffer;
	buffer.aliasBuffer(*this);
	return buffer;
}
void ext::vulkan::Buffer::aliasBuffer( const ext::vulkan::Buffer& buffer ) {
	this->aliased = true;
	this->device = buffer.device;
	this->buffer = buffer.buffer;
	this->memory = buffer.memory;
	this->descriptor = buffer.descriptor;
	this->alignment = buffer.alignment;
	this->stride = buffer.alignment;
	this->address = buffer.address;
	this->mapped = buffer.mapped;
	this->usage = buffer.usage;
	this->memoryProperties = buffer.memoryProperties;
	this->allocation = buffer.allocation;
	this->allocationInfo = buffer.allocationInfo;
	this->count = buffer.count;
	this->written = buffer.written;
	this->staged = buffer.staged;
}

void* ext::vulkan::Buffer::map( VkDeviceSize size, VkDeviceSize offset ) {
	if ( !mapped ) VK_CHECK_RESULT(vmaMapMemory( allocator, allocation, &mapped ));
	return static_cast<char*>(mapped) + offset;;
}
void ext::vulkan::Buffer::unmap() {
	if ( !mapped ) return;
	vmaUnmapMemory( allocator, allocation );
	mapped = nullptr;
}

void ext::vulkan::Buffer::updateDescriptor( VkDeviceSize size, VkDeviceSize offset ) {
	descriptor.offset = offset;
	descriptor.buffer = buffer;
	descriptor.range = size;
}

void ext::vulkan::Buffer::allocate( VkBufferCreateInfo bufferCreateInfo ) {
	VmaAllocationCreateInfo allocCreateInfo = {};
	
	allocCreateInfo.usage = VMA_MEMORY_USAGE_CPU_TO_GPU;
	allocCreateInfo.preferredFlags = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT;
	if ( bufferCreateInfo.usage & (VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_INDEX_BUFFER_BIT) ) {
		allocCreateInfo.usage = VMA_MEMORY_USAGE_GPU_ONLY;
	}

	vmaCreateBufferWithAlignment( allocator, &bufferCreateInfo, &allocCreateInfo, alignment, &buffer, &allocation, &allocationInfo );
	VK_REGISTER_HANDLE( buffer );
}

size_t ext::vulkan::Buffer::getAddress() const {
//	if ( !(usage & VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT) ) UF_MSG_DEBUG("CALLING GETADDRESS ON BUFFER WITHOUT ADDRESS BIT: {}", fmt::ptr(this->buffer));
	if ( this->address ) return this->address;

	VkBufferDeviceAddressInfoKHR info{};
	info.sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO;
	info.buffer = buffer;
	return (this->address = vkGetBufferDeviceAddressKHR(this->device ? *this->device : ext::vulkan::device, &info));
}

size_t ext::vulkan::Buffer::getLength( ) const {
	return allocationInfo.size;
}
size_t ext::vulkan::Buffer::getOffset( size_t i ) const {
	return this->stride * i;
}

ext::vulkan::Buffer::~Buffer() {
//	this->destroy(); // bad, will get called every time a vector gets resized, buffer owners will clean it up in their own destructors
}
void ext::vulkan::Buffer::initialize( ext::vulkan::Device& device, size_t alignment, bool staged ) {
	this->device = &device;
	this->alignment = alignment;
	this->staged = staged;
}
void ext::vulkan::Buffer::destroy(bool defer) {
	if ( !device || aliased ) return;
	if ( defer ) {
		ext::vulkan::mutex.lock();
		device->transient.buffers.emplace_back(*this);
		ext::vulkan::mutex.unlock();
	} else if ( buffer ) {
		vmaDestroyBuffer( allocator, buffer, allocation );
		VK_UNREGISTER_HANDLE( buffer );
	}

	this->buffer = {};
	this->memory = {};
	this->descriptor = {};
//	this->alignment = {};
	this->stride = {};
	this->address = {};
	this->mapped = {};
//	this->usage = {};
//	this->memoryProperties = {};
	this->allocation = {};
	this->allocationInfo = {};
	this->count = 1;
	this->written = false;
	this->staged = VK_DEFAULT_STAGE_BUFFERS;
}
void ext::vulkan::Buffer::initialize( const void* data, VkDeviceSize length, VkBufferUsageFlags usage, VkMemoryPropertyFlags memoryProperties ) {
	if ( !device ) device = &ext::vulkan::device;
	if ( this->staged ) usage |= VK_BUFFER_USAGE_TRANSFER_DST_BIT; // implicitly set properties

	// assume all UBOs are dynamic
	auto totalLength = (this->stride = length);
	if ( VK_UBO_USE_N_BUFFERS && usage & VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT ) {
		this->count = ext::vulkan::swapchain.buffers;
		this->alignment = device->properties.limits.minUniformBufferOffsetAlignment;
		this->stride = ALIGNED_SIZE( length, this->alignment );
		totalLength = this->stride * this->count;
	}
	VK_CHECK_RESULT(device->createBuffer( nullptr, totalLength, usage, memoryProperties, *this ));

	if ( length != totalLength ) this->updateDescriptor( length, 0 );
	if ( data && length ) update( data, length );
}
bool ext::vulkan::Buffer::update( const void* data, VkDeviceSize length ) const {
//	if ( !data || !length ) return false;
	if ( !length ) return false;
	if ( !buffer ) return false;

	// to-do: fix this because it's a thorn in my side when a mesh needs to update
	if ( length > allocationInfo.size ) {
		//UF_MSG_WARNING("Buffer update of {} exceeds buffer size of {}", length, allocationInfo.size);

		auto savedUsage = usage;
		auto savedMemProps = memoryProperties;
		auto savedAlignment = alignment;

		Buffer& self = const_cast<Buffer&>(*this);

		self.destroy(true);
		self.initialize(*device, savedAlignment);
		self.initialize(data, length, savedUsage, savedMemProps);

		return true;
	}
	if ( !data ) return false;

	bool broadcast = !this->written;
	this->written = true;

	STATIC_THREAD_LOCAL(uf::stl::vector<VkBufferCopy>, regions);

	if ( broadcast && this->count > 1 ) {
		for ( size_t i = 0; i < this->count; i++ ) {
			auto& region = regions.emplace_back();
			region.size = length;
			region.srcOffset = 0;
			region.dstOffset = this->getOffset( i );
		}
	} else {
		auto& region = regions.emplace_back();
		region.size = length;
		region.srcOffset = 0;
		region.dstOffset = 0;
		if ( VK_UBO_USE_N_BUFFERS && this->count == ext::vulkan::swapchain.buffers ) {
			region.dstOffset = this->getOffset( states::currentBuffer );
		}
	}

	if ( !this->staged ) {
		auto* self = const_cast<ext::vulkan::Buffer*>(this);
		void* map = self->map(length);
		for ( const auto& region : regions ) memcpy(static_cast<char*>(map) + region.dstOffset, data, length);
		self->unmap();
		return false;
	}

	ext::vulkan::Device* device = this->device ? this->device : &ext::vulkan::device;
	
	Buffer staging = device->fetchTransientBuffer(
		data, length,
		VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
	);
	auto commandBuffer = device->fetchCommandBuffer( QueueEnum::TRANSFER ); // waits on finish
		device->UF_CHECKPOINT_MARK( commandBuffer, pod::Checkpoint::GENERIC, "copyBuffer" );
		vkCmdCopyBuffer(commandBuffer, staging.buffer, buffer, regions.size(), regions.data());
	device->flushCommandBuffer(commandBuffer);
	return false;
}

//
// Buffers
void ext::vulkan::Buffers::initialize( Device& device, bool staged ) {
	this->device = &device;
	this->staged = staged;
}
/*
ext::vulkan::Buffers::~Buffers() {
	for ( auto& buffer : buffers ) buffer.destroy();
	buffers.clear();
}
*/
void ext::vulkan::Buffers::destroy(bool defer) {
	for ( auto& buffer : buffers ) if ( buffer.device ) buffer.destroy(defer);
	buffers.clear();
}

size_t ext::vulkan::Buffers::initializeBuffer( const void* data, VkDeviceSize length, VkBufferUsageFlags usage, VkMemoryPropertyFlags memoryProperties ) {
	size_t index = buffers.size();
	auto& buffer = buffers.emplace_back();
	buffer.initialize( *device, requestedAlignment, this->staged );
	buffer.initialize( data, length, usage, memoryProperties );
	return index;
}
bool ext::vulkan::Buffers::updateBuffer( const void* data, VkDeviceSize length, const Buffer& buffer ) const {
	return buffer.update( data, length );
}
#endif