#include "core/VulkanWrapper.h"

#include <core/Logger.h>
#include <core/Utility.h>
#include <stdio.h>

#include "core/internal/GLFW.h"

typedef struct {
    VkFormat format;
    VkSwapchainKHR handle;
    u32 amount;
    VkImage* images;
    VkImageView* imageViews;
    VkCommandPool commandPool;
    VkCommandBuffer* commandsBuffers;
    VkFramebuffer* framebuffers;
} Swapchain;

typedef struct {
    // basic setup
    VkInstance instance;
    VkSurfaceKHR surface;
    VkDebugUtilsMessengerEXT debugMessenger;
    VkDebugReportCallbackEXT debugReportCallback;
    // render setup
    u32 graphicsFamily;
    VkPhysicalDevice physicalDevice;
    VkDevice device;
    VkQueue graphicsQueue;
    VkSemaphore semaphore;
    VkSemaphore renderSemaphore;
    Swapchain swapchain;
    VkRenderPass renderPass;
    u32 width;
    u32 height;
} VulkanData;
static VulkanData vk = {0};

#define VK_ERROR_CASE(error)                                                   \
    case error: return #error

static const char* getVulkanErrorString(VkResult r) {
    switch(r) {
        VK_ERROR_CASE(VK_ERROR_OUT_OF_HOST_MEMORY);
        VK_ERROR_CASE(VK_ERROR_OUT_OF_DEVICE_MEMORY);
        VK_ERROR_CASE(VK_ERROR_INITIALIZATION_FAILED);
        VK_ERROR_CASE(VK_ERROR_LAYER_NOT_PRESENT);
        VK_ERROR_CASE(VK_ERROR_EXTENSION_NOT_PRESENT);
        VK_ERROR_CASE(VK_ERROR_INCOMPATIBLE_DRIVER);
        default: return "unknown";
    }
}

#define VK_ASSERT(a)                                                           \
    do {                                                                       \
        VkResult vkResult = (a);                                               \
        if(vkResult != VK_SUCCESS) {                                           \
            printf("Vulkan error at [%s:%d]: %s\n",                            \
                   getShortFileName(__FILE__), __LINE__,                       \
                   getVulkanErrorString(vkResult));                            \
            return true;                                                       \
        }                                                                      \
    } while(false)

static PFN_vkVoidFunction getVulkanFunction(const char* name) {
    return vkGetInstanceProcAddr(vk.instance, name);
}

#define GET_VULKAN_FUNCTION(name) ((PFN_##name)getVulkanFunction(#name))

static bool initInstance() {
    VkApplicationInfo i = {.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO,
                           .pApplicationName = "Vulkan",
                           .applicationVersion = VK_MAKE_VERSION(1, 0, 0),
                           .pEngineName = "Kajetan",
                           .engineVersion = VK_MAKE_VERSION(0, 0, 1),
                           .apiVersion = VK_API_VERSION_1_1};
    VkInstanceCreateInfo ci = {
        .sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
        .pApplicationInfo = &i,
        .enabledLayerCount = 1,
        .ppEnabledLayerNames = (const char*[]){"VK_LAYER_KHRONOS_validation"},
        .enabledExtensionCount = 4,
        .ppEnabledExtensionNames =
            (const char*[]){"VK_KHR_surface", "VK_KHR_wayland_surface",
                            VK_EXT_DEBUG_UTILS_EXTENSION_NAME,
                            VK_EXT_DEBUG_REPORT_EXTENSION_NAME}};
    VK_ASSERT(vkCreateInstance(&ci, nullptr, &vk.instance));
    return false;
}

static VKAPI_ATTR VkBool32 onVulkanDebugMessenger VKAPI_CALL(
    VkDebugUtilsMessageSeverityFlagBitsEXT, VkDebugUtilsMessageTypeFlagsEXT,
    const VkDebugUtilsMessengerCallbackDataEXT* data, void*) {
    LOG_WARNING("Vulkan validation layer message: %s", data->pMessage);
    return false;
}

static bool initDebugMessenger() {
    VkDebugUtilsMessengerCreateInfoEXT info = {
        .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_MESSENGER_CREATE_INFO_EXT,
        .messageSeverity = VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT |
                           VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT,
        .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,
        .pfnUserCallback = onVulkanDebugMessenger};
    auto f = GET_VULKAN_FUNCTION(vkCreateDebugUtilsMessengerEXT);
    if(f == nullptr) {
        LOG_WARNING("Could not find debug util messenger function");
        return false;
    }
    VK_ASSERT(f(vk.instance, &info, nullptr, &vk.debugMessenger));
    return false;
}

static VKAPI_ATTR VkBool32 onVulkanDebugReport
VKAPI_CALL(VkDebugReportFlagsEXT, VkDebugReportObjectTypeEXT, uint64_t, size_t,
           int32_t, const char* pLayerPrefix, const char* pMessage, void*) {
    LOG_WARNING("Vulkan debug message '%s': %s", pLayerPrefix, pMessage);
    return false;
}

static bool initDebugReportCallback() {
    VkDebugReportCallbackCreateInfoEXT info = {
        .sType = VK_STRUCTURE_TYPE_DEBUG_REPORT_CREATE_INFO_EXT,
        .flags = VK_DEBUG_REPORT_WARNING_BIT_EXT |
                 VK_DEBUG_REPORT_PERFORMANCE_WARNING_BIT_EXT |
                 VK_DEBUG_REPORT_ERROR_BIT_EXT,
        .pfnCallback = onVulkanDebugReport};
    auto f = GET_VULKAN_FUNCTION(vkCreateDebugReportCallbackEXT);
    if(f == nullptr) {
        LOG_WARNING("Could not find debug report function");
        return false;
    }
    VK_ASSERT(f(vk.instance, &info, nullptr, &vk.debugReportCallback));
    return false;
}

static bool initDebugging() {
    return initDebugMessenger() || initDebugReportCallback();
}

static void destroyDebugMessenger() {
    auto f = GET_VULKAN_FUNCTION(vkDestroyDebugUtilsMessengerEXT);
    if(f == nullptr) {
        LOG_WARNING("Could not find debug util messenger destroy function");
        return;
    }
    f(vk.instance, vk.debugMessenger, nullptr);
}

static void destroyDebugReportCallback() {
    auto f = GET_VULKAN_FUNCTION(vkDestroyDebugReportCallbackEXT);
    if(f == nullptr) {
        LOG_WARNING("Could not find debug report destroy function");
        return;
    }
    f(vk.instance, vk.debugReportCallback, nullptr);
}

static bool choosePhysicalDevice(VkPhysicalDevice* devices, u32 amount) {
    VK_ASSERT(vkEnumeratePhysicalDevices(vk.instance, &amount, devices));
    int bestPoints = 0;
    LOG_INFO("Found %u devices", amount);
    for(u32 i = 0; i < amount; i++) {
        int points = 0;
        VkPhysicalDeviceProperties p;
        vkGetPhysicalDeviceProperties(devices[i], &p);
        if(p.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) {
            points += 100;
        } else if(p.deviceType == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU) {
            points += 50;
        } else if(p.deviceType == VK_PHYSICAL_DEVICE_TYPE_VIRTUAL_GPU) {
            points += 20;
        }
        if(points > bestPoints) {
            bestPoints = points;
            vk.physicalDevice = devices[i];
            LOG_INFO("Best Device: %s", p.deviceName);
        }
    }
    return bestPoints == 0;
}

static bool initPhysicalDevice() {
    u32 deviceCount = 0;
    VK_ASSERT(vkEnumeratePhysicalDevices(vk.instance, &deviceCount, nullptr));
    if(deviceCount == 0) {
        LOG_ERROR("No physical device was found");
        return true;
    }
    VkPhysicalDevice* devices =
        coreAllocate(sizeof(VkPhysicalDevice) * deviceCount);
    bool r = choosePhysicalDevice(devices, deviceCount);
    coreFree(devices);
    return r;
}

static u32 findQueueFamilies(VkPhysicalDevice pd, VkQueueFlags flags) {
    u32 count = 0;
    vkGetPhysicalDeviceQueueFamilyProperties(pd, &count, nullptr);
    VkQueueFamilyProperties* properties =
        coreAllocate(sizeof(VkQueueFamilyProperties) * count);
    vkGetPhysicalDeviceQueueFamilyProperties(pd, &count, properties);
    for(u32 i = 0; i < count; i++) {
        if(properties[i].queueCount != 0 &&
           (properties[i].queueFlags & flags) == flags) {
            coreFree(properties);
            return i;
        }
    }
    coreFree(properties);
    return (u32)-1;
}

static bool initDevice() {
    VkPhysicalDeviceFeatures deviceFeatures = {0};
    vkGetPhysicalDeviceFeatures(vk.physicalDevice, &deviceFeatures);

    vk.graphicsFamily =
        findQueueFamilies(vk.physicalDevice, VK_QUEUE_GRAPHICS_BIT);
    if(vk.graphicsFamily == (u32)-1) {
        LOG_ERROR("Cannot find queue family");
        return true;
    }
    VkDeviceQueueCreateInfo dqInfo = {
        .sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
        .queueFamilyIndex = vk.graphicsFamily,
        .queueCount = 1,
        .pQueuePriorities = (float[]){1.0f}};
    VkDeviceCreateInfo deviceCreateInfo = {
        .sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
        .queueCreateInfoCount = 1,
        .pQueueCreateInfos = &dqInfo,
        .enabledExtensionCount = 1,
        .ppEnabledExtensionNames =
            (const char*[]){VK_KHR_SWAPCHAIN_EXTENSION_NAME},
        .pEnabledFeatures = &deviceFeatures};
    VK_ASSERT(vkCreateDevice(vk.physicalDevice, &deviceCreateInfo, nullptr,
                             &vk.device));
    vkGetDeviceQueue(vk.device, vk.graphicsFamily, 0, &vk.graphicsQueue);
    if(vk.graphicsQueue == VK_NULL_HANDLE) {
        LOG_ERROR("Cannot get device graphics queue");
        return true;
    }
    return false;
}

static bool initSurface() {
    VK_ASSERT(glfwCreateWindowSurface(vk.instance, getWindow(), nullptr,
                                      &vk.surface));
    return false;
}

static u32 getSwapImageCount(const VkSurfaceCapabilitiesKHR* caps) {
    u32 c = caps->minImageCount + 1;
    // according to VkSurfaceCapabilitiesKHR doc:
    // maxImageCount is 0 when there is no strict limit
    if(caps->maxImageCount != 0 && c > caps->maxImageCount) {
        return caps->maxImageCount;
    }
    return c;
}

static bool getSurfaceFormat(VkSurfaceFormatKHR* sf) {
    u32 c = 0;
    VK_ASSERT(vkGetPhysicalDeviceSurfaceFormatsKHR(vk.physicalDevice,
                                                   vk.surface, &c, nullptr));
    VkSurfaceFormatKHR* formats = coreAllocate(sizeof(VkSurfaceFormatKHR) * c);
    VK_ASSERT(vkGetPhysicalDeviceSurfaceFormatsKHR(vk.physicalDevice,
                                                   vk.surface, &c, formats));
    for(u32 i = 0; i < c; i++) {
        if(formats[i].format == VK_FORMAT_B8G8R8A8_UNORM &&
           formats[i].colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR) {
            *sf = formats[i];
            coreFree(formats);
            return false;
        }
    }
    coreFree(formats);
    LOG_ERROR("Could not find any matching surface format");
    return true;
}

static bool getPresentMode(VkPresentModeKHR* pm) {
    u32 c = 0;
    VK_ASSERT(vkGetPhysicalDeviceSurfacePresentModesKHR(
        vk.physicalDevice, vk.surface, &c, nullptr));
    VkPresentModeKHR* modes = coreAllocate(sizeof(VkPresentModeKHR) * c);
    VK_ASSERT(vkGetPhysicalDeviceSurfacePresentModesKHR(vk.physicalDevice,
                                                        vk.surface, &c, modes));
    for(u32 i = 0; i < c; i++) {
        if(modes[i] == VK_PRESENT_MODE_MAILBOX_KHR) {
            *pm = modes[i];
            coreFree(modes);
            return false;
        }
    }
    coreFree(modes);
    LOG_ERROR("Could not find any matching present modes");
    return true;
}

static bool initSwapchain() {
    VkSurfaceCapabilitiesKHR caps = {0};
    VK_ASSERT(vkGetPhysicalDeviceSurfaceCapabilitiesKHR(vk.physicalDevice,
                                                        vk.surface, &caps));
    VkSurfaceFormatKHR sf = {0};
    if(getSurfaceFormat(&sf)) {
        return true;
    }
    vk.swapchain.format = sf.format;
    VkPresentModeKHR pm = {0};
    if(getPresentMode(&pm)) {
        return true;
    }
    VkSwapchainCreateInfoKHR ci = {
        .sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
        .surface = vk.surface,
        .minImageCount = getSwapImageCount(&caps),
        .imageFormat = sf.format,
        .imageColorSpace = sf.colorSpace,
        .imageExtent = {.width = vk.width, .height = vk.height},
        .imageArrayLayers = 1,
        .imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT |
                      VK_IMAGE_USAGE_TRANSFER_DST_BIT,
        .imageSharingMode = VK_SHARING_MODE_EXCLUSIVE,
        .queueFamilyIndexCount = 1,
        .pQueueFamilyIndices = &vk.graphicsFamily,
        .preTransform = caps.currentTransform,
        .compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
        .presentMode = pm,
        .clipped = VK_TRUE,
        .oldSwapchain = VK_NULL_HANDLE};
    VK_ASSERT(
        vkCreateSwapchainKHR(vk.device, &ci, nullptr, &vk.swapchain.handle));
    return false;
}

static bool createImageView(VkImage image, VkImageView* view) {
    VkImageViewCreateInfo info = {
        .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
        .image = image,
        .viewType = VK_IMAGE_VIEW_TYPE_2D,
        .format = VK_FORMAT_B8G8R8A8_UNORM,
        .subresourceRange = {.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
                             .baseMipLevel = 0,
                             .levelCount = 1,
                             .baseArrayLayer = 0,
                             .layerCount = 1}};
    VK_ASSERT(vkCreateImageView(vk.device, &info, nullptr, view));
    return false;
}

static void destroySwapchainImages() {
    for(size_t x = 0; x < vk.swapchain.amount; x++) {
        vkDestroyImageView(vk.device, vk.swapchain.imageViews[x], nullptr);
    }
    coreFree(vk.swapchain.images);
    coreFree(vk.swapchain.imageViews);
}

static bool initSwapchainImages() {
    u32 c = 0;
    VK_ASSERT(
        vkGetSwapchainImagesKHR(vk.device, vk.swapchain.handle, &c, nullptr));
    LOG_INFO("Found %u images", c);
    vk.swapchain.amount = c;
    vk.swapchain.images = coreAllocate(sizeof(VkImage) * c);
    vk.swapchain.imageViews = coreZeroAllocate(sizeof(VkImageView) * c);
    VK_ASSERT(vkGetSwapchainImagesKHR(vk.device, vk.swapchain.handle, &c,
                                      vk.swapchain.images));
    for(u32 x = 0; x < c; x++) {
        if(createImageView(vk.swapchain.images[x],
                           vk.swapchain.imageViews + x)) {
            return true;
        }
    }
    return false;
}

static bool initSemaphore(VkSemaphore* s) {
    VkSemaphoreCreateInfo info = {.sType =
                                      VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO};
    VK_ASSERT(vkCreateSemaphore(vk.device, &info, nullptr, s));
    return false;
}

static void destroySemaphore(VkSemaphore s) {
    vkDestroySemaphore(vk.device, s, nullptr);
}

static bool checkPresentationSupport() {
    VkBool32 b = false;
    VK_ASSERT(vkGetPhysicalDeviceSurfaceSupportKHR(
        vk.physicalDevice, vk.graphicsFamily, vk.surface, &b));
    if(!b) {
        LOG_ERROR("Presentation is not supported");
        return true;
    }
    return false;
}

static bool initCommandBuffers() {
    VkCommandPoolCreateInfo info = {
        .sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
        .queueFamilyIndex = vk.graphicsFamily};
    VK_ASSERT(vkCreateCommandPool(vk.device, &info, nullptr,
                                  &vk.swapchain.commandPool));
    vk.swapchain.commandsBuffers =
        coreAllocate(sizeof(VkCommandBuffer) * vk.swapchain.amount);
    VkCommandBufferAllocateInfo a = {
        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
        .commandPool = vk.swapchain.commandPool,
        .level = VK_COMMAND_BUFFER_LEVEL_PRIMARY,
        .commandBufferCount = vk.swapchain.amount};
    VK_ASSERT(
        vkAllocateCommandBuffers(vk.device, &a, vk.swapchain.commandsBuffers));
    return false;
}

static void destroyCommandBuffers() {
    if(vk.swapchain.commandPool != VK_NULL_HANDLE && vk.swapchain.amount > 0) {
        vkFreeCommandBuffers(vk.device, vk.swapchain.commandPool,
                             vk.swapchain.amount, vk.swapchain.commandsBuffers);
    }
    coreFree(vk.swapchain.commandsBuffers);
    vkDestroyCommandPool(vk.device, vk.swapchain.commandPool, nullptr);
}

static bool initRenderPass() {
    VkAttachmentDescription c = {
        .format = vk.swapchain.format,
        .samples = VK_SAMPLE_COUNT_1_BIT,
        .loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR,
        .storeOp = VK_ATTACHMENT_STORE_OP_STORE,
        .stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE,
        .stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE,
        .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED,
        .finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR};
    VkAttachmentReference ca = {
        .attachment = 0, .layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL};
    VkSubpassDescription subpass = {.pipelineBindPoint =
                                        VK_PIPELINE_BIND_POINT_GRAPHICS,
                                    .colorAttachmentCount = 1,
                                    .pColorAttachments = &ca};
    VkRenderPassCreateInfo info = {
        .sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO,
        .attachmentCount = 1,
        .pAttachments = &c,
        .subpassCount = 1,
        .pSubpasses = &subpass};
    VK_ASSERT(vkCreateRenderPass(vk.device, &info, nullptr, &vk.renderPass));
    return false;
}

bool initVulkan() {
    vk.width = 400;
    vk.height = 300;
    return initInstance() || initDebugging() || initSurface() ||
           initPhysicalDevice() || initDevice() || checkPresentationSupport() ||
           initSwapchain() || initSwapchainImages() ||
           initSemaphore(&vk.semaphore) || initSemaphore(&vk.renderSemaphore) ||
           initCommandBuffers() || initRenderPass();
}

static bool fillCommandBuffer(u32 index) {
    VkCommandBufferBeginInfo info = {
        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
        .flags = VK_COMMAND_BUFFER_USAGE_SIMULTANEOUS_USE_BIT};
    VK_ASSERT(vkBeginCommandBuffer(vk.swapchain.commandsBuffers[index], &info));
    return false;
}

void renderVulkan() {
    for(u32 i = 0; i < vk.swapchain.amount; i++) {
        (void)fillCommandBuffer;
        // fillCommandBuffer(i);
    }
}

void destroyVulkan(void) {
    if(vk.device != VK_NULL_HANDLE) {
        vkDestroyRenderPass(vk.device, vk.renderPass, nullptr);
        destroyCommandBuffers();
        destroySemaphore(vk.semaphore);
        destroySemaphore(vk.renderSemaphore);
        destroySwapchainImages();
        vkDestroySwapchainKHR(vk.device, vk.swapchain.handle, nullptr);
    }
    vkDestroyDevice(vk.device, nullptr);
    if(vk.instance != VK_NULL_HANDLE) {
        vkDestroySurfaceKHR(vk.instance, vk.surface, nullptr);
        destroyDebugMessenger();
        destroyDebugReportCallback();
    }
    vkDestroyInstance(vk.instance, nullptr);
}