// =============================================
//  Aster: device.cpp
//  Copyright (c) 2020-2024 Anish Bhobe
// =============================================

#include "device.h"

#include "context.h"
#include "window.h"

#include <map>
#include <set>
#include <vector>

Device::Device(Device &&_other) noexcept :
		physical_device{ _other.physical_device }, device{ std::exchange(_other.device, nullptr) }, queues{ _other.queues }, allocator{ std::exchange(_other.allocator, nullptr) }, name{ std::move(_other.name) } {}

Device &Device::operator=(Device &&_other) noexcept {
	if (this == &_other)
		return *this;
	physical_device = _other.physical_device;
	device = std::exchange(_other.device, nullptr);
	queues = _other.queues;
	allocator = std::exchange(_other.allocator, nullptr);
	name = std::move(_other.name);
	return *this;
}

Device::Device(const std::string_view &_name, Context *_context, const PhysicalDevice &_physical_device_info, const vk::PhysicalDeviceFeatures &_enabled_features) :
		physical_device{ _physical_device_info },
		name{ _name } {
	const auto &physical_device = _physical_device_info.device;
	const auto &queue_families = _physical_device_info.queue_families;

	// Logical Device
	std::map<u32, u16> unique_queue_families;
	unique_queue_families[queue_families.graphics_idx]++;
	unique_queue_families[queue_families.present_idx]++;
	unique_queue_families[queue_families.transfer_idx]++;
	unique_queue_families[queue_families.compute_idx]++;

	std::array<f32, 4> queue_priority = { 1.0f, 1.0f, 1.0f, 1.0f };
	std::vector<vk::DeviceQueueCreateInfo> queue_create_infos;
	for (auto &[index_, count_] : unique_queue_families) {
		queue_create_infos.push_back({
				.queueFamilyIndex = index_,
				.queueCount = count_,
				.pQueuePriorities = queue_priority.data(),
		});
	}
	try {
		device = physical_device.createDevice({
				.queueCreateInfoCount = cast<u32>(queue_create_infos.size()),
				.pQueueCreateInfos = queue_create_infos.data(),
				.enabledLayerCount = _context->enable_validation_layers ? cast<u32>(_context->validation_layers.size()) : 0,
				.ppEnabledLayerNames = _context->enable_validation_layers ? _context->validation_layers.data() : nullptr,
				.enabledExtensionCount = cast<u32>(_context->device_extensions.size()),
				.ppEnabledExtensionNames = _context->device_extensions.data(),
				.pEnabledFeatures = &_enabled_features,
		});
	} catch (const std::exception &err) {
		ERROR("Failed to create a logical device with "s + err.what());
		throw err;
	}
	INFO("Logical Device Created!");

	VmaAllocatorCreateInfo allocator_create_info = {
		.physicalDevice = *physical_device,
		.device = *device,
		.instance = *_context->instance,
	};

	auto result = cast<vk::Result>(vmaCreateAllocator(&allocator_create_info, &allocator));
	if (failed(result)) {
		ERROR("Memory allocator creation failed with "s + vk::to_string(result));
		throw std::runtime_error("Memory allocator creation failed with "s + vk::to_string(result));
	}
	VERBOSE("Memory Allocator Created");

	INFO(fmt::format("Created Device '{}' Successfully", _name));

	// Setup queues
	{
		u32 compute_idx = --unique_queue_families[queue_families.compute_idx];
		u32 transfer_idx = --unique_queue_families[queue_families.transfer_idx];
		u32 present_idx = --unique_queue_families[queue_families.present_idx];
		u32 graphics_idx = --unique_queue_families[queue_families.graphics_idx];

		queues.graphics = device.getQueue(queue_families.graphics_idx, graphics_idx);
		queues.present = device.getQueue(queue_families.present_idx, present_idx);
		queues.transfer = device.getQueue(queue_families.transfer_idx, transfer_idx);
		queues.compute = device.getQueue(queue_families.graphics_idx, compute_idx);
		INFO(fmt::format("Graphics Queue Index: ({}, {})", queue_families.graphics_idx, graphics_idx));
		INFO(fmt::format("Present Queue Index: ({}, {})", queue_families.present_idx, present_idx));
		INFO(fmt::format("Transfer Queue Index: ({}, {})", queue_families.transfer_idx, transfer_idx));
		INFO(fmt::format("Compute Queue Index: ({}, {})", queue_families.compute_idx, compute_idx));
	}

	// transfer_cmd_pool = device.createCommandPool({
	// 		.flags = vk::CommandPoolCreateFlagBits::eTransient,
	// 		.queueFamilyIndex = queue_families.transfer_idx,
	// });
	// if (failed(result)) {
	// 	allocator.destroy();
	// 	device.destroy();
	// 	return Err::make(fmt::format("Transfer command pool creation failed with {}" CODE_LOC, to_cstr(result)), result);
	// }
	// VERBOSE("Transfer Command Pool Created");
	//
	// vk::CommandPool graphics_cmd_pool;
	// tie(result, graphics_cmd_pool) = device.createCommandPool({
	// 		.flags = vk::CommandPoolCreateFlagBits::eTransient | vk::CommandPoolCreateFlagBits::eResetCommandBuffer,
	// 		.queueFamilyIndex = queue_families.graphics_idx,
	// });
	// if (failed(result)) {
	// 	device.destroyCommandPool(transfer_cmd_pool);
	// 	allocator.destroy();
	// 	device.destroy();
	// 	return Err::make(fmt::format("Graphics command pool creation failed with {}" CODE_LOC, to_cstr(result)), result);
	// }
	// VERBOSE("Graphics Command Pool Created");

	// transfer_cmd_pool,
	// graphics_cmd_pool
	// };

	set_name(_name);
	// final_device.set_object_name(transfer_cmd_pool, "Async transfer command pool");
	// final_device.set_object_name(graphics_cmd_pool, "Single use Graphics command pool");
}

Device::~Device() {
	if (allocator) {
		vmaDestroyAllocator(allocator);
	}
	INFO("Device '" + name + "' Destroyed");
}
//
// Res<vk::CommandBuffer> Device::alloc_temp_command_buffer(vk::CommandPool _pool) const {
// 	vk::CommandBuffer cmd;
// 	vk::CommandBufferAllocateInfo cmd_buf_alloc_info = {
// 		.commandPool = _pool,
// 		.level = vk::CommandBufferLevel::ePrimary,
// 		.commandBufferCount = 1,
// 	};
// 	if (const auto result = device.allocateCommandBuffers(&cmd_buf_alloc_info, &cmd); failed(result)) {
// 		return Err::make(fmt::format("Temp Command buffer allocation failed with {}" CODE_LOC, to_cstr(result)), result);
// 	}
// 	return cmd;
// }
//
// Res<SubmitTask<Buffer>> Device::upload_data(const Borrowed<Buffer> &_host_buffer, Buffer &&_staging_buffer) {
// 	ERROR_IF(!(_host_buffer->usage & vk::BufferUsageFlagBits::eTransferDst), fmt::format("Buffer {} is not a transfer dst. Use vk::BufferUsageFlagBits::eTransferDst during creation", _host_buffer->name.data()))
// 	ELSE_IF_ERROR(!(_staging_buffer.usage & vk::BufferUsageFlagBits::eTransferSrc), fmt::format("Buffer {} is not a transfer src. Use vk::BufferUsageFlagBits::eTransferSrc during creation", _staging_buffer.name.data()))
// 	ELSE_IF_WARN(_host_buffer->memory_usage != vma::MemoryUsage::eGpuOnly, fmt::format("Memory {} is not GPU only. Upload not required", _host_buffer->name.data()))
// 	ELSE_IF_WARN(_host_buffer->memory_usage != vma::MemoryUsage::eCpuOnly, fmt::format("Memory {} is not CPU only. Staging should ideally be a CPU only buffer", _staging_buffer.name.data()));
//
// 	vk::CommandBuffer cmd;
// 	vk::CommandBufferAllocateInfo allocate_info = {
// 		.commandPool = transfer_cmd_pool,
// 		.level = vk::CommandBufferLevel::ePrimary,
// 		.commandBufferCount = 1,
// 	};
//
// 	auto result = device.allocateCommandBuffers(&allocate_info, &cmd);
// 	if (failed(result)) {
// 		return Err::make(fmt::format("Transfer command pool allocation failed with {}" CODE_LOC, to_cstr(result)), result);
// 	}
// 	set_object_name(cmd, fmt::format("{} transfer command", _host_buffer->name.data()));
//
// 	result = cmd.begin({
// 			.flags = vk::CommandBufferUsageFlagBits::eOneTimeSubmit,
// 	});
// 	if (failed(result)) {
// 		return Err::make(fmt::format("Command buffer begin failed with {}" CODE_LOC, to_cstr(result)), result);
// 	}
//
// 	cmd.copyBuffer(_staging_buffer.buffer, _host_buffer->buffer, { { .srcOffset = 0, .dstOffset = 0, .size = cast<u32>(_staging_buffer.size) } });
// 	result = cmd.end();
// 	if (failed(result)) {
// 		return Err::make(fmt::format("Command buffer end failed with {}" CODE_LOC, to_cstr(result)), result);
// 	}
//
// 	return SubmitTask<Buffer>::create(borrow(this), std::move(_staging_buffer), queues.transfer, transfer_cmd_pool, { cmd });
// }
//
// Res<SubmitTask<Buffer>> Device::upload_data(const Borrowed<Buffer> &_host_buffer, const std::span<u8> &_data) {
// 	ERROR_IF(!(_host_buffer->usage & vk::BufferUsageFlagBits::eTransferDst), fmt::format("Buffer {} is not a transfer dst. Use vk::BufferUsageFlagBits::eTransferDst during creation", _host_buffer->name.data()))
// 	ELSE_IF_WARN(_host_buffer->memory_usage != vma::MemoryUsage::eGpuOnly, fmt::format("Memory {} is not GPU only. Upload not required", _host_buffer->name.data()));
//
// 	if (auto res = Buffer::create("_stage " + _host_buffer->name, borrow(this), _data.size(), vk::BufferUsageFlagBits::eTransferSrc, vma::MemoryUsage::eCpuOnly)) {
// 		return Err::make("Staging buffer creation failed", std::move(res.error()));
// 	} else {
// 		if (auto result = update_data(borrow(res.value()), _data)) {
// 			return upload_data(_host_buffer, std::move(res.value()));
// 		} else {
// 			return Err::make(std::move(result.error()));
// 		}
// 	}
// }
//
// Res<> Device::update_data(const Borrowed<Buffer> &_host_buffer, const std::span<u8> &_data) const {
// 	if (_host_buffer->memory_usage != vma::MemoryUsage::eCpuToGpu &&
// 			_host_buffer->memory_usage != vma::MemoryUsage::eCpuOnly) {
// 		return Err::make("Memory is not on CPU so mapping can't be done. Use upload_data" CODE_LOC);
// 	}
//
// 	auto [result, mapped_memory] = allocator.mapMemory(_host_buffer->allocation);
// 	if (failed(result)) {
// 		return Err::make(fmt::format("Memory mapping failed with {}" CODE_LOC, to_cstr(result)), result);
// 	}
// 	memcpy(mapped_memory, _data.data(), _data.size());
// 	allocator.unmapMemory(_host_buffer->allocation);
//
// 	return {};
// }

void Device::set_name(const std::string_view &_name) {
	VERBOSE(fmt::format("Device {} -> {}", name.data(), _name.data()));
	name = _name;
	set_object_name(*physical_device.device, fmt::format("{} GPU", _name.data()));
	set_object_name(*device, fmt::format("{} Device", _name.data()));
}