// =============================================
//  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{ std::move(_other.physical_device) },
		device{ std::exchange(_other.device, nullptr) },
		allocator{ std::exchange(_other.allocator, nullptr) },
		name{ std::move(_other.name) } {}

Device &Device::operator=(Device &&_other) noexcept {
	if (this == &_other)
		return *this;
	physical_device = std::move(_other.physical_device);
	device = std::exchange(_other.device, nullptr);
	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, std::set<QueueAllocation> &&_enabled_queues) :
		physical_device{ _physical_device_info },
		name{ _name } {
	const auto &physical_device = _physical_device_info.device;

	// Logical Device
	std::vector queue_priority(_enabled_queues.size(), 1.0f);
	std::vector<vk::DeviceQueueCreateInfo> queue_create_infos;
	queue_create_infos.reserve(_enabled_queues.size());
	for (auto &[family_, count_] : _enabled_queues) {
		queue_create_infos.push_back({
				.queueFamilyIndex = family_,
				.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;
	}
	INFO("Logical Device Created!");

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

	allocator = nullptr;
	const 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));

	set_name(_name);
}

Device::~Device() {
	if (allocator) {
		vmaDestroyAllocator(allocator);
	}
	INFO("Device '" + name + "' Destroyed");
}

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