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

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_3,
    };

    // 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);

    // Create the memory allocator
    VmaAllocatorCreateInfo allocator_info = {
        .physicalDevice = physical_device,
        .device = device,
        .instance = instance,
    };
    CHECK_VULKAN(vmaCreateAllocator(
        &allocator_info,
        &allocator));
}

Device::~Device() {
    vmaDestroyAllocator(allocator);
    vkDestroyDevice(device, VK_NULL_HANDLE);
    vkDestroyInstance(instance, VK_NULL_HANDLE);
}


Buffer_t::~Buffer_t() {
    vmaDestroyBuffer(allocator, buffer, allocation);
}

Buffer Device::create_buffer(VkDeviceSize size,
                             VkBufferUsageFlags usage,
                             VmaMemoryUsage memory_usage) {
    VkBufferCreateInfo buffer_info = {
        .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
        .size = size,
        .usage = usage,
    };
    VmaAllocationCreateInfo allocation_info = {
        .usage = memory_usage,
    };
    Buffer_t buffer = {
        .allocator = this->allocator,
        .flags = usage,
        .size = size,
    };
    CHECK_VULKAN(vmaCreateBuffer(
        allocator,
        &buffer_info,
        &allocation_info,
        &buffer.buffer,
        &buffer.allocation,
        VK_NULL_HANDLE));
    return std::make_shared<Buffer_t>(buffer);
}

Texture2D_t::~Texture2D_t() {
    vmaDestroyImage(allocator, image, allocation);
    // TODO: optionally destroy image view, if created
}

Texture2D Device::create_texture(VkExtent2D extent,
                                 VkFormat format,
                                 VkImageUsageFlags usage,
                                 VmaMemoryUsage memory_usage) {
    VkImageCreateInfo image_info = {
        .sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
        .imageType = VK_IMAGE_TYPE_2D,
        .format = format,
        .extent = {extent.width, extent.height, 1},
        .mipLevels = 1,
        .arrayLayers = 1,
        .samples = VK_SAMPLE_COUNT_1_BIT,
        .tiling = VK_IMAGE_TILING_OPTIMAL,
        .usage = usage,
    };
    VmaAllocationCreateInfo allocation_info = {
        .usage = memory_usage,
    };
    Texture2D_t texture = {
        .allocator = this->allocator,
        .flags = usage,
        .extent = extent,
    };
    CHECK_VULKAN(vmaCreateImage(
        allocator,
        &image_info,
        &allocation_info,
        &texture.image,
        &texture.allocation,
        VK_NULL_HANDLE));
    return std::make_shared<Texture2D_t>(texture);
}

} // namespace iris