#ifndef VULKAN_DEVICE_HPP_
#define VULKAN_DEVICE_HPP_

#include <memory>
#include <vector>
#include <set>
#include <optional>

#include <Utilities/Extern.hpp>
#include <Vk/Instance/Instance.hpp>
#include <Vk/Surface/Surface.hpp>
#include <vulkan/vulkan_core.h>

namespace Vk
{
    /**
     * @brief Represents a Vulkan logical device.
     */
    class Device
    {
    public:
        /**
         * @brief Represents swapchain support details.
         */
        struct SwapchainSupportDetails
        {
            VkSurfaceCapabilitiesKHR capabilities;
            std::vector<VkSurfaceFormatKHR> formats;
            std::vector<VkPresentModeKHR> presentModes;
        };

        /**
         * @brief Represents the queue families on this device.
         */
        struct QueueFamilies
        {
            std::optional<uint32_t> graphicsFamily;
            std::optional<uint32_t> presentFamily;

            bool isComplete()
            {
                return graphicsFamily.has_value() && presentFamily.has_value();
            }
        };

    private:
        VkDevice m_device;
        std::shared_ptr<Vk::Instance> m_instance;
        std::shared_ptr<Vk::Surface> m_surface;
        VkPhysicalDevice m_physicalDevice;
        QueueFamilies m_queueFamilies;
        VkPhysicalDeviceMemoryProperties m_memProperties;
        VmaAllocator m_allocator;

        static const std::vector<const char*> deviceExtensions;

        VkPhysicalDevice pickPhysicalDevice();
        QueueFamilies findQueueFamilies(VkPhysicalDevice device);
        bool areDeviceExtensionsSupported(VkPhysicalDevice device);
        SwapchainSupportDetails querySwapchainSupport(VkPhysicalDevice device);
        std::optional<int> rankPhysicalDevice(VkPhysicalDevice device);

    public:
        VkPhysicalDeviceFeatures features;
        VkPhysicalDeviceProperties properties;

        glm::uvec3 maxWorkGroupCount;
        glm::uvec3 maxWorkGroupSize;
        uint32_t maxWorkGroupInvocations;

        /**
         * @brief Represents a device queue.
         */
        class Queue
        {
        private:
            VkQueue m_queue;

        public:
            Queue() = default;
            Queue(const Queue &) = delete;
            
            friend class Vk::Device;

            /**
             * @brief Get the wrapped queue.
             */
            VkQueue inner();
        };

        Queue graphicsQueue;
        Queue presentQueue;

        /**
         * @brief Construct the logical device.
         * 
         * This constructor enumerates the physical devices of the instance, and finds
         * a suitable one to create a logical device.
         * 
         * @param instance The Vulkan instance.
         * @param surface The window surface.
         */
        Device(std::shared_ptr<Vk::Instance> instance, std::shared_ptr<Vk::Surface> surface);
        Device(const Device &) = delete;
        ~Device();

        /**
         * @brief Get the wrapped device.
         */
        VkDevice inner();

        /**
         * @brief Return the details of swapchain support on this device.
         */
        SwapchainSupportDetails getSwapchainSupportDetails();

        /**
         * @brief Find a supported image format.
         * 
         * Crashes if the format cannot be found.
         * 
         * @param candidates Candidates, from most to least desirable.
         * @param tiling The tiling for the format.
         * @param features The required features.
         * @return The format.
         */
        VkFormat findSupportedFormat(
            const std::vector<VkFormat>& candidates,
            VkImageTiling tiling,
            VkFormatFeatureFlags features
        );

        /**
         * @brief Returns the format that should be used for depth images.
         * 
         * Crashes if no suitable format was found.
         */
        VkFormat findDepthFormat();

        /**
         * @brief Return information about the queue families.
         */
        QueueFamilies getQueueFamilies();

        /**
         * @brief Wait for all operations on this device to complete.
         */
        void waitIdle();

        /**
         * @brief Get the VMA allocator.
         */
        VmaAllocator getAllocator();

        /**
         * @brief Return the number of memory heaps.
         */
        uint32_t getNumMemoryHeaps();

        /**
         * @brief Returns whether or not timer queries should be enabled, based on support and options.
         */
        bool areTimerQueriesEnabled();
    };
}

std::ostream& operator<<(std::ostream &os, const Vk::Device::SwapchainSupportDetails &details);

#endif