#include "vulkan_helper.h"
#include "vulkan/vulkan_core.h"
#include <cstdint>
#include <cstdlib>
#include <iostream>

namespace iris {

Device::Device(std::vector<std::string>  layers, std::vector<std::string> instance_extensions) {
    VkApplicationInfo app_info = {
        .sType = VK_STRUCTURE_TYPE_APPLICATION_INFO,
        .pApplicationName = "IrisRenderer",
        .applicationVersion = VK_MAKE_VERSION(1, 0, 0),
        .pEngineName = "No Engine",
        .engineVersion = VK_MAKE_VERSION(1, 0, 0),
        .apiVersion = VK_API_VERSION_1_0,
    };

    // Create the Vulkan instance
    uint32_t enabled_layer_count = layers.size();
    std::vector<const char*> layers_cstr;
    for (const auto& layer : layers) {
        layers_cstr.push_back(layer.c_str());
    }

    uint32_t enabled_extension_count = instance_extensions.size();
    std::vector<const char*> instance_extensions_cstr;
    for (const auto& extension : instance_extensions) {
        instance_extensions_cstr.push_back(extension.c_str());
    }

    VkInstanceCreateInfo instance_info = {
        .sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
        .pApplicationInfo = &app_info,
        .enabledLayerCount = enabled_layer_count,
        .ppEnabledLayerNames = enabled_layer_count == 0 ? VK_NULL_HANDLE : layers_cstr.data(),
        .enabledExtensionCount = enabled_extension_count,
        .ppEnabledExtensionNames = enabled_extension_count == 0 ? VK_NULL_HANDLE : instance_extensions_cstr.data(),
    };
    CHECK_VULKAN(vkCreateInstance(
        &instance_info,
        VK_NULL_HANDLE,
        &instance));

    // Enumerate and select the physical device
    uint32_t physical_device_count = 0;
    CHECK_VULKAN(vkEnumeratePhysicalDevices(
        instance,
        &physical_device_count,
        VK_NULL_HANDLE));
    std::vector<VkPhysicalDevice> physical_devices(physical_device_count);
    CHECK_VULKAN(vkEnumeratePhysicalDevices(
        instance,
        &physical_device_count,
        physical_devices.data()));
    // For now, just select the first physical device, optionally check capabilities of the device
    physical_device = physical_devices[0];
    {
        uint32_t device_extension_count = 0;
        CHECK_VULKAN(vkEnumerateDeviceExtensionProperties(
            physical_device,
            VK_NULL_HANDLE,
            &device_extension_count,
            VK_NULL_HANDLE));
        std::vector<VkExtensionProperties> device_extensions(device_extension_count);
        CHECK_VULKAN(vkEnumerateDeviceExtensionProperties(
            physical_device,
            VK_NULL_HANDLE,
            &device_extension_count,
            device_extensions.data()));
        
        bool has_raytracing = false;
        bool has_acceleration_structure = false;
        for (const auto& extension : device_extensions) {
            if (std::string(extension.extensionName) == VK_KHR_RAY_TRACING_PIPELINE_EXTENSION_NAME) {
                has_raytracing = true;
            }
            if (std::string(extension.extensionName) == VK_KHR_ACCELERATION_STRUCTURE_EXTENSION_NAME) {
                has_acceleration_structure = true;
            }
        }

        if (!has_raytracing || !has_acceleration_structure) {
            // TODO throw an exception
            std::cerr << "Physical device does not support ray tracing extensions" << std::endl;
            abort();
        }
    }

    // Create the logical device
    float queue_priority = 1.0f;
    main_queue_family_index = 0; // TODO: query capabilities to find a proper queue index
    VkDeviceQueueCreateInfo queue_info = {
        .sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
        .queueFamilyIndex = main_queue_family_index, 
        .queueCount = 1,
        .pQueuePriorities = &queue_priority,
    };
    
    VkPhysicalDeviceAccelerationStructureFeaturesKHR acceleration_structure_features = {
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ACCELERATION_STRUCTURE_FEATURES_KHR,
        .accelerationStructure = VK_TRUE,
    };
    VkPhysicalDeviceRayTracingPipelineFeaturesKHR raytracing_features = {
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_RAY_TRACING_PIPELINE_FEATURES_KHR,
        .pNext = &acceleration_structure_features,
        .rayTracingPipeline = VK_TRUE,
    };
    VkPhysicalDeviceFeatures2 device_features = {
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2,
        .pNext = &raytracing_features,
        .features {
            .samplerAnisotropy = VK_TRUE,
        }
    };

    constexpr char *device_extensions[] = {
        VK_EXT_SHADER_ATOMIC_FLOAT_EXTENSION_NAME,
        VK_KHR_RAY_TRACING_PIPELINE_EXTENSION_NAME,
        VK_KHR_ACCELERATION_STRUCTURE_EXTENSION_NAME,
        VK_KHR_DEFERRED_HOST_OPERATIONS_EXTENSION_NAME
    };
    VkDeviceCreateInfo device_info = {
        .sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
        .pNext = &device_features,
        .queueCreateInfoCount = 1,
        .pQueueCreateInfos = &queue_info,
        .enabledExtensionCount = sizeof(device_extensions) / sizeof(device_extensions[0]),
        .ppEnabledExtensionNames = device_extensions,
    };
    CHECK_VULKAN(vkCreateDevice(
        physical_device,
        &device_info,
        VK_NULL_HANDLE,
        &device));
    
    // Get the graphics queue
    vkGetDeviceQueue(
        device,
        main_queue_family_index,
        0,
        &graphics_queue);
}

} // namespace iris