#include "iwa/device.hpp"

#include "iwa/log.hpp"
#include "iwa/instance.hpp"

namespace iwa
{
namespace
{
void buildDefaultDeviceExtensionList(std::vector<ExtensionInfo>& outExtensions) noexcept
{
    outExtensions.push_back({.name = VK_KHR_SWAPCHAIN_EXTENSION_NAME, .required = true});
}

bool checkQueueFamilies(const PhysicalDeviceInfo& deviceInfo, [[maybe_unused]] int& score)
{
    return deviceInfo.graphicsQueueFamily != std::numeric_limits<std::uint32_t>::max()
           && deviceInfo.computeQueueFamily != std::numeric_limits<std::uint32_t>::max();
}

bool checkDeviceExtensions(const std::vector<ExtensionInfo>& extensions, const PhysicalDeviceInfo& deviceInfo, int& score)
{
    auto isExtensionSupported = [&deviceInfo](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;
    };

    for (const ExtensionInfo& extInfo : extensions)
    {
        if (isExtensionSupported(extInfo.name))
        {
            score += 10;
        }
        else if (extInfo.required)
        {
            logVerbose("Vulkan device {} not supported as it is missing required extension {}.", deviceInfo.properties.deviceName, extInfo.name);
            return false;
        }
    }

    return true;
}

int scorePhysicalDevice(
        const PhysicalDeviceCriteria& deviceCriteria,
        const std::vector<ExtensionInfo>& extensions,
        const PhysicalDeviceInfo& deviceInfo)
{
    int score = 0;

    switch (deviceInfo.properties.deviceType)
    {
        case vk::PhysicalDeviceType::eDiscreteGpu:
            score = 1000;
            break;
        case vk::PhysicalDeviceType::eIntegratedGpu:
            score = 100;
            break;
        default: break;
    }

#define CHECK_FEATURE_STRUCT(infoName, requiredName) \
    for (auto feature : deviceCriteria.requiredName) \
    { \
        if (!(deviceInfo.infoName.*feature)) \
        { \
            return -1; \
        } \
    }
    CHECK_FEATURE_STRUCT(features, requiredFeatures)
    CHECK_FEATURE_STRUCT(vulkan11Features, requiredVulkan11Features)
    CHECK_FEATURE_STRUCT(vulkan12Features, requiredVulkan12Features)
    CHECK_FEATURE_STRUCT(vulkan13Features, requiredVulkan13Features)
    CHECK_FEATURE_STRUCT(accelerationStructureFeatures, requiredAccelerationStructureFeatures)
    CHECK_FEATURE_STRUCT(rayTracingPipelineFeatures, requiredRayTracingPipelineFeatures)
    CHECK_FEATURE_STRUCT(meshShaderFeatures, requredMeshShaderFeatures)

#undef CHECK_FEATURE_STRUCT

    if (!checkQueueFamilies(deviceInfo, score))
    {
        logVerbose("Vulkan device {} not supported as no suitable queue family configuration could be found.", deviceInfo.properties.deviceName);
        return -1;
    }

    if (!checkDeviceExtensions(extensions, deviceInfo, score))
    {
        return -1;
    }

    return score;
}

std::size_t findBestPhysicalDevice(
        PhysicalDeviceCriteria& deviceCriteria,
        const std::vector<ExtensionInfo>& extensions,
        const std::vector<PhysicalDeviceInfo>& deviceInfos)
{
    int bestScore = -1;
    std::size_t bestIdx = 0;

    for (std::size_t idx = 0; idx < deviceInfos.size(); ++idx)
    {
        const int score = scorePhysicalDevice(deviceCriteria, extensions, deviceInfos[idx]);
        if (score > bestScore)
        {
            bestScore = score;
            bestIdx = idx;
        }
    }

    if (bestScore < 0)
    {
        logAndDie("Could not find a suitable Vulkan device!");
    }

    return bestIdx;
}

std::vector<const char*> buildDeviceExtensionNameList(const PhysicalDeviceInfo& deviceInfo, std::vector<ExtensionInfo>& extensions)
{
    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;
    };

    std::vector<const char*> enabledExtensions;
    for (ExtensionInfo& extInfo : extensions)
    {
        const bool supported = isExtensionSupported(extInfo.name);
        if (!supported && extInfo.required)
        {
            // this shouldn't be possible, we already checked when scoring the device
            logAndDie("Required Vulkan device extension not supported: {}.", extInfo.name);
        }
        if (supported) {
            enabledExtensions.push_back(extInfo.name);
        }
        extInfo.enabled = supported;
    }
    return enabledExtensions;
}

std::vector<vk::DeviceQueueCreateInfo> buildQueueCreateInfoList(const PhysicalDeviceInfo& deviceInfo)
{
    static const float QUEUE_PRIORITY_ONE = 1.0f;
    std::vector<vk::DeviceQueueCreateInfo> createInfos;

    createInfos.emplace_back()
            .setQueueFamilyIndex(deviceInfo.graphicsQueueFamily)
            .setQueueCount(1)
            .setPQueuePriorities(&QUEUE_PRIORITY_ONE);

    return createInfos;
}
} // namespace

Device::Device(ObjectPtr<Instance> owner, DeviceCreationArgs args)
    : super_t(std::move(owner)),
      mExtensions(std::move(args.extensions))
{
    if (!args.flags.noDefaultExtensions)
    {
        buildDefaultDeviceExtensionList(mExtensions);
    }

    const std::vector<PhysicalDeviceInfo>& physicalDevices = getOwner()->getPhysicalDevices();
    const std::size_t physicalDeviceIdx = findBestPhysicalDevice(args.physicalDeviceCriteria, mExtensions, physicalDevices);
    mDeviceInfo = &physicalDevices[physicalDeviceIdx];
    const std::vector<const char*> enabledExtensions = buildDeviceExtensionNameList(*mDeviceInfo, mExtensions);
    const std::vector<vk::DeviceQueueCreateInfo> queueCreateInfos = buildQueueCreateInfoList(*mDeviceInfo);

    void* pNext = nullptr;
#define APPEND_FEATURE_STRUCT(type, name, requiredArray)               \
    type name{.pNext = pNext};                                         \
    if (!args.physicalDeviceCriteria.requiredArray.empty())            \
    {                                                                  \
        for (auto feature : args.physicalDeviceCriteria.requiredArray) \
        {                                                              \
            (name).*(feature) = VK_TRUE;                               \
        }                                                              \
        pNext = &(name);                                               \
    }
    APPEND_FEATURE_STRUCT(vk::PhysicalDeviceVulkan11Features, vulkan11Features, requiredVulkan11Features)
    APPEND_FEATURE_STRUCT(vk::PhysicalDeviceVulkan12Features, vulkan12Features, requiredVulkan12Features)
    APPEND_FEATURE_STRUCT(vk::PhysicalDeviceVulkan13Features, vulkan13Features, requiredVulkan13Features)
    APPEND_FEATURE_STRUCT(vk::PhysicalDeviceAccelerationStructureFeaturesKHR, accelerationStructureFeatures, requiredAccelerationStructureFeatures)
    APPEND_FEATURE_STRUCT(vk::PhysicalDeviceRayTracingPipelineFeaturesKHR, rayTracingPipelineFeature, requiredRayTracingPipelineFeatures)
    APPEND_FEATURE_STRUCT(vk::PhysicalDeviceMeshShaderFeaturesEXT, meshShaderFeatures, requredMeshShaderFeatures)
#undef APPEND_FEATURE_STRUCT

    vk::PhysicalDeviceFeatures enabledFeatures; // NOLINT(*-const-correctness) false positive
    for (auto feature : args.physicalDeviceCriteria.requiredFeatures)
    {
        enabledFeatures.*feature = VK_TRUE;
    }

    const vk::DeviceCreateInfo deviceCreateInfo =
    {
        .pNext = pNext,
        .queueCreateInfoCount = static_cast<std::uint32_t>(queueCreateInfos.size()),
        .pQueueCreateInfos = queueCreateInfos.data(),
        .enabledExtensionCount = static_cast<std::uint32_t>(enabledExtensions.size()),
        .ppEnabledExtensionNames = enabledExtensions.data(),
        .pEnabledFeatures = &enabledFeatures
    };
    mHandle = mDeviceInfo->device.createDevice(deviceCreateInfo);
    if (args.flags.singleDevice)
    {
        VULKAN_HPP_DEFAULT_DISPATCHER.init(mHandle);
    }

    mGraphicsQueue = mHandle.getQueue(mDeviceInfo->graphicsQueueFamily, 0);
    if (mDeviceInfo->graphicsQueueFamily == mDeviceInfo->computeQueueFamily)
    {
        mComputeQueue = mGraphicsQueue;
    }
    else
    {
        mComputeQueue = mHandle.getQueue(mDeviceInfo->computeQueueFamily, 0);
    }

    getOwner()->getMainTaskLoop().addTask(c_updateLoop(), &mUpdateLoopHandle);
    getOwner()->deviceCreated.emit(*this);
}

Device::~Device() noexcept
{
    mUpdateLoopHandle.cancel();
    mScratchCommandPools.clear(); // hack2: due to the below hack the pools would get destroyed after the device, reset them manually
    if (mHandle)
    {
        mHandle.waitIdle();
        mPendingScratchCmdBuffers.clear();
        getOwner()->queueDelete([handle=mHandle]()
        {
            handle.destroy();
        });
    }
}

ScratchCommandPool::ScratchCommandPool(Device& device)
{
    mCommandPool = device.createChild<CommandPool>(CommandPoolCreationArgs{
        .flags = vk::CommandPoolCreateFlagBits::eResetCommandBuffer,
        .queueFamilyIndex = device.getDeviceInfo().graphicsQueueFamily
    });
}

ObjectPtr<CommandBuffer> ScratchCommandPool::allocateCommandBuffer()
{
    for (Buffer& buffer : mBuffers)
    {
        if (buffer.doneFuture->ready())
        {
            buffer.doneFuture = std::make_shared<mijin::Future<void>>();
            buffer.cmdBuffer->getVkHandle().reset();
            return buffer.cmdBuffer;
        }
    }

    // nothing found, allocate a new one
    return mBuffers.emplace_back(mCommandPool->allocateCommandBuffer(), std::make_shared<mijin::Future<void>>()).cmdBuffer;
}

mijin::FuturePtr<void> ScratchCommandPool::getFuture(const ObjectPtr<CommandBuffer>& cmdBuffer) noexcept
{
    for (const Buffer& buffer : mBuffers)
    {
        if (buffer.cmdBuffer == cmdBuffer)
        {
            return buffer.doneFuture;
        }
    }
    logAndDie("Someone passed an invalid cmdBuffer to getFuture()!");
}

ObjectPtr<DeviceMemory> Device::allocateDeviceMemory(const DeviceMemoryAllocationArgs& args)
{
    return createChild<DeviceMemory>(args);
}

ObjectPtr<CommandBuffer> Device::beginScratchCommandBuffer()
{
    ObjectPtr<CommandBuffer> cmdBuffer;
    {
        const std::shared_lock readLock(mScratchCommandPoolsMutex);
        auto it = mScratchCommandPools.find(std::this_thread::get_id());
        if (it != mScratchCommandPools.end())
        {
            cmdBuffer = it->second.allocateCommandBuffer();
        }
    }

    if (cmdBuffer == nullptr)
    {
        // create the scratch command pool
        const std::unique_lock writeLock(mScratchCommandPoolsMutex);
        ScratchCommandPool& pool = mScratchCommandPools.emplace(std::this_thread::get_id(), ScratchCommandPool(*this)).first->second;
        cmdBuffer = pool.allocateCommandBuffer();

        // hack: an object should normally not contain pointers to its own child as that leads to cyclic dependencies
        //       decrease reference count by 1 to migitate that
        decreaseReferenceCount();
    }

    const vk::CommandBufferBeginInfo beginInfo =
    {
        .flags = vk::CommandBufferUsageFlagBits::eOneTimeSubmit
    };
    cmdBuffer->getVkHandle().begin(beginInfo);
    return cmdBuffer;
}

mijin::FuturePtr<void> Device::endScratchCommandBuffer(ObjectPtr<CommandBuffer> cmdBuffer)
{
    mijin::FuturePtr<void> future;
    {
        const std::shared_lock readLock(mScratchCommandPoolsMutex);
        future = mScratchCommandPools.at(std::this_thread::get_id()).getFuture(cmdBuffer);
    }
    cmdBuffer->getVkHandle().end();

    getOwner()->runOnMainThread([cmdBuffer = std::move(cmdBuffer), self = getPointer(), future]() mutable
    {
        ObjectPtr<Fence> doneFence = Fence::create(self);

        const vk::SubmitInfo submitInfo =
        {
            .commandBufferCount = 1,
            .pCommandBuffers = &cmdBuffer->getVkHandle()
        };
        self->getGraphicsQueue().submit(submitInfo, *doneFence);
        self->mPendingScratchCmdBuffers.push_back({
            .cmdBuffer = std::move(cmdBuffer),
            .doneFence = std::move(doneFence),
            .future    = std::move(future)
        });
    });
    return future;
}

void Device::queueDelete(std::function<void()> deleter) noexcept
{
    getOwner()->queueDelete(std::move(deleter));
}

mijin::Task<> Device::c_updateLoop() noexcept
{
    while(!getOwner()->isQuitRequested())
    {
        for (auto it = mPendingScratchCmdBuffers.begin(); it != mPendingScratchCmdBuffers.end();)
        {
            if (it->doneFence->isDone())
            {
                it->future->set();
                it = mPendingScratchCmdBuffers.erase(it);
            } else
            {
                ++it;
            }
        }
        co_await mijin::c_suspend();
    }
}
}