#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>

using namespace deimos;

void DeimosFormat(IWriter* writer, VkResult res)
{
#define PRINT_RESULT(X) case VK_##X: writer->Write(AsBytes(StringView(#X))); break; 
    switch (res)
    {
        PRINT_RESULT(SUCCESS);
        PRINT_RESULT(NOT_READY);
        PRINT_RESULT(TIMEOUT);
        PRINT_RESULT(EVENT_SET);
        PRINT_RESULT(EVENT_RESET);
        PRINT_RESULT(INCOMPLETE);
        PRINT_RESULT(ERROR_OUT_OF_HOST_MEMORY);
        PRINT_RESULT(ERROR_OUT_OF_DEVICE_MEMORY);
        PRINT_RESULT(ERROR_INITIALIZATION_FAILED);
        PRINT_RESULT(ERROR_DEVICE_LOST);
        PRINT_RESULT(ERROR_MEMORY_MAP_FAILED);
        PRINT_RESULT(ERROR_LAYER_NOT_PRESENT);
        PRINT_RESULT(ERROR_EXTENSION_NOT_PRESENT);
        PRINT_RESULT(ERROR_FEATURE_NOT_PRESENT);
        PRINT_RESULT(ERROR_INCOMPATIBLE_DRIVER);
        PRINT_RESULT(ERROR_TOO_MANY_OBJECTS);
        PRINT_RESULT(ERROR_FORMAT_NOT_SUPPORTED);
        PRINT_RESULT(ERROR_FRAGMENTED_POOL);
        PRINT_RESULT(ERROR_UNKNOWN);
        PRINT_RESULT(ERROR_OUT_OF_POOL_MEMORY);
        PRINT_RESULT(ERROR_INVALID_EXTERNAL_HANDLE);
        PRINT_RESULT(ERROR_FRAGMENTATION);
        PRINT_RESULT(ERROR_INVALID_OPAQUE_CAPTURE_ADDRESS);
        PRINT_RESULT(PIPELINE_COMPILE_REQUIRED);
        PRINT_RESULT(ERROR_SURFACE_LOST_KHR);
        PRINT_RESULT(ERROR_NATIVE_WINDOW_IN_USE_KHR);
        PRINT_RESULT(SUBOPTIMAL_KHR);
        PRINT_RESULT(ERROR_OUT_OF_DATE_KHR);
        PRINT_RESULT(ERROR_INCOMPATIBLE_DISPLAY_KHR);
        PRINT_RESULT(ERROR_VALIDATION_FAILED_EXT);
        PRINT_RESULT(ERROR_IMAGE_USAGE_NOT_SUPPORTED_KHR);
        PRINT_RESULT(ERROR_NOT_PERMITTED_KHR);
        PRINT_RESULT(ERROR_FULL_SCREEN_EXCLUSIVE_MODE_LOST_EXT);
        PRINT_RESULT(THREAD_IDLE_KHR);
        PRINT_RESULT(THREAD_DONE_KHR);
        PRINT_RESULT(OPERATION_DEFERRED_KHR);
        PRINT_RESULT(OPERATION_NOT_DEFERRED_KHR);
        PRINT_RESULT(ERROR_INCOMPATIBLE_SHADER_BINARY_EXT);
        PRINT_RESULT(ERROR_INVALID_SHADER_NV);
        PRINT_RESULT(ERROR_VIDEO_PICTURE_LAYOUT_NOT_SUPPORTED_KHR);
        PRINT_RESULT(ERROR_VIDEO_PROFILE_OPERATION_NOT_SUPPORTED_KHR);
        PRINT_RESULT(ERROR_VIDEO_PROFILE_FORMAT_NOT_SUPPORTED_KHR);
        PRINT_RESULT(ERROR_VIDEO_PROFILE_CODEC_NOT_SUPPORTED_KHR);
        PRINT_RESULT(ERROR_VIDEO_STD_VERSION_NOT_SUPPORTED_KHR);
        PRINT_RESULT(ERROR_INVALID_DRM_FORMAT_MODIFIER_PLANE_LAYOUT_EXT);
        PRINT_RESULT(ERROR_INVALID_VIDEO_STD_PARAMETERS_KHR);
        PRINT_RESULT(ERROR_COMPRESSION_EXHAUSTED_EXT);
        PRINT_RESULT(RESULT_MAX_ENUM);
    }
}

namespace
{

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

const VkAllocationCallbacks* kVkAlloc = nullptr;

struct VulkanSwapchain: public RenderSwapchain
{
    VkSwapchainKHR            swapchain{};
    VkFormat                  format{};
    gsl::owner<Span<VkImage>> images;
};

class VulkanBackendImpl : public IVulkanBackend, public IRenderBackend
{
    Allocator*       m_allocator;
    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;
        VkCommandBuffer cmd_buffer = VK_NULL_HANDLE;
    };

    FrameResources   m_frame_resources;

public:
    VulkanBackendImpl(
        Allocator* allocator,
        VulkanApi* vk,
        VkInstance instance,
        VkPhysicalDevice physical_device,
        uint32_t queue_family,
        VkSurfaceKHR surface,
        VkDevice device) :
        m_allocator{allocator},
        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: $", res);
            return;
        }

        VkCommandBufferAllocateInfo alloc_info{};
        alloc_info.sType              = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
        alloc_info.pNext              = nullptr;
        alloc_info.commandPool        = m_frame_resources.cmd_pool;
        alloc_info.level              = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
        alloc_info.commandBufferCount = 1;

        res = m_vk->AllocateCommandBuffers(m_device, &alloc_info, &m_frame_resources.cmd_buffer);
        if (res != VK_SUCCESS)
        {
            log_api->LogError("Couldn't allocate Vulkan command buffer: $", res);
        }

        VkCommandBufferBeginInfo begin_info{};
        begin_info.sType             = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
        begin_info.pNext             = nullptr;
        begin_info.flags             = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
        begin_info.pInheritanceInfo  = nullptr;

        res = m_vk->BeginCommandBuffer(m_frame_resources.cmd_buffer, &begin_info);
        if (res != VK_SUCCESS)
        {
            log_api->LogError("Couldn't begin Vulkan command buffer recording: $", res);
        }
    }

    void EndFrame() override
    {
        VkResult res = m_vk->EndCommandBuffer(m_frame_resources.cmd_buffer);
        if (res != VK_SUCCESS)
        {
            log_api->LogError("Couldn't end Vulkan command buffer recording: $", res);
        }

        VkCommandBufferSubmitInfo cmd_buffer_submit_info{};
        cmd_buffer_submit_info.sType         = VK_STRUCTURE_TYPE_COMMAND_BUFFER_SUBMIT_INFO;
        cmd_buffer_submit_info.pNext         = nullptr;
        cmd_buffer_submit_info.commandBuffer = m_frame_resources.cmd_buffer;
        cmd_buffer_submit_info.deviceMask    = 0;

        VkSubmitInfo2 submit_info{};
        submit_info.sType                    = VK_STRUCTURE_TYPE_SUBMIT_INFO_2;
        submit_info.pNext                    = nullptr;
        submit_info.flags                    = 0;
        submit_info.waitSemaphoreInfoCount   = 0;
        submit_info.commandBufferInfoCount   = 1;
        submit_info.pCommandBufferInfos      = &cmd_buffer_submit_info;
        submit_info.signalSemaphoreInfoCount = 0;

        res = m_vk->QueueSubmit2(m_queue, 1, &submit_info, VK_NULL_HANDLE);
        if (res != VK_SUCCESS)
        {
            log_api->LogError("Couldn't submit Vulkan queue: $", res);
        }
        
        m_vk->DeviceWaitIdle(m_device);
        m_vk->DestroyCommandPool(m_device, m_frame_resources.cmd_pool, kVkAlloc);
    }
    
    StatusOr<gsl::owner<RenderSwapchain*>> CreateSwapchain() override
    {
        auto temp_alloc = temp_api->Acquire();
        
        VkSurfaceCapabilitiesKHR caps;
        VkResult res = m_vk->GetPhysicalDeviceSurfaceCapabilitiesKHR(m_physical_device, m_surface, &caps);
        if (res != VK_SUCCESS)
        {
            return RuntimeError("Couldn't retrieve Vulkan surface capabitilies");
        }

        uint32_t format_count = 0;
        uint32_t present_mode_count = 0;
        
        m_vk->GetPhysicalDeviceSurfaceFormatsKHR(m_physical_device, m_surface, &format_count, nullptr);
        m_vk->GetPhysicalDeviceSurfacePresentModesKHR(m_physical_device, m_surface, &present_mode_count, nullptr);
        
        auto formats       = temp_alloc.allocator().NewArray<VkSurfaceFormatKHR>(format_count);
        auto present_modes = temp_alloc.allocator().NewArray<VkPresentModeKHR>(present_mode_count);

        m_vk->GetPhysicalDeviceSurfaceFormatsKHR(m_physical_device, m_surface, &format_count, formats.data());
        m_vk->GetPhysicalDeviceSurfacePresentModesKHR(m_physical_device, m_surface, &present_mode_count, present_modes.data());

        VkPresentModeKHR present_mode = VK_PRESENT_MODE_FIFO_KHR;
        if (present_modes.Contains(VK_PRESENT_MODE_MAILBOX_KHR))
        {
            present_mode = VK_PRESENT_MODE_MAILBOX_KHR;
        }

        static constexpr VkFormat kAcceptableFormats[] = {
            VK_FORMAT_R8G8B8_SRGB,
            VK_FORMAT_B8G8R8_SRGB,
            VK_FORMAT_R8G8B8A8_SRGB,
            VK_FORMAT_B8G8R8A8_SRGB,
        };

        bool has_format = false;
        VkSurfaceFormatKHR format;
        for (const auto& candidate: kAcceptableFormats)
        {
            for (const auto& f: formats)
            {
                if (f.format == candidate && f.colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR)
                {
                    has_format = true;
                    format = f;
                    break;
                }
            }

            if (has_format) { break; }
        }


        if (!has_format)
        {
            return RuntimeError("No acceptable Vulkan swapchain format found");
        }

        const VkSwapchainCreateInfoKHR create_info = {
            .sType                        = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
            .pNext                        = nullptr,
            .flags                        = 0,
            .surface                      = m_surface,
            .minImageCount                = Clamp<uint32_t>(2, caps.minImageCount, caps.maxImageCount),
            .imageFormat                  = format.format,
            .imageColorSpace              = format.colorSpace,
            .imageExtent                  = caps.currentExtent,
            .imageArrayLayers             = 1,
            .imageUsage                   = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT,
            .imageSharingMode             = VK_SHARING_MODE_EXCLUSIVE,
            .queueFamilyIndexCount        = 0,
            .pQueueFamilyIndices          = nullptr,
            .preTransform                 = VK_SURFACE_TRANSFORM_IDENTITY_BIT_KHR,
            .compositeAlpha               = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
            .presentMode                  = present_mode,
            .clipped                      = VK_TRUE,
            .oldSwapchain                 = VK_NULL_HANDLE,
        };

        VkSwapchainKHR vk_swapchain{};
        res = m_vk->CreateSwapchainKHR(m_device, &create_info, kVkAlloc, &vk_swapchain);
        if (res != VK_SUCCESS)
        {
            return RuntimeError("Error when creating Vulkan swapchain");
        }

        auto* swapchain = m_allocator->New<VulkanSwapchain>();
        swapchain->swapchain  = vk_swapchain;
        swapchain->format     = format.format;

        uint32_t image_count = 0;
        res = m_vk->GetSwapchainImagesKHR(m_device, vk_swapchain, &image_count, nullptr);
        if (res != VK_SUCCESS)
        {
            return RuntimeError("Error when retrieving Vulkan swapchain images");
        }

        swapchain->images = m_allocator->NewArray<VkImage>(image_count);
        m_vk->GetSwapchainImagesKHR(m_device, vk_swapchain, &image_count, swapchain->images.data());

        log_api->LogInfo("Vulkan swapchain created with $ images, format $, present mode $", image_count, (int)format.format, (int)present_mode);

        return swapchain;
    }
    
    void DestroySwapchain(gsl::owner<RenderSwapchain*> swapchain) override
    {
        auto* vk_swapchain = (VulkanSwapchain*)swapchain; // NOLINT
        m_vk->DestroySwapchainKHR(m_device, vk_swapchain->swapchain, kVkAlloc);
        m_allocator->DeleteArray(vk_swapchain->images);
        m_allocator->Delete(swapchain);
    }
};

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 VkPhysicalDeviceSynchronization2Features synchronization2_features{
        .sType                    = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SYNCHRONIZATION_2_FEATURES,
        .pNext                    = nullptr,
        .synchronization2         = VK_TRUE,
    };
    
    const VkDeviceCreateInfo create_info{
        .sType                    = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
        .pNext                    = &synchronization2_features,
        .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>(
        allocator,
        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