#include "deimos/vulkan/vulkan_backend.h"

#include "deimos/vulkan/vulkan_loader.h"

#include <deimos/core/api_registry.h>
#include <deimos/core/allocator.h>
#include <deimos/core/temp_allocator.h>
#include <deimos/core/log.h>
#include <deimos/core/os.h>
#include <deimos/render/backend.h>

namespace
{

using namespace deimos;

LogApi*           log_api;
OsApi*            os_api;
ApiRegistry*      api_registry;
TempAllocatorApi* temp_api;
VulkanLoaderApi*  vulkan_loader_api;

const VkAllocationCallbacks* kVkAlloc = nullptr;

class VulkanBackendImpl : public IVulkanBackend, public IRenderBackend
{
    VulkanApi*       m_vk;
    VkInstance       m_instance;
    VkPhysicalDevice m_physical_device;
    uint32_t         m_queue_family;
    VkSurfaceKHR     m_surface;
    VkDevice         m_device;
    VkQueue          m_queue = VK_NULL_HANDLE;

    struct FrameResources
    {
        VkCommandPool cmd_pool = VK_NULL_HANDLE;
    };

    FrameResources   m_frame_resources;

public:
    VulkanBackendImpl(
        VulkanApi* vk,
        VkInstance instance,
        VkPhysicalDevice physical_device,
        uint32_t queue_family,
        VkSurfaceKHR surface,
        VkDevice device) :
        m_vk{vk}, m_instance{instance}, m_physical_device{physical_device},
        m_queue_family{queue_family}, m_surface{surface}, m_device{device}
    {
        m_vk->GetDeviceQueue(m_device, m_queue_family, 0, &m_queue);
        Ensures(m_queue != VK_NULL_HANDLE);
    }

    IRenderBackend* AsRenderBackend() override
    {
        return this;
    }

    void BeginFrame() override
    {
        VkCommandPoolCreateInfo create_info;
        create_info.sType            = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
        create_info.pNext            = nullptr;
        create_info.flags            = 0;
        create_info.queueFamilyIndex = m_queue_family;

        VkResult res = m_vk->CreateCommandPool(m_device, &create_info, kVkAlloc, &m_frame_resources.cmd_pool);
        if (res != VK_SUCCESS)
        {
            log_api->LogError("Couldn't allocate Vulkan command pool: $", (int)res);
            return;
        }
    }

    void EndFrame() override
    {
        m_vk->DeviceWaitIdle(m_device);
        m_vk->DestroyCommandPool(m_device, m_frame_resources.cmd_pool, kVkAlloc);
    }
};

StatusOr<VkInstance> CreateInstance(VulkanApi* vk)
{
    const VkApplicationInfo application_info{
        .sType              = VK_STRUCTURE_TYPE_APPLICATION_INFO,
        .pNext              = nullptr,
        .pApplicationName   = "Deimos game",
        .applicationVersion = 0,
        .pEngineName        = "Deimos engine",
        .engineVersion      = 0,
        .apiVersion         = VK_API_VERSION_1_3,
    };

    const char* extensions[]{
        VK_KHR_SURFACE_EXTENSION_NAME,
        VK_KHR_WIN32_SURFACE_EXTENSION_NAME,
    };

    // @Todo Select layers & extensions based on config
    const char* layers[]{
        "VK_LAYER_KHRONOS_validation",
        "VK_LAYER_LUNARG_monitor",
    };
    
    const VkInstanceCreateInfo create_info{
        .sType                   = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
        .pNext                   = nullptr,
        .flags                   = 0,
        .pApplicationInfo        = &application_info,
        .enabledLayerCount       = (uint32_t)ArraySize(layers),
        .ppEnabledLayerNames     = layers,
        .enabledExtensionCount   = (uint32_t)ArraySize(extensions),
        .ppEnabledExtensionNames = extensions,
    };

    VkInstance instance{};
    const VkResult res = vk->CreateInstance(&create_info, kVkAlloc, &instance);
    if (res != VK_SUCCESS)
    {
        return RuntimeError("vkCreateInstance failed");
    }

    return instance;
}

StatusOr<VkSurfaceKHR> CreateSurface(VulkanApi* vk, VkInstance instance, OsWindow* window)
{
    const VkWin32SurfaceCreateInfoKHR create_info{
        .sType                = VK_STRUCTURE_TYPE_WIN32_SURFACE_CREATE_INFO_KHR,
        .pNext                = nullptr,
        .flags                = 0,
        .hinstance            = os_api->window->Win32Hinstance(window),
        .hwnd                 = os_api->window->Win32Hwnd(window),
    };

    VkSurfaceKHR surface{};
    const VkResult res = vk->CreateWin32SurfaceKHR(instance, &create_info, kVkAlloc, &surface);
    if (res != VK_SUCCESS)
    {
        return RuntimeError("vkCreateWin32SurfaceKHR failed");
    }

    return surface;
}

StatusOr<uint32_t> FindQueueFamily(VulkanApi* vk, VkPhysicalDevice physical_device, VkSurfaceKHR surface)
{
    auto temp_alloc = temp_api->Acquire();

    uint32_t queue_family_count = 0;
    vk->GetPhysicalDeviceQueueFamilyProperties(physical_device, &queue_family_count, nullptr);
    if (queue_family_count == 0) 
    {
        return RuntimeError("No queue on this physical device");
    }

    log_api->LogInfo("Physical device has $ queue families", queue_family_count);

    auto queue_families = temp_alloc.allocator().NewArray<VkQueueFamilyProperties>(queue_family_count);
    vk->GetPhysicalDeviceQueueFamilyProperties(physical_device, &queue_family_count, queue_families.data());

    for (uint32_t index = 0; index < queue_family_count; ++index)
    {
        const auto& prps = queue_families[index];

        const bool supports_graphics_compute = (prps.queueFlags & VK_QUEUE_GRAPHICS_BIT) != 0 && (prps.queueFlags & VK_QUEUE_COMPUTE_BIT) != 0;

        VkBool32 vk_supports_presentation{};
        const VkResult presentation_res = vk->GetPhysicalDeviceSurfaceSupportKHR(physical_device, index, surface, &vk_supports_presentation);
        const bool supports_presentation = presentation_res == VK_SUCCESS && vk_supports_presentation == VK_TRUE;
        
        if (supports_graphics_compute && supports_presentation)
        {
            return index;
        }
    }

    return RuntimeError("Couldn't find a suitable queue");
}

Status FindPhysicalDevice(
    VulkanApi* vk,
    VkInstance instance,
    VkSurfaceKHR surface,
    VkPhysicalDevice* out_physical_device,
    uint32_t* out_queue_family)
{
    auto temp_alloc = temp_api->Acquire();
    
    uint32_t physical_device_count = 0;
    vk->EnumeratePhysicalDevices(instance, &physical_device_count, nullptr);
    if (physical_device_count == 0)
    {
        return RuntimeError("No Vulkan device found");
    }

    log_api->LogInfo("Found $ physical devices", physical_device_count);

    auto physical_devices = temp_alloc.allocator().NewArray<VkPhysicalDevice>(physical_device_count);
    vk->EnumeratePhysicalDevices(instance, &physical_device_count, physical_devices.data());

    for (VkPhysicalDevice physical_device: physical_devices)
    {
        {
            VkPhysicalDeviceVulkan12Properties prps12{};
            prps12.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_PROPERTIES;
            prps12.pNext = nullptr;
    
            VkPhysicalDeviceProperties2 prps{};
            prps.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2;
            prps.pNext = &prps12;
    
            vk->GetPhysicalDeviceProperties2(physical_device, &prps);
            
            log_api->LogInfo("Trying Vulkan device $ ($, $)",
                prps.properties.deviceName, prps12.driverName, prps12.driverInfo);
        }

        auto maybe_queue_family = FindQueueFamily(vk, physical_device, surface);
        if (maybe_queue_family.ok())
        {
            log_api->LogInfo("This device is compatible, choosing it");
            *out_physical_device = physical_device;
            *out_queue_family = maybe_queue_family.value();
            return {};
        }
        
        log_api->LogInfo("Incompatible because: $", maybe_queue_family);
    }

    return RuntimeError("No suitable device found");
}

StatusOr<VkDevice> CreateDevice(VulkanApi* vk, VkPhysicalDevice physical_device, uint32_t queue_family)
{
    const float queue_priority = 1.0F;
    
    const VkDeviceQueueCreateInfo queue_create_info{
        .sType                    = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
        .pNext                    = nullptr,
        .flags                    = 0,
        .queueFamilyIndex         = queue_family,
        .queueCount               = 1,
        .pQueuePriorities         = &queue_priority,
    };
    
    const char* extensions[]{
        VK_KHR_SWAPCHAIN_EXTENSION_NAME,
    };
    
    const VkDeviceCreateInfo create_info{
        .sType                    = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
        .pNext                    = nullptr,
        .flags                    = 0,
        .queueCreateInfoCount     = 1,
        .pQueueCreateInfos        = &queue_create_info,
        .enabledLayerCount        = 0,
        .ppEnabledLayerNames      = nullptr,
        .enabledExtensionCount    = ArraySize(extensions),
        .ppEnabledExtensionNames  = extensions,
        .pEnabledFeatures         = nullptr,
    };

    VkDevice device = VK_NULL_HANDLE;
    const VkResult res = vk->CreateDevice(physical_device, &create_info, kVkAlloc, &device);
    if (res != VK_SUCCESS)
    {
        return RuntimeError("vkCreateDeviceFailed");
    }

    return device;
}

StatusOr<gsl::owner<IVulkanBackend*>> InitializeVulkan(Allocator* allocator, OsWindow* window)
{
    auto* vk = vulkan_loader_api->LoadEntry(allocator);

    VkInstance instance{};
    {
        StatusOr<VkInstance> s = CreateInstance(vk);
        if (!s.ok()) { return s.status(); }
        instance = s.value();
        log_api->LogInfo("Vulkan instance created");
    }

    vulkan_loader_api->LoadInstance(vk, instance);

    VkSurfaceKHR surface{};
    {
        StatusOr<VkSurfaceKHR> s = CreateSurface(vk, instance, window);
        if (!s.ok()) { return s.status(); }
        surface = s.value();
        log_api->LogInfo("Vulkan surface created");
    }

    VkPhysicalDevice physical_device{};
    uint32_t queue_family{};
    {
        Status s = FindPhysicalDevice(vk, instance, surface, &physical_device, &queue_family);
        if (!s.ok()) { return s; }
        log_api->LogInfo("Vulkan queue family: $", queue_family);
    }

    VkDevice device = VK_NULL_HANDLE;
    {
        StatusOr<VkDevice> s = CreateDevice(vk, physical_device, queue_family);
        if (!s.ok()) { return s.status(); }
        device = s.value();
        log_api->LogInfo("Vulkan device created");
    }

    vulkan_loader_api->LoadDevice(vk, device);

    return allocator->New<VulkanBackendImpl>(
        vk, instance, physical_device,
        queue_family, surface, device);
}

} // anonymous namespace

namespace deimos
{

class VulkanBackendApiImpl : public VulkanBackendApi
{
public:
    StatusOr<gsl::owner<IVulkanBackend*>> CreateBackend(Allocator* allocator, OsWindow* window) override
    {
        return InitializeVulkan(allocator, window);
    }
};

void RegisterVulkanBackendApi(ApiRegistry* registry)
{
    api_registry = registry;
    log_api = registry->Get<LogApi>();
    os_api = registry->Get<OsApi>();
    temp_api = registry->Get<TempAllocatorApi>();
    vulkan_loader_api = registry->Get<VulkanLoaderApi>();

    auto* allocator_api = registry->Get<AllocatorApi>();

    auto* impl = allocator_api->system->New<VulkanBackendApiImpl>();
    registry->Set(impl);
    
    log_api->LogInfo("Vulkan backend API registered");
}

} // namespace deimos