#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->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);
}
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->address = buffer.address;
	this->mapped = buffer.mapped;
	this->usage = buffer.usage;
	this->memoryProperties = buffer.memoryProperties;
	this->allocation = buffer.allocation;
	this->allocationInfo = buffer.allocationInfo;
}

void* ext::vulkan::Buffer::map( VkDeviceSize size, VkDeviceSize offset ) {
	if ( !mapped ) VK_CHECK_RESULT(vmaMapMemory( allocator, allocation, &mapped ));
	return mapped;
}
void ext::vulkan::Buffer::unmap() {
	if ( !mapped ) return;
	vmaUnmapMemory( allocator, allocation );
	mapped = nullptr;
}
/*
void* ext::vulkan::Buffer::map( VkDeviceSize size, VkDeviceSize offset ) const {
	void* mapped{};
	VK_CHECK_RESULT(vmaMapMemory( allocator, allocation, &mapped ));
	return mapped;
}
void ext::vulkan::Buffer::unmap() const {
	vmaUnmapMemory( allocator, allocation );
}
VkResult ext::vulkan::Buffer::bind( VkDeviceSize offset ) {
	return VK_SUCCESS;
}
VkResult ext::vulkan::Buffer::flush( VkDeviceSize size, VkDeviceSize offset ) const {
	return VK_SUCCESS;
}

VkResult ext::vulkan::Buffer::invalidate( VkDeviceSize size, VkDeviceSize offset ) {
	return VK_SUCCESS;
}
void ext::vulkan::Buffer::copyTo( void* data, VkDeviceSize size ) {
	assert(mapped);
	memcpy(mapped, data, size);
}
*/

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 );
}

size_t ext::vulkan::Buffer::getAddress() {
//	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::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 vkGetBufferDeviceAddressKHR(this->device ? *this->device : ext::vulkan::device, &info);
}

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 ) {
	this->device = &device;
	this->alignment = alignment;
}
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 );
	}

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

//	if ( usage != VK_BUFFER_USAGE_TRANSFER_SRC_BIT ) usage |= VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT;

	VK_CHECK_RESULT(device->createBuffer(
		nullptr,
		length,
		usage,
		memoryProperties,
		*this
	));

	if ( data && length ) update( data, length, stage );
	
/*
	{
		uf::stl::string type;

		if ( usage & VK_BUFFER_USAGE_TRANSFER_SRC_BIT ) type += "transfer_src,";
		if ( usage & VK_BUFFER_USAGE_TRANSFER_DST_BIT ) type += "transfer_dst,";
		if ( usage & VK_BUFFER_USAGE_UNIFORM_TEXEL_BUFFER_BIT ) type += "uniform_texel,";
		if ( usage & VK_BUFFER_USAGE_STORAGE_TEXEL_BUFFER_BIT ) type += "storage_texel,";
		if ( usage & VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT ) type += "uniform,";
		if ( usage & VK_BUFFER_USAGE_STORAGE_BUFFER_BIT ) type += "storage,";
		if ( usage & VK_BUFFER_USAGE_INDEX_BUFFER_BIT ) type += "index,";
		if ( usage & VK_BUFFER_USAGE_VERTEX_BUFFER_BIT ) type += "vertex,";
		if ( usage & VK_BUFFER_USAGE_INDIRECT_BUFFER_BIT ) type += "indirect,";
		if ( usage & VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_STORAGE_BIT_KHR ) type += "acceleration_structure,";
		if ( usage & VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT ) type += "address,";
		if ( usage & VK_BUFFER_USAGE_SHADER_BINDING_TABLE_BIT_KHR ) type += "binding_table,";

		UF_MSG_DEBUG("CREATING BUFFER: {}: {}({})", fmt::ptr(this->buffer), type, usage);
	}
*/

	this->usage = usage;
}
bool ext::vulkan::Buffer::update( const void* data, VkDeviceSize length, bool stage ) const {
//	if ( !data || !length ) return false;
	if ( !length ) return false;
	if ( !buffer ) return false;
	if ( length > allocationInfo.size ) {
		UF_MSG_WARNING("Buffer update of {} exceeds buffer size of {}", length, allocationInfo.size );

		Buffer& b = *const_cast<Buffer*>(this);
		b.destroy(true);
		b.initialize( data, length, usage, memoryProperties, stage );
		return true;
	}
	if ( !data ) return false;
	if ( !stage ) {
		auto* self = const_cast<ext::vulkan::Buffer*>(this);
		void* map = self->map();
		memcpy(map, data, length);
		self->unmap();
		return false;
	}

	ext::vulkan::Device* device = this->device ? this->device : &ext::vulkan::device;
	
	Buffer staging = device->fetchTransientBuffer(
//	Buffer staging = device->createBuffer(
		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
		VkBufferCopy region = {};
		region.size = length;
		device->UF_CHECKPOINT_MARK( commandBuffer, pod::Checkpoint::GENERIC, "copyBuffer" );
		vkCmdCopyBuffer(commandBuffer, staging.buffer, buffer, 1, &region);
	device->flushCommandBuffer(commandBuffer);
	return false;
}

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

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