#include <deimos/core/api_registry.h>
#include <deimos/core/log.h>
#include <deimos/core/temp_allocator.h>
#include <deimos/core/os.h>
#include <deimos/core/status.h>
#include <deimos/vulkan/vulkan.h>

using namespace deimos;

static LogApi* log_api;
static TempAllocatorApi* temp_api;
static OsApi* os_api;

static const VkAllocationCallbacks* kVkAlloc = nullptr;

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 InternalError("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 InternalError("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 InternalError("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 InternalError("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 InternalError("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 InternalError("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 InternalError("vkCreateDeviceFailed");
    }

    return device;
}

Status InitializeVulkan(ApiRegistry* api_registry, OsWindow* window)
{
    RegisterVulkanLoaderApi(api_registry);
    auto* loader_api = api_registry->Get<VulkanLoaderApi>();
    auto* vk = loader_api->LoadEntry();

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

    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");
    }

    loader_api->LoadDevice(vk, device);

    VkQueue queue = VK_NULL_HANDLE;
    vk->GetDeviceQueue(device, queue_family, 0, &queue);
    Ensures(queue != VK_NULL_HANDLE);

    vk->DestroyDevice(device, kVkAlloc);
    vk->DestroySurfaceKHR(instance, surface, kVkAlloc);
    vk->DestroyInstance(instance, kVkAlloc);
    
    return {};
}

int main(int /* argc */, char* /* argv */[])
{
    auto* api_registry = InitializeGlobalApiRegistry();
    log_api = api_registry->Get<LogApi>();
    temp_api = api_registry->Get<TempAllocatorApi>();
    os_api = api_registry->Get<OsApi>();
    
    log_api->LogInfo("Base APIs registered");

    OsWindow* window{};
    {
        auto s = os_api->window->Create("Deimos", 1280, 720);
        if (!s.ok())
        {
            log_api->LogError("Window error: $", s);
            return 1;
        }
        window = s.value();
    }

    const Status s = InitializeVulkan(api_registry, window);
    if (!s.ok())
    {
        log_api->LogError("Vulkan initializaiton error: $", s);
        return 1;
    }

    while (!os_api->window->QuitRequested(window))
    {
        os_api->window->Update(window);
    }

    log_api->LogInfo("Goodbye");

    return 0;
}