#include "VKDevice.hpp"
#include "MayaFlux/Core/GlobalGraphicsInfo.hpp"
#include "MayaFlux/Journal/Archivist.hpp"

#include "set"

namespace MayaFlux::Core {

VKDevice::~VKDevice()
{
    cleanup();
}

VKDevice::VKDevice(VKDevice&& other) noexcept
    : m_physical_device(other.m_physical_device)
    , m_logical_device(other.m_logical_device)
    , m_graphics_queue(other.m_graphics_queue)
    , m_compute_queue(other.m_compute_queue)
    , m_transfer_queue(other.m_transfer_queue)
    , m_queue_families(other.m_queue_families)
{
    other.m_physical_device = VK_NULL_HANDLE;
    other.m_logical_device = VK_NULL_HANDLE;
    other.m_graphics_queue = VK_NULL_HANDLE;
    other.m_compute_queue = VK_NULL_HANDLE;
    other.m_transfer_queue = VK_NULL_HANDLE;
}

VKDevice& VKDevice::operator=(VKDevice&& other) noexcept
{
    if (this != &other) {
        cleanup();
        m_physical_device = other.m_physical_device;
        m_logical_device = other.m_logical_device;
        m_graphics_queue = other.m_graphics_queue;
        m_compute_queue = other.m_compute_queue;
        m_transfer_queue = other.m_transfer_queue;
        m_queue_families = other.m_queue_families;

        other.m_physical_device = VK_NULL_HANDLE;
        other.m_logical_device = VK_NULL_HANDLE;
        other.m_graphics_queue = VK_NULL_HANDLE;
        other.m_compute_queue = VK_NULL_HANDLE;
        other.m_transfer_queue = VK_NULL_HANDLE;
    }
    return *this;
}

bool VKDevice::initialize(vk::Instance instance, vk::SurfaceKHR /*temp_surface*/, const GraphicsBackendInfo& backend_info)
{
    if (!pick_physical_device(instance, nullptr)) {
        return false;
    }

    return create_logical_device(instance, backend_info);
}

void VKDevice::cleanup()
{
    if (m_logical_device) {
        m_logical_device.destroy();
        m_logical_device = nullptr;
        MF_INFO(Journal::Component::Core, Journal::Context::GraphicsBackend, "Vulkan logical device destroyed.");
    }
    m_physical_device = nullptr;
    m_graphics_queue = nullptr;
    m_compute_queue = nullptr;
    m_transfer_queue = nullptr;
    m_queue_families = {};
}

bool VKDevice::pick_physical_device(vk::Instance instance, vk::SurfaceKHR /*temp_surface*/)
{
    vk::Instance vk_instance(instance);
    auto devices = vk_instance.enumeratePhysicalDevices();

    if (devices.empty()) {
        error<std::runtime_error>(Journal::Component::Core, Journal::Context::GraphicsBackend,
            std::source_location::current(),
            "Failed to find GPUs with Vulkan support!");
    }

    for (const auto& device : devices) {
        QueueFamilyIndices indices = find_queue_families(device, nullptr);

        if (indices.graphics_family.has_value()) {
            auto available_extensions = device.enumerateDeviceExtensionProperties();
            bool supports_swapchain = false;
            bool supports_mesh_shader = false;

            for (const auto& ext : available_extensions) {
                if (strcmp(ext.extensionName, VK_KHR_SWAPCHAIN_EXTENSION_NAME) == 0) {
                    supports_swapchain = true;
                }
                if (strcmp(ext.extensionName, VK_EXT_MESH_SHADER_EXTENSION_NAME) == 0) {
                    supports_mesh_shader = true;
                }
            }

            if (!supports_swapchain) {
                MF_WARN(Journal::Component::Core, Journal::Context::GraphicsBackend,
                    "Physical device {} does not support VK_KHR_swapchain - skipping",
                    device.getProperties().deviceName.data());
                continue;
            }

            if (!supports_mesh_shader) {
                MF_WARN(Journal::Component::Core, Journal::Context::GraphicsBackend,
                    "Physical device {} does not support VK_EXT_mesh_shader - "
                    "mesh shading features will be unavailable",
                    device.getProperties().deviceName.data());
            } else {
                vk::PhysicalDeviceMeshShaderFeaturesEXT mesh_features;
                vk::PhysicalDeviceFeatures2 temp_features;
                temp_features.pNext = &mesh_features;
                device.getFeatures2(&temp_features);

                if (mesh_features.meshShader == VK_FALSE || mesh_features.taskShader == VK_FALSE) {
                    MF_INFO(Journal::Component::Core, Journal::Context::GraphicsBackend,
                        "Physical device {} supports VK_EXT_mesh_shader extension but features disabled",
                        device.getProperties().deviceName.data());
                    supports_mesh_shader = false;
                }
            }

            m_physical_device = device;
            m_queue_families = indices;
            m_supports_mesh_shaders = supports_mesh_shader;

            vk::PhysicalDeviceProperties props = device.getProperties();
            MF_LOG(Journal::Component::Core, Journal::Context::GraphicsBackend,
                "Selected GPU: {}", props.deviceName.data());
            return true;
        }
    }

    error<std::runtime_error>(Journal::Component::Core, Journal::Context::GraphicsBackend,
        std::source_location::current(),
        "Failed to find a suitable GPU!");
    return false;
}

QueueFamilyIndices VKDevice::find_queue_families(vk::PhysicalDevice device, vk::SurfaceKHR surface)
{
    QueueFamilyIndices indices;
    auto queue_families = device.getQueueFamilyProperties();

    int i = 0;
    for (const auto& queue_family : queue_families) {
        if (queue_family.queueCount > 0 && queue_family.queueFlags & vk::QueueFlagBits::eGraphics) {
            indices.graphics_family = i;
        }

        if (queue_family.queueCount > 0 && queue_family.queueFlags & vk::QueueFlagBits::eCompute && !(queue_family.queueFlags & vk::QueueFlagBits::eGraphics)) {
            indices.compute_family = i;
        }

        if (queue_family.queueCount > 0 && queue_family.queueFlags & vk::QueueFlagBits::eTransfer && !(queue_family.queueFlags & vk::QueueFlagBits::eGraphics) && !(queue_family.queueFlags & vk::QueueFlagBits::eCompute)) {
            indices.transfer_family = i;
        }

        if (surface && queue_family.queueCount > 0) {
            vk::Bool32 presentSupport = device.getSurfaceSupportKHR(i, surface);
            if (presentSupport) {
                indices.present_family = i;
                MF_INFO(Journal::Component::Core, Journal::Context::GraphicsBackend,
                    "Found presentation support in queue family {}", i);
            }
        }

        i++;
    }

    if (indices.graphics_family.has_value()) {
        if (!indices.compute_family.has_value()) {
            indices.compute_family = indices.graphics_family;
        }
        if (!indices.transfer_family.has_value()) {
            indices.transfer_family = indices.graphics_family;
        }
    }

    return indices;
}

bool VKDevice::update_presentation_queue(vk::SurfaceKHR surface)
{
    if (!surface) {
        MF_WARN(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "No surface provided for presentation support check");
        return false;
    }

    auto queue_families = m_physical_device.getQueueFamilyProperties();

    for (uint32_t i = 0; i < queue_families.size(); i++) {
        if (queue_families[i].queueCount > 0) {
            vk::Bool32 presentSupport = m_physical_device.getSurfaceSupportKHR(i, surface);
            if (presentSupport) {
                if (i == m_queue_families.graphics_family.value()) {
                    m_queue_families.present_family = i;
                    m_presentation_initialized = true;
                    MF_INFO(Journal::Component::Core, Journal::Context::GraphicsBackend,
                        "Graphics queue family {} supports presentation", i);
                    return true;
                }
            }
        }
    }

    for (uint32_t i = 0; i < queue_families.size(); i++) {
        if (queue_families[i].queueCount > 0) {
            vk::Bool32 presentSupport = m_physical_device.getSurfaceSupportKHR(i, surface);
            if (presentSupport) {
                m_queue_families.present_family = i;
                m_presentation_initialized = true;
                MF_INFO(Journal::Component::Core, Journal::Context::GraphicsBackend,
                    "Found presentation support in queue family {}", i);
                return true;
            }
        }
    }

    MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
        "No queue family with presentation support found!");
    return false;
}

void VKDevice::query_supported_extensions()
{
    std::vector<vk::ExtensionProperties> availableExtensions = m_physical_device.enumerateDeviceExtensionProperties();

    MF_LOG(Journal::Component::Core, Journal::Context::GraphicsBackend, "Available physical device extensions:");
    for (const auto& extension : availableExtensions) {
        std::cout << "\t- " << extension.extensionName << " (Version: " << extension.specVersion << ")\n";
    }
    MF_LOG(Journal::Component::Core, Journal::Context::GraphicsBackend, "End of list.");
}

bool VKDevice::create_logical_device(vk::Instance /*instance*/, const GraphicsBackendInfo& backend_info)
{
    if (!m_queue_families.graphics_family.has_value()) {
        error<std::runtime_error>(Journal::Component::Core, Journal::Context::GraphicsBackend,
            std::source_location::current(),
            "No graphics queue family found!");
    }

    std::set<uint32_t> unique_queue_families;
    unique_queue_families.insert(m_queue_families.graphics_family.value());

    if (m_queue_families.compute_family.has_value()) {
        unique_queue_families.insert(m_queue_families.compute_family.value());
    }

    if (m_queue_families.transfer_family.has_value() && m_queue_families.transfer_family != m_queue_families.graphics_family) {
        unique_queue_families.insert(m_queue_families.transfer_family.value());
    }

    std::vector<vk::DeviceQueueCreateInfo> queue_create_infos;
    float queue_priority = 1.0F;

    for (uint32_t queue_family : unique_queue_families) {
        vk::DeviceQueueCreateInfo queue_create_info {};
        queue_create_info.queueFamilyIndex = queue_family;
        queue_create_info.queueCount = 1;
        queue_create_info.pQueuePriorities = &queue_priority;
        queue_create_infos.push_back(queue_create_info);
    }

    vk::PhysicalDeviceFeatures device_features {};
    device_features.samplerAnisotropy = backend_info.required_features.sampler_anisotropy;
    device_features.geometryShader = backend_info.required_features.geometry_shaders;
    device_features.tessellationShader = backend_info.required_features.tessellation_shaders;
    device_features.multiViewport = backend_info.required_features.multi_viewport;
    device_features.fillModeNonSolid = backend_info.required_features.fill_mode_non_solid;

    std::vector<const char*> device_extensions = { VK_KHR_SWAPCHAIN_EXTENSION_NAME };

#ifdef MAYAFLUX_PLATFORM_MACOS
    auto available_exts = m_physical_device.enumerateDeviceExtensionProperties();
    if (std::ranges::any_of(available_exts, [](const auto& ext) {
            return strcmp(ext.extensionName, "VK_KHR_portability_subset") == 0;
        })) {
        device_extensions.push_back("VK_KHR_portability_subset");
    }

    auto feature_chain = vk::StructureChain {
        vk::PhysicalDeviceFeatures2 {},
        vk::PhysicalDeviceVulkan13Features {},
        vk::PhysicalDeviceVulkan12Features {}
    };

#else
    if (m_supports_mesh_shaders) {
        device_extensions.push_back(VK_EXT_MESH_SHADER_EXTENSION_NAME);
    }

    auto feature_chain = vk::StructureChain {
        vk::PhysicalDeviceFeatures2 {},
        vk::PhysicalDeviceVulkan13Features {},
        vk::PhysicalDeviceVulkan12Features {},
        vk::PhysicalDeviceMeshShaderFeaturesEXT {}
    };

    if (!m_supports_mesh_shaders) {
        feature_chain.unlink<vk::PhysicalDeviceMeshShaderFeaturesEXT>();
    } else {
        feature_chain.get<vk::PhysicalDeviceMeshShaderFeaturesEXT>().taskShader = VK_TRUE;
        feature_chain.get<vk::PhysicalDeviceMeshShaderFeaturesEXT>().meshShader = VK_TRUE;
    }
#endif

    feature_chain.get<vk::PhysicalDeviceFeatures2>().features = device_features;
    feature_chain.get<vk::PhysicalDeviceVulkan13Features>().dynamicRendering = VK_TRUE;
    feature_chain.get<vk::PhysicalDeviceVulkan13Features>().synchronization2 = VK_TRUE;
    feature_chain.get<vk::PhysicalDeviceVulkan12Features>().bufferDeviceAddress = VK_TRUE;

    for (const auto& ext : backend_info.required_extensions) {
        device_extensions.push_back(ext.c_str());
    }

    vk::DeviceCreateInfo create_info {};
    create_info.queueCreateInfoCount = static_cast<uint32_t>(queue_create_infos.size());
    create_info.pQueueCreateInfos = queue_create_infos.data();
    create_info.pNext = &feature_chain.get<vk::PhysicalDeviceFeatures2>();
    create_info.enabledExtensionCount = static_cast<uint32_t>(device_extensions.size());
    create_info.ppEnabledExtensionNames = device_extensions.data();

    try {
        m_logical_device = m_physical_device.createDevice(create_info);
        VULKAN_HPP_DEFAULT_DISPATCHER.init(m_logical_device);

    } catch (const std::exception& e) {
        error_rethrow(Journal::Component::Core, Journal::Context::GraphicsBackend,
            std::source_location::current(),
            "Failed to create logical device: {}", e.what());
    }

    m_graphics_queue = m_logical_device.getQueue(m_queue_families.graphics_family.value(), 0);

    if (backend_info.enable_compute_queue && m_queue_families.compute_family.has_value()) {
        m_compute_queue = m_logical_device.getQueue(m_queue_families.compute_family.value(), 0);
    } else {
        m_compute_queue = m_graphics_queue;
    }

    if (backend_info.enable_transfer_queue && m_queue_families.transfer_family.has_value()) {
        m_transfer_queue = m_logical_device.getQueue(m_queue_families.transfer_family.value(), 0);
    } else {
        m_transfer_queue = m_graphics_queue;
    }

    return true;
}

}
