#include "iwa/instance.hpp"

#include <mijin/detect.hpp>
#include <mijin/debug/assert.hpp>
#include "iwa/addon.hpp"
#include "iwa/log.hpp"
#include "iwa/window.hpp"

#include <vulkan/vulkan_from_string.inl>

VULKAN_HPP_DEFAULT_DISPATCH_LOADER_DYNAMIC_STORAGE

namespace iwa
{
IWA_REGISTER_CLASS(Instance)

namespace
{
void buildDefaultInstanceExtensionList(std::vector<ExtensionInfo>& outExtensions) noexcept
{
#if !defined(KAZAN_RELEASE)
    outExtensions.push_back({.name = VK_EXT_DEBUG_UTILS_EXTENSION_NAME, .required = false});
#endif
    // if (!getEngineOptions().headless)
    {
        outExtensions.push_back({.name = VK_KHR_SURFACE_EXTENSION_NAME, .required = true});
#if MIJIN_TARGET_OS == MIJIN_OS_LINUX
        // There are two possible APIs for X11 surfaces (SDL supports both)
        // just try to enable both and hope SDL will be happy.
        // Both are optional, in case one is not supported, but in the end
        // one of them will be needed.
        outExtensions.push_back({.name = "VK_KHR_xcb_surface", .required = false});
        outExtensions.push_back({.name = "VK_KHR_xlib_surface", .required = false});
#elif MIJIN_TARGET_OS == MIJIN_OS_WINDOWS
        outExtensions.push_back({.name = "VK_KHR_win32_surface", .required = true});
#endif
    }
}

void buildDefaultInstanceLayerList(std::vector<LayerInfo>& outLayers) noexcept
{
#if !defined(KAZAN_RELEASE)
    outLayers.push_back({.name = "VK_LAYER_KHRONOS_validation", .required = false});
#else
    (void) outLayers;
#endif
}

std::vector<const char*> checkInstanceExtensions(std::vector<ExtensionInfo>& extensions)
{
    std::vector<vk::ExtensionProperties> properties = vk::enumerateInstanceExtensionProperties(nullptr);

    auto isExtensionSupported = [&properties](const char* extension)
    {
        for (const vk::ExtensionProperties& props : properties)
        {
            if (std::strncmp(props.extensionName, extension, VK_MAX_EXTENSION_NAME_SIZE) == 0)
            {
                return true;
            }
        }
        return false;
    };

    std::vector<const char*> enabledExtensions;
    for (ExtensionInfo& extInfo : extensions)
    {
        const bool supported = isExtensionSupported(extInfo.name);
        if (!supported && extInfo.required)
        {
            logAndDie("Required Vulkan instance extension not supported: {}.", extInfo.name);
        }
        if (supported) {
            enabledExtensions.push_back(extInfo.name);
        }
        extInfo.enabled = supported;
    }
    return enabledExtensions;
}

std::vector<const char*> checkInstanceLayers(std::vector<LayerInfo>& layers)
{
    std::vector<const char*> enabledLayers;
#if !defined(KAZAN_RELEASE)
    std::vector<vk::LayerProperties> properties = vk::enumerateInstanceLayerProperties();

    auto isLayerSupported = [&properties](const char* extension)
    {
        for (const vk::LayerProperties& props : properties)
        {
            if (std::strncmp(props.layerName, extension, VK_MAX_EXTENSION_NAME_SIZE) == 0)
            {
                return true;
            }
        }
        return false;
    };

    for (LayerInfo& layerInfo : layers)
    {
        const bool supported = isLayerSupported(layerInfo.name);
        if (!supported && layerInfo.required)
        {
            logAndDie("Required Vulkan instance layer not supported: {}.", layerInfo.name);
        }
        if (supported) {
            enabledLayers.push_back(layerInfo.name);
        }
        layerInfo.enabled = supported;
    }
#endif
    return enabledLayers;
}
vk::Bool32 VKAPI_PTR vulkanMessengerCallback(
        VkDebugUtilsMessageSeverityFlagBitsEXT severity,
        VkDebugUtilsMessageTypeFlagsEXT /* types */,
        const VkDebugUtilsMessengerCallbackDataEXT* pCallbackData,
        void* /* pUserData */
)
{
    switch (pCallbackData->messageIdNumber)
    {
    case 1303270965: // We use general layout to copy texture to mips. I know it's not optimal, but it's allowed and
                     // works for me ¯\_(ツ)_/¯
    case 148949623:  // The spec says that a) you can't create descriptor sets from layouts with the push-descriptor bit
                     // set and b) you have to bind a descriptor set when using said set layout in vkCmdPushDescriptorSetKHR.
                     // So you have to use a set, but can't create it... Something is off.
                     // Apart from that c) it also says that dstSet is ignored when using vkCmdPushDescriptorSetKHR, so
                     // maybe the validation layers are just wrong.
                     // a) https://registry.khronos.org/vulkan/specs/1.3-extensions/man/html/VkDescriptorSetAllocateInfo.html#VUID-VkDescriptorSetAllocateInfo-pSetLayouts-00308
                     // b) https://registry.khronos.org/vulkan/specs/1.3-extensions/man/html/VkWriteDescriptorSet.html#VUID-VkWriteDescriptorSet-dstSet-00320
                     // c) https://registry.khronos.org/vulkan/specs/1.3-extensions/man/html/vkCmdPushDescriptorSetKHR.html
        return VK_FALSE;
    }
    logMsg("VK> {}", pCallbackData->pMessage);
    if (/* g_breakOnVulkanError && */ severity == VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT)
    {
        if (const mijin::Result<mijin::Stacktrace> trace = mijin::captureStacktrace(0); trace.isSuccess())
        {
            logMsg("{}", *trace);
        }
        MIJIN_TRAP();
    }

    return VK_FALSE;
}

std::pair<std::uint32_t, std::uint32_t> detectQueueFamilies(vk::PhysicalDevice physicalDevice, vk::SurfaceKHR dummySurface)
{
    std::vector<vk::QueueFamilyProperties> queueFamilyProperties = physicalDevice.getQueueFamilyProperties();

    for (std::uint32_t idx = 0; idx < queueFamilyProperties.size(); ++idx)
    {
        const vk::QueueFamilyProperties& properties = queueFamilyProperties[idx];
        const vk::Bool32 surfaceSupport = dummySurface ? physicalDevice.getSurfaceSupportKHR(idx, dummySurface) : VK_TRUE; // no dummy surface -> no surface support needed
        if (surfaceSupport
            && (properties.queueFlags & vk::QueueFlagBits::eGraphics)
            && (properties.queueFlags & vk::QueueFlagBits::eCompute))
        {
            return std::make_pair(idx, idx);
        }
    }

    return std::make_pair(std::numeric_limits<std::uint32_t>::max(), std::numeric_limits<std::uint32_t>::max());
}

PhysicalDeviceInfo getPhysicalDeviceInfo(vk::PhysicalDevice physicalDevice, vk::SurfaceKHR dummySurface)
{
    PhysicalDeviceInfo deviceInfo;

    deviceInfo.device = physicalDevice;
    deviceInfo.memoryProperties = physicalDevice.getMemoryProperties();
    deviceInfo.extensions = physicalDevice.enumerateDeviceExtensionProperties();

    auto isExtensionSupported = [&](const char* extension)
    {
        for (const vk::ExtensionProperties& props : deviceInfo.extensions)
        {
            if (std::strncmp(props.extensionName, extension, VK_MAX_EXTENSION_NAME_SIZE) == 0) {
                return true;
            }
        }
        return false;
    };
    const bool rayTracingSupported = isExtensionSupported(VK_KHR_ACCELERATION_STRUCTURE_EXTENSION_NAME)
                                  && isExtensionSupported(VK_KHR_RAY_TRACING_PIPELINE_EXTENSION_NAME)
                                  && isExtensionSupported(VK_KHR_DEFERRED_HOST_OPERATIONS_EXTENSION_NAME);
    const bool meshShadersSupported = isExtensionSupported(VK_EXT_MESH_SHADER_EXTENSION_NAME);

    vk::PhysicalDeviceProperties2 deviceProperties = {
            .pNext = &deviceInfo.rayTracingProperties
    };
    physicalDevice.getProperties2(&deviceProperties);
    deviceInfo.properties = deviceProperties.properties;

#define ADD_TO_NEXT_CHAIN(struct_name) \
    *ppNext = &(struct_name); ppNext = &(struct_name).pNext
    vk::PhysicalDeviceFeatures2 deviceFeatures;
    void** ppNext = &deviceFeatures.pNext;
    ADD_TO_NEXT_CHAIN(deviceInfo.vulkan11Features);
    ADD_TO_NEXT_CHAIN(deviceInfo.vulkan12Features);
    ADD_TO_NEXT_CHAIN(deviceInfo.vulkan13Features);
    if (rayTracingSupported)
    {
        ADD_TO_NEXT_CHAIN(deviceInfo.accelerationStructureFeatures);
        ADD_TO_NEXT_CHAIN(deviceInfo.rayTracingPipelineFeatures);
    }
    if (meshShadersSupported) {
        ADD_TO_NEXT_CHAIN(deviceInfo.meshShaderFeatures);
    }

#undef ADD_TO_NEXT_CHAIN
    physicalDevice.getFeatures2(&deviceFeatures);
    deviceInfo.features = deviceFeatures.features;

    // fill the renderer features
    deviceInfo.availableFeatures.rayTracing = rayTracingSupported;
    deviceInfo.availableFeatures.meshShaders = meshShadersSupported && deviceInfo.meshShaderFeatures.meshShader;

    deviceInfo.surfaceCapabilities = physicalDevice.getSurfaceCapabilitiesKHR(dummySurface);
    std::tie(deviceInfo.graphicsQueueFamily, deviceInfo.computeQueueFamily) = detectQueueFamilies(physicalDevice, dummySurface);

    return deviceInfo;
}

std::vector<PhysicalDeviceInfo> getPhysicalDeviceInfos(Instance& instance)
{
    const std::vector<vk::PhysicalDevice> devices = instance.getVkHandle().enumeratePhysicalDevices();
    std::vector<PhysicalDeviceInfo> deviceInfos;
    deviceInfos.reserve(devices.size());

    const ObjectPtr<Window> dummyWindow = instance.createWindow({.flags={.hidden=true}});
    const vk::SurfaceKHR dummySurface = dummyWindow->getVkSurface();

    for (const vk::PhysicalDevice device : devices)
    {
        deviceInfos.push_back(getPhysicalDeviceInfo(device, dummySurface));
    }

    return deviceInfos;
}
} // namespace

Instance::Instance(InstanceCreationArgs args)
    : super_t(nullptr),
      mExtensions(std::move(args.extensions)),
      mLayers(std::move(args.layers))
{
    setMainThread();

    const AddonInitArgs addonInitArgs = {
        .instance = getPointer(),
        .instanceCreationArgs = args
    };
    for (Addon* addon : getAddons())
    {
        addon->init(addonInitArgs);
    }

    VULKAN_HPP_DEFAULT_DISPATCHER.init();

    if (!args.flags.noDefaultExtensions)
    {
        buildDefaultInstanceExtensionList(mExtensions);
    }
    if (!args.flags.noDefaultLayers)
    {
        buildDefaultInstanceLayerList(mLayers);
    }

    const std::vector<const char*> enabledLayerNames = checkInstanceLayers(mLayers);
    const std::vector<const char*> enabledExtensionNames = checkInstanceExtensions(mExtensions);
    const vk::InstanceCreateInfo createInfo =
    {
        .pApplicationInfo = &args.applicationInfo,
        .enabledLayerCount = static_cast<std::uint32_t>(enabledLayerNames.size()),
        .ppEnabledLayerNames = enabledLayerNames.data(),
        .enabledExtensionCount = static_cast<std::uint32_t>(enabledExtensionNames.size()),
        .ppEnabledExtensionNames = enabledExtensionNames.data()
    };
    mHandle = vk::createInstance(createInfo);
    VULKAN_HPP_DEFAULT_DISPATCHER.init(mHandle);

    // dump info about enabled extensions and layers
    logMsg("Enabled Vulkan instance extensions: {}", enabledExtensionNames);
    logMsg("Enabled Vulkan instance layers: {}", enabledLayerNames);

    // also dump the actual instance version
    const std::uint32_t apiVersion = vk::enumerateInstanceVersion();
    logMsg("Vulkan instance version: {}.{}.{}.\n",
       VK_API_VERSION_MAJOR(apiVersion),
       VK_API_VERSION_MINOR(apiVersion),
       VK_API_VERSION_PATCH(apiVersion)
    );

    if (isExtensionEnabled(VK_EXT_DEBUG_UTILS_EXTENSION_NAME))
    {
        // create debug messenger
        const vk::DebugUtilsMessengerCreateInfoEXT dumCreateInfo =
        {
            .messageSeverity = vk::DebugUtilsMessageSeverityFlagBitsEXT::eWarning | vk::DebugUtilsMessageSeverityFlagBitsEXT::eError,
            .messageType = vk::DebugUtilsMessageTypeFlagBitsEXT::eGeneral | vk::DebugUtilsMessageTypeFlagBitsEXT::ePerformance | vk::DebugUtilsMessageTypeFlagBitsEXT::eValidation,
            .pfnUserCallback = &vulkanMessengerCallback
        };
        mDebugMessenger = mHandle.createDebugUtilsMessengerEXT(dumCreateInfo);
    }

    mPhysicalDevices = getPhysicalDeviceInfos(*this);
    mWorkerTaskLoop.start(10);
}

Instance::~Instance() noexcept
{
    for (Addon* addon : getAddons())
    {
        addon->cleanup();
    }
    while (!mDeleteQueue.empty())
    {
        runDeleters(true);
    }

    if (mDebugMessenger)
    {
        mHandle.destroyDebugUtilsMessengerEXT(mDebugMessenger);
    }
    mHandle.destroy();
}

bool Instance::isExtensionEnabled(const char* name) const noexcept
{
    for (const ExtensionInfo& extInfo : mExtensions)
    {
        if (std::strcmp(extInfo.name, name) == 0)
        {
            return extInfo.enabled;
        }
    }
    return false;
}

bool Instance::isLayerEnabled(const char* name) const noexcept
{
    for (const LayerInfo& layerInfo : mLayers)
    {
        if (std::strcmp(layerInfo.name, name) == 0)
        {
            return layerInfo.enabled;
        }
    }
    return false;
}

ObjectPtr<Window> Instance::createWindow(const WindowCreationArgs& args)
{
    return createChild<Window>(args);
}

ObjectPtr<Device> Instance::createDevice(DeviceCreationArgs args)
{
    return createChild<Device>(std::move(args));
}

void Instance::queueDelete(deleter_t deleter) noexcept
{
    MIJIN_ASSERT(deleter, "Don't pass an empty function into queueDelete()!");
    if (!deleter)
    {
        return;
    }

    runOnMainThread([deleter = std::move(deleter), this]() mutable
    {
        static const int DELETE_DELAY_FRAMES = 5; // TODO: ?
        mDeleteQueue.push_back({
           .deleter = std::move(deleter),
           .remainingFrames = DELETE_DELAY_FRAMES
        });
    });
}

void Instance::tickDeleteQueue()
{
    for (DeleteQueueEntry& entry : mDeleteQueue)
    {
        --entry.remainingFrames;
    }
    runDeleters();
}

void Instance::setMainThread()
{
    mMainThread = std::this_thread::get_id();
}

void Instance::requestQuit() noexcept
{
    if (mQuitRequested)
    {
        return;
    }
    
    mQuitRequested = true;
    quitRequested.emit();

    mMainTaskLoop.addTask([](Instance* instance) -> mijin::Task<>
    {
        for (int tick = 0; tick < 100; ++tick)
        {
            if (instance->getMainTaskLoop().getNumTasks() == 1)
            {
                // just us left, bye bye!
                co_return;
            }
            co_await mijin::c_suspend();
        }
        const std::size_t activeTasks = instance->getMainTaskLoop().getActiveTasks() - 1; // substract yourself
        if (activeTasks > 0)
        {
            logMsg("{} tasks did not finish when shutting down!", activeTasks - 1);
#if MIJIN_COROUTINE_ENABLE_DEBUG_INFO
            const std::vector<mijin::TaskHandle> tasks = instance->getMainTaskLoop().getAllTasks();
            for (auto [index, handle] : mijin::enumerate(tasks))
            {
                if (handle == mijin::getCurrentTask())
                {
                    continue;
                }
                const mijin::Optional<mijin::Stacktrace> stack = handle.getCreationStack();
                if (!stack.empty())
                {
                    logMsg("Task {} creation stack:\n{}", index, *stack);
                }
            }
#endif // MIJIN_COROUTINE_ENABLE_DEBUG_INFO
            MIJIN_TRAP();
        }
        else
        {
            const std::size_t totalTasks = instance->getMainTaskLoop().getNumTasks() - 1;
            logMsg("{} tasks still waiting on program exit.", totalTasks);
        }
        instance->getMainTaskLoop().cancelAllTasks();
        co_return;
    }(this));
}

void Instance::runDeleters(bool runAll)
{
    for (auto it = mDeleteQueue.begin(); it != mDeleteQueue.end();)
    {
        if (runAll || it->remainingFrames <= 0)
        {
            it->deleter();
            it = mDeleteQueue.erase(it);
        }
        else
        {
            ++it;
        }
    }
}
}