#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

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

enum {
    QUEUE_FAMILY_COUNT = 2,
    
    QUEUE_FAMILY_GRAPHICS_FAMILY = 1 << 0,
    QUEUE_FAMILY_PRESENT_FAMILY  = 1 << 1,
    
    QUEUE_FAMILY_ALL = QUEUE_FAMILY_GRAPHICS_FAMILY |
                       QUEUE_FAMILY_PRESENT_FAMILY,
};

struct queue_family_idx {
    uint32_t graphics_family;
    uint32_t present_family;
};

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

// if the function has 'has' or 'is', then
// 0 - fail; 1 - success (sometimes -1 - fatal error)

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 is_device_suitable(VkPhysicalDevice phy_device, VkSurfaceKHR surface);

// does not follow the general scheme, see implementation for details
static int find_queue_families(VkPhysicalDevice phy_device, VkSurfaceKHR surface, struct queue_family_idx *queue_family);

#ifdef DEBUG
static int create_debug_messenger(device_t device, struct device_info* info);
static int has_validation_layer_support(const char * const *layers, uint32_t nlayers);
static void populate_debug_messenger_info(VkDebugUtilsMessengerCreateInfoEXT *info);

VkResult CreateDebugUtilsMessengerEXT(
    VkInstance instance,
    const VkDebugUtilsMessengerCreateInfoEXT* pCreateInfo,
    const VkAllocationCallbacks* pAllocator,
    VkDebugUtilsMessengerEXT* pDebugMessenger);
void DestroyDebugUtilsMessengerEXT(
    VkInstance instance,
    VkDebugUtilsMessengerEXT debugMessenger,
    const VkAllocationCallbacks* pAllocator);
static VKAPI_ATTR VkBool32 VKAPI_CALL debug_callback(
    VkDebugUtilsMessageSeverityFlagBitsEXT messageSeverity,
    VkDebugUtilsMessageTypeFlagsEXT messageType,
    const VkDebugUtilsMessengerCallbackDataEXT* pCallbackData,
    void* pUserData);
#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));
   
    CHECK(create_instance);
    CHECK(create_surface);
#ifdef DEBUG
    CHECK(create_debug_messenger);
#endif
    CHECK(create_physical_device);
    CHECK(create_logical_device);
    
    return device;    
fail:
    device_destroy(device);
    return NULL;
}

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

    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)
{
    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(uint32_t i = 0; i < info->ext_count; i++) {
        extensions[i] = (char *)info->extensions[i];
    }
        
    // we just added one more extension
    create_info.enabledExtensionCount = info->ext_count+1;
    create_info.ppEnabledExtensionNames = (const char * const *)extensions;

    // validation layer support
    char *validation_layers[] = {
        "VK_LAYER_KHRONOS_validation"
    };
    
    uint32_t nlayers = ARR_SIZE(validation_layers);

    if(!has_validation_layer_support((const char * const *)validation_layers, nlayers)) {
        err("validation_layer_support: failed");
        goto fail;
    }
        
    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};
    populate_debug_messenger_info(&msg_info);
        
    create_info.pNext = &msg_info;
#endif 

    if(vkCreateInstance(&create_info, NULL, &device->instance) != VK_SUCCESS) {
        err("vkCreateInstance: failed");
        goto fail;
    }
    
#ifdef DEBUG
    free(extensions);
#endif

    return 0;
fail:
    return 1;
}

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;
        
    device->physical_device = VK_NULL_HANDLE;
    
    uint32_t dev_count = 0;
    vkEnumeratePhysicalDevices(device->instance, &dev_count, NULL);

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

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

    for(uint32_t i = 0; i < dev_count; i++) {
        if(is_device_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 fail;
    }
    
    return 0;
fail:
    return 1;
}

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

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

    // queue infos
    VkDeviceQueueCreateInfo *queue_infos = xcalloc(QUEUE_FAMILY_COUNT, sizeof(VkDeviceQueueCreateInfo));
    memset(queue_infos, 0, QUEUE_FAMILY_COUNT * sizeof(VkDeviceQueueCreateInfo));
    
    uint32_t *queue_indices = xcalloc(QUEUE_FAMILY_COUNT, sizeof(uint32_t));
    queue_indices[0] = indices.graphics_family;
    queue_indices[1] = indices.present_family;

    float queue_priority = 1.0f;
    for(int i = 0; i < QUEUE_FAMILY_COUNT; i++)
    {
        queue_infos[i].sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
        queue_infos[i].queueFamilyIndex = queue_indices[i];
        queue_infos[i].queueCount = 1;
        queue_infos[i].pQueuePriorities = &queue_priority;
    }

    // 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 = QUEUE_FAMILY_COUNT;

    create_info.pEnabledFeatures = &device_features;
    
    create_info.enabledExtensionCount = 0;
 
    // 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 fail;
    }

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

    free(queue_infos);
    free(queue_indices);

    return 0;
fail:
    return 1;
}

// returns a bitmask if a given family is set, or not
static int find_queue_families(VkPhysicalDevice phy_device, VkSurfaceKHR surface, struct queue_family_idx *queue_family)
{
    int ret = 0;
    
    uint32_t count = 0;
    vkGetPhysicalDeviceQueueFamilyProperties(phy_device, &count, NULL);

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

    for(uint32_t i = 0; i < count; i++) {
        if(properties[i].queueFlags & VK_QUEUE_COMPUTE_BIT) {
            ret |= QUEUE_FAMILY_GRAPHICS_FAMILY;
            if(queue_family != NULL) {
                queue_family->graphics_family = i;
            }
        }

        VkBool32 present_support = 0;
        vkGetPhysicalDeviceSurfaceSupportKHR(phy_device, i, surface, &present_support);
        if(present_support) {
            ret |= QUEUE_FAMILY_PRESENT_FAMILY;
            if(queue_family != NULL){
                queue_family->present_family = i;
            }
        }
    }

    free(properties);
    return ret;
}

static int is_device_suitable(VkPhysicalDevice phy_device, VkSurfaceKHR surface)
{
    if(find_queue_families(phy_device, surface, NULL) == QUEUE_FAMILY_ALL) {
        return 1;
    }

    return 0;
}


#ifdef DEBUG

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

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

    return 0;
}

static void populate_debug_messenger_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_callback;
}

static int has_validation_layer_support(const char * const *layers, uint32_t nlayers)
{
    (void)layers;
    (void)nlayers;
    
    // uint32_t navaliable = 0;
    // vkEnumerateInstanceLayerProperties(&navaliable, NULL);

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

static VKAPI_ATTR VkBool32 VKAPI_CALL debug_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