#include <string.h>
#include <vulkan/vulkan.h>

#include "device.h"
#include "common.h"

// TODO: rename device.c/.h
// TODO: add error checking
// TODO: add log output

#define CLAMP(val, min, max) (((val) < (min)) ? (min) : ((val) > (max) ? (max) : (val)))

#ifdef DEBUG
char *validation_layers[] = {
    "VK_LAYER_KHRONOS_validation"
};
#endif

char *device_extensions[] = {
    VK_KHR_SWAPCHAIN_EXTENSION_NAME
};

struct queue_family_idx {
    u32 graphics_family;
    bool has_graphics_family;
    
    u32 present_family;
    bool has_present_family;
};

struct swap_chain_support_details {
    VkSurfaceCapabilitiesKHR capabilities;
    
    VkSurfaceFormatKHR *formats;
    u32 nformats;
    
    VkPresentModeKHR   *present_modes;
    u32 npresent_modes;
};

// most function that return an int follow this scheme:
// 0 - success; 1 - error; (sometimes -1 - fatal error)

// -- Major Functions ---
static int create_instance(device_t device, struct device_info *info);
static int create_surface(device_t device, struct device_info *info);
static int create_physical_device(device_t device, struct device_info *info);
static int create_logical_device(device_t device, struct device_info *info);
static int create_swap_chain(device_t device, struct device_info *info);

// --- Helper Function ---
static bool device_is_suitable(VkPhysicalDevice phy_device, VkSurfaceKHR surface);
static bool device_has_extension_support(VkPhysicalDevice phy_device);
static int device_queue_families(VkPhysicalDevice phy_device, VkSurfaceKHR surface, struct queue_family_idx *queue_family);

static int swap_chain_support(VkPhysicalDevice phy_device, VkSurfaceKHR surface, struct swap_chain_support_details *details);
static void swap_chain_free_support_details(struct swap_chain_support_details *details);
static int swap_chain_choose_format(VkSurfaceFormatKHR* formats, u32 nformats, VkSurfaceFormatKHR *format);
static int swap_chain_choose_present_mode(VkPresentModeKHR *modes, u32 nmodes, VkPresentModeKHR *mode);
static int swap_chain_get_extent(VkSurfaceCapabilitiesKHR capabilities, VkExtent2D *extent);

// --- Debug Functions ---
#ifdef DEBUG
static bool instance_has_validation_layers(const char * const *layers, u32 nlayers);

static int create_debug_messenger(device_t device, struct device_info* info);
static void debug_messenger_populate_info(VkDebugUtilsMessengerCreateInfoEXT *info);

static VKAPI_ATTR VkBool32 VKAPI_CALL debug_messenger_callback(
    VkDebugUtilsMessageSeverityFlagBitsEXT messageSeverity,
    VkDebugUtilsMessageTypeFlagsEXT messageType,
    const VkDebugUtilsMessengerCallbackDataEXT* pCallbackData,
    void* pUserData);

// Vulkan Wrappers
VkResult CreateDebugUtilsMessengerEXT(
    VkInstance instance,
    const VkDebugUtilsMessengerCreateInfoEXT* pCreateInfo,
    const VkAllocationCallbacks* pAllocator,
    VkDebugUtilsMessengerEXT* pDebugMessenger);
void DestroyDebugUtilsMessengerEXT(
    VkInstance instance,
    VkDebugUtilsMessengerEXT debugMessenger,
    const VkAllocationCallbacks* pAllocator);
#endif

#define CHECK(func) if(func(device, info)) { goto fail; }

device_t device_create(struct device_info *info)
{
    device_t device = xmalloc(sizeof(struct device));
    device->instance = VK_NULL_HANDLE;
    device->logical_device = VK_NULL_HANDLE;
    device->swap_chain = VK_NULL_HANDLE;
   
    CHECK(create_instance);
    CHECK(create_surface);
#ifdef DEBUG
    CHECK(create_debug_messenger);
#endif
    CHECK(create_physical_device);
    CHECK(create_logical_device);
    CHECK(create_swap_chain)
    
    return device;    
fail:
    device_destroy(device);
    return NULL;
}

void device_destroy(device_t device)
{
    if(!device) return;

    free(device->swap_chain_images);
    vkDestroySwapchainKHR(device->logical_device, device->swap_chain, NULL);
    
    vkDestroyDevice(device->logical_device, NULL);
    
    #ifdef DEBUG
    DestroyDebugUtilsMessengerEXT(device->instance, device->debug_messenger, NULL);
    #endif
    
    vkDestroySurfaceKHR(device->instance, device->surface, NULL);
    vkDestroyInstance(device->instance, NULL);
    
    free(device);
}

static int create_instance(device_t device, struct device_info *info)
{
    int ret = 1;
    
    VkApplicationInfo app_info = {0};
    app_info.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
    app_info.pEngineName = "Engine";
    app_info.engineVersion = VK_MAKE_VERSION(1, 0, 0);
    app_info.apiVersion = VK_API_VERSION_1_0;

    app_info.pApplicationName = info->name;
    app_info.applicationVersion = info->version;

    VkInstanceCreateInfo create_info = {0};
    create_info.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
    create_info.pApplicationInfo = &app_info;
    create_info.enabledLayerCount = 0;
            
#ifdef DEBUG
    // add debug utils extensions for debug callback
    char **extensions = xcalloc(info->ext_count+1, sizeof(char *));
    extensions[info->ext_count] = VK_EXT_DEBUG_UTILS_EXTENSION_NAME;
        
    for(u32 i = 0; i < info->ext_count; i++) {
        extensions[i] = (char *)info->extensions[i];
    }
        
    create_info.enabledExtensionCount = info->ext_count+1;
    create_info.ppEnabledExtensionNames = (const char * const *)extensions;

    // validation layer support
    u32 nlayers = ARR_SIZE(validation_layers);

    if(!instance_has_validation_layers((const char * const *)validation_layers, nlayers)) {
        err("validation_layer_support: failed");
        goto exit;
    }
        
    create_info.ppEnabledLayerNames = (const char * const *)validation_layers;
    create_info.enabledLayerCount = nlayers;

    // add the debug messenger for instance creation and destruction
    VkDebugUtilsMessengerCreateInfoEXT msg_info = {0};
    debug_messenger_populate_info(&msg_info);
        
    create_info.pNext = &msg_info;
#endif 

    if(vkCreateInstance(&create_info, NULL, &device->instance) != VK_SUCCESS) {
        err("vkCreateInstance: failed");
        goto exit;
    }
    


    ret = 0;
exit:
    
#ifdef DEBUG
    free(extensions);
#endif
    return ret;
}

static int create_surface(device_t device, struct device_info *info)
{
    if(info->surface_func(device->instance, &device->surface)) {
        err("Couldn't create a VkSurfaceKHR");
        return 1;
    }
    
    return 0;
}

static int create_physical_device(device_t device, struct device_info *info)
{
    (void)info;
    int ret = 1;
        
    device->physical_device = VK_NULL_HANDLE;
    
    u32 dev_count = 0;
    vkEnumeratePhysicalDevices(device->instance, &dev_count, NULL);

    if(dev_count == 0) {
        err("No physical devices could be found!");
        goto exit;
    }

    VkPhysicalDevice *devices = xcalloc(dev_count, sizeof(VkPhysicalDevice));
    vkEnumeratePhysicalDevices(device->instance, &dev_count, devices);

    for(u32 i = 0; i < dev_count; i++) {
        if(device_is_suitable(devices[i], device->surface)) {
            device->physical_device = devices[i];
            break;
        }
    }

    free(devices);

    if(device->physical_device == VK_NULL_HANDLE) {
        err("No suitable physical device could be found");
        goto exit;
    }

    ret = 0;
exit:
    return ret;
}

static int create_logical_device(device_t device, struct device_info *info)
{
    (void)info;
    int ret = 1;

    // queue family data
    struct queue_family_idx indices = {0};
    device_queue_families(device->physical_device, device->surface, &indices);

    // queue infos
    #define QUEUE_FAMILY_COUNT 2
    
    VkDeviceQueueCreateInfo queue_infos[QUEUE_FAMILY_COUNT] = {0};
    u32 queue_indices[QUEUE_FAMILY_COUNT] = {
        indices.graphics_family, indices.present_family,
    };

    u32 unique_queue_family_count = 0;

    float queue_priority = 1.0f;
    
    // basically get add only the unique queue families
    for(size_t i = 0; i < QUEUE_FAMILY_COUNT; i++)
    {
        bool unique = true;
        for(size_t j = 0; j < i; j++) {
            if(queue_indices[i] == queue_indices[j]) {
                unique = false;
                break;
            }
        }
        
        if(!unique) continue;
        
        queue_infos[unique_queue_family_count].sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
        queue_infos[unique_queue_family_count].queueFamilyIndex = queue_indices[i];
        queue_infos[unique_queue_family_count].queueCount = 1;
        queue_infos[unique_queue_family_count].pQueuePriorities = &queue_priority;

        unique_queue_family_count++;
    }

    // device features
    VkPhysicalDeviceFeatures device_features = {0};

    // logical device create info
    VkDeviceCreateInfo create_info = {0};
    create_info.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
    
    create_info.pQueueCreateInfos = queue_infos;
    create_info.queueCreateInfoCount = unique_queue_family_count;

    create_info.pEnabledFeatures = &device_features;

    create_info.ppEnabledExtensionNames = (const char * const *)device_extensions;
    create_info.enabledExtensionCount = ARR_SIZE(device_extensions);
 
    // validation layers can be set, but
    // newer implementations will ignore them
    // so i wont bother adding them
    create_info.enabledLayerCount = 0;

    if(vkCreateDevice(device->physical_device, &create_info, NULL, &device->logical_device) != VK_SUCCESS) {
        err("vkCreateDevice: failed");
        goto exit;
    }

    vkGetDeviceQueue(device->logical_device, indices.graphics_family, 0, &device->graphics_queue);
    vkGetDeviceQueue(device->logical_device, indices.present_family, 0, &device->present_queue);

    ret = 0;
exit:
    return ret;
}

static int create_swap_chain(device_t device, struct device_info *info)
{
    (void)info;
    int ret = 1;
    
    struct swap_chain_support_details details;
    swap_chain_support(device->physical_device, device->surface, &details);

    VkSurfaceFormatKHR surface_format;
    VkPresentModeKHR   present_mode;
    VkExtent2D         extent;
    u32                image_count = 3;
    
    swap_chain_choose_format(details.formats, details.nformats, &surface_format);
    swap_chain_choose_present_mode(details.present_modes, details.npresent_modes, &present_mode);
    swap_chain_get_extent(details.capabilities, &extent);

    if(details.capabilities.maxImageCount > 0) {
        image_count = CLAMP(image_count, details.capabilities.minImageCount, details.capabilities.maxImageCount);
    }

    // start filling in the create info
    VkSwapchainCreateInfoKHR create_info = {0};
    create_info.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;
    
    create_info.surface = device->surface;
    create_info.minImageCount = image_count;
    create_info.imageFormat = surface_format.format;
    create_info.imageColorSpace = surface_format.colorSpace;
    create_info.imageExtent = extent;
    create_info.imageArrayLayers = 1;
    create_info.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;

    struct queue_family_idx indices;
    device_queue_families(device->physical_device, device->surface, &indices);

    u32 queue_indices[2] = { indices.graphics_family, indices.present_family };

    // set the sharing mode of the images
    if(indices.graphics_family != indices.present_family) {
        create_info.imageSharingMode = VK_SHARING_MODE_CONCURRENT;
        create_info.queueFamilyIndexCount = 2;
        create_info.pQueueFamilyIndices = queue_indices;
    } else {
        create_info.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
        create_info.queueFamilyIndexCount = 0;
        create_info.pQueueFamilyIndices = NULL;
    }

    create_info.preTransform = details.capabilities.currentTransform;
    create_info.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;

    create_info.presentMode = present_mode;
    create_info.clipped = VK_TRUE;

    create_info.oldSwapchain = NULL;

    if (vkCreateSwapchainKHR(device->logical_device, &create_info, NULL, &device->swap_chain) != VK_SUCCESS) {
        goto exit;
    }

    vkGetSwapchainImagesKHR(device->logical_device, device->swap_chain, &device->nimages, NULL);
    
    device->swap_chain_images = xcalloc(device->nimages, sizeof(*device->swap_chain_images));
    vkGetSwapchainImagesKHR(device->logical_device, device->swap_chain, &device->nimages, device->swap_chain_images);

    ret = 0;
exit:
    swap_chain_free_support_details(&details);
    return ret;
}

static int device_queue_families(VkPhysicalDevice phy_device, VkSurfaceKHR surface, struct queue_family_idx *queue_family)
{   
    u32 count = 0;
    vkGetPhysicalDeviceQueueFamilyProperties(phy_device, &count, NULL);

    VkQueueFamilyProperties *properties = xcalloc(count, sizeof(VkQueueFamilyProperties));
    vkGetPhysicalDeviceQueueFamilyProperties(phy_device, &count, properties);

    for(u32 i = 0; i < count; i++) {
        if(properties[i].queueFlags & VK_QUEUE_COMPUTE_BIT) {
            queue_family->graphics_family = i;
            queue_family->has_graphics_family = true;
        }

        VkBool32 present_support = false;
        vkGetPhysicalDeviceSurfaceSupportKHR(phy_device, i, surface, &present_support);
        
        if(present_support) {
            queue_family->present_family = i;
            queue_family->has_present_family = true;
        }
    }

    free(properties);
    return 0;
}

static bool device_is_suitable(VkPhysicalDevice phy_device, VkSurfaceKHR surface)
{
    struct queue_family_idx indices;
    struct swap_chain_support_details details;
        
    return
        (device_queue_families(phy_device, surface, &indices),
         indices.has_graphics_family && indices.has_present_family)
        &&
        (device_has_extension_support(phy_device))
        &&
        (swap_chain_support(phy_device, surface, &details),
         swap_chain_free_support_details(&details),
         (details.nformats > 0) && (details.npresent_modes > 0));
}

static bool device_has_extension_support(VkPhysicalDevice phy_device)
{
    u32 count = 0;
    vkEnumerateDeviceExtensionProperties(phy_device, NULL, &count, NULL);

    VkExtensionProperties *properties = xcalloc(count, sizeof(VkExtensionProperties));
    vkEnumerateDeviceExtensionProperties(phy_device, NULL, &count, properties);

    for(size_t i = 0; i < ARR_SIZE(device_extensions); i++)
    {
        bool present = false;
        for(u32 j = 0; j < count; j++)
            if(strncmp(device_extensions[i], properties[j].extensionName, 256) == 0) {
                present = true;
            }
       
        if(!present) {
            free(properties);
            return false;
        }
    }
    
    free(properties);
    return true;
}

static int swap_chain_support(VkPhysicalDevice phy_device, VkSurfaceKHR surface, struct swap_chain_support_details *details)
{
    vkGetPhysicalDeviceSurfaceCapabilitiesKHR(phy_device, surface, &details->capabilities);

    vkGetPhysicalDeviceSurfaceFormatsKHR(phy_device, surface, &details->nformats, NULL);
    if(details->nformats > 0) {
        details->formats = xcalloc(details->nformats, sizeof(*details->formats));
        vkGetPhysicalDeviceSurfaceFormatsKHR(phy_device, surface, &details->nformats, details->formats);
    }

    vkGetPhysicalDeviceSurfacePresentModesKHR(phy_device, surface, &details->npresent_modes, NULL);
    if(details->npresent_modes > 0) {
        details->present_modes = xcalloc(details->npresent_modes, sizeof(*details->present_modes));
        vkGetPhysicalDeviceSurfacePresentModesKHR(phy_device, surface, &details->npresent_modes, details->present_modes);
    }

    return 0;    
}

static void swap_chain_free_support_details(struct swap_chain_support_details *details)
{
    if(details->nformats > 0) free(details->formats);     
    if(details->npresent_modes > 0) free(details->present_modes);
}

static int swap_chain_choose_format(VkSurfaceFormatKHR* formats, u32 nformats, VkSurfaceFormatKHR *format)
{
    for(u32 i = 0; i < nformats; i++)
        if(formats[i].format == VK_FORMAT_B8G8R8A8_SRGB &&
           formats[i].colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR) {
            *format = formats[i];
            return 0;
        }

    *format = formats[0];
    return 0;
}

static int swap_chain_choose_present_mode(VkPresentModeKHR *modes, u32 nmodes, VkPresentModeKHR *mode)
{
    for(u32 i = 0; i < nmodes; i++) 
        if(modes[i] == VK_PRESENT_MODE_MAILBOX_KHR) {
            *mode = modes[i];
            return 0;
        }

    *mode = VK_PRESENT_MODE_FIFO_KHR;
    return 0;    
}

static int swap_chain_get_extent(VkSurfaceCapabilitiesKHR capabilities, VkExtent2D *extent)
{
    if(capabilities.currentExtent.width != UINT32_MAX) {
        *extent = capabilities.currentExtent;
        return 0;
    }
    
    // TODO implement
    err("Not Implemented");
    return 1;    
}

#ifdef DEBUG

static bool instance_has_validation_layers(const char * const *layers, u32 nlayers)
{
    (void)layers;
    (void)nlayers;
    
    // u32 navaliable = 0;
    // vkEnumerateInstanceLayerProperties(&navaliable, NULL);

    // VkLayerProperties *available_layers = xcalloc(navaliable, sizeof(VkLayerProperties));
    // vkEnumerateInstanceLayerProperties(&navaliable, available_layers);
    return true;
}

static int create_debug_messenger(device_t device, struct device_info *info)
{
    (void)info;
    
    VkDebugUtilsMessengerCreateInfoEXT cinfo = {0};
    debug_messenger_populate_info(&cinfo);

    if(CreateDebugUtilsMessengerEXT(device->instance, &cinfo, NULL, &device->debug_messenger) != VK_SUCCESS) {
        info("CreateDebugUtilMessengerEXT: failed");
        return 1;
    }

    return 0;
}

static void debug_messenger_populate_info(VkDebugUtilsMessengerCreateInfoEXT *info)
{
    info->sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_MESSENGER_CREATE_INFO_EXT;
    info->messageSeverity =
        VK_DEBUG_UTILS_MESSAGE_SEVERITY_VERBOSE_BIT_EXT |
        VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT |
        VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT;
    info->messageType =
        VK_DEBUG_UTILS_MESSAGE_TYPE_GENERAL_BIT_EXT |
        VK_DEBUG_UTILS_MESSAGE_TYPE_VALIDATION_BIT_EXT |
        VK_DEBUG_UTILS_MESSAGE_TYPE_PERFORMANCE_BIT_EXT;
    info->pfnUserCallback = debug_messenger_callback;
}

static VKAPI_ATTR VkBool32 VKAPI_CALL debug_messenger_callback(
    VkDebugUtilsMessageSeverityFlagBitsEXT messageSeverity,
    VkDebugUtilsMessageTypeFlagsEXT messageType,
    const VkDebugUtilsMessengerCallbackDataEXT* pCallbackData,
    void* pUserData)
{
    (void)messageSeverity;
    (void)messageType;
    (void)pUserData;
    
    warn("Validation Layer: %s", pCallbackData->pMessage);
    return VK_FALSE;
}

VkResult CreateDebugUtilsMessengerEXT(VkInstance instance, const VkDebugUtilsMessengerCreateInfoEXT* pCreateInfo, const VkAllocationCallbacks* pAllocator, VkDebugUtilsMessengerEXT* pDebugMessenger) {
    PFN_vkCreateDebugUtilsMessengerEXT f = (PFN_vkCreateDebugUtilsMessengerEXT)vkGetInstanceProcAddr(instance, "vkCreateDebugUtilsMessengerEXT");
    if (f != NULL) {
        return f(instance, pCreateInfo, pAllocator, pDebugMessenger);
    } else {
        return VK_ERROR_EXTENSION_NOT_PRESENT;
    }
}

void DestroyDebugUtilsMessengerEXT(VkInstance instance, VkDebugUtilsMessengerEXT debugMessenger, const VkAllocationCallbacks* pAllocator) {
    PFN_vkDestroyDebugUtilsMessengerEXT f = (PFN_vkDestroyDebugUtilsMessengerEXT)vkGetInstanceProcAddr(instance, "vkDestroyDebugUtilsMessengerEXT");
    if (f != NULL) {
        f(instance, debugMessenger, pAllocator);
    }
}

#endif