//
// Created by Vicente Ferrari Smith on 12.02.26.
//

#define VOLK_IMPLEMENTATION
#include <Volk/volk.h>
#define VMA_IMPLEMENTATION
#include <vma/vk_mem_alloc.h>

#include "vulkan.h"
#include <print>
#include <vector>

#include "../graphics.h"

VkInstance instance{};
VkPhysicalDevice physicalDevice{};
VkDevice device{};
VkQueue graphics_queue{};
uint32_t queueFamily{};

VkSurfaceKHR surface{};
VkDebugUtilsMessengerEXT debugMessenger{};

VmaAllocator allocator{};

constexpr uint32_t MAX_FRAMES_IN_FLIGHT = 2;
constexpr uint32_t MAX_VERTICES_PER_BATCH = 65536;

VkSwapchainKHR swapchain;
VkExtent2D swapchain_extent;
VkSurfaceFormatKHR swapchain_format{
    VK_FORMAT_B8G8R8A8_UNORM,
    VK_COLOR_SPACE_SRGB_NONLINEAR_KHR
};
std::vector<VkSemaphore> renderFinished;

std::vector<VkImage> images;
std::vector<VkImageView> imageViews;
std::vector<VkImageLayout> imageLayouts;

void upload_texture(
    const int w,
    const int h,
    const void* pixels,
    Texture *texture)
{
    VkDeviceSize imageSize = w * h * 4;

    // --- 1. Create Staging Buffer (CPU Visible) ---
    VkBuffer stagingBuffer;
    VmaAllocation stagingAlloc;

    VkBufferCreateInfo stagingBufferInfo = { .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO };
    stagingBufferInfo.size = imageSize;
    stagingBufferInfo.usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT;

    VmaAllocationCreateInfo stagingAllocCreateInfo = {
        .flags = VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT | VMA_ALLOCATION_CREATE_MAPPED_BIT,
        .usage = VMA_MEMORY_USAGE_AUTO,
    };

    VmaAllocationInfo stagingResultInfo;
    vmaCreateBuffer(allocator, &stagingBufferInfo, &stagingAllocCreateInfo, &stagingBuffer, &stagingAlloc, &stagingResultInfo);

    // Copy raw pixels into the mapped memory provided by VMA
    memcpy(stagingResultInfo.pMappedData, pixels, imageSize);

    // --- 2. Create GPU Image (Device Local / Tiled) ---
    VkExtent3D imageExtent = { (uint32_t) w, (uint32_t) h, 1 };

    VkImageCreateInfo imageInfo = {
        .sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
        .imageType = VK_IMAGE_TYPE_2D,
        .format = VK_FORMAT_R8G8B8A8_UNORM,
        .extent = imageExtent,
        .mipLevels = 1,
        .arrayLayers = 1,
        .samples = VK_SAMPLE_COUNT_1_BIT,
        .tiling = VK_IMAGE_TILING_OPTIMAL,
        .usage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
        .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED
    };

    VmaAllocationCreateInfo imageAllocCreateInfo = {
        .usage = VMA_MEMORY_USAGE_AUTO,
        .priority = 1.0f,
    };

    vmaCreateImage(allocator, &imageInfo, &imageAllocCreateInfo, image, allocation, nullptr);

    // --- 3. The Transfer ---
    immediate_submit([&](VkCommandBuffer cmd) {
        // Transition image from UNDEFINED to TRANSFER_DST
        transition_image_layout(cmd, *image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);

        VkBufferImageCopy copyRegion = {};
        copyRegion.bufferOffset = 0;
        copyRegion.bufferRowLength = 0;
        copyRegion.bufferImageHeight = 0;
        copyRegion.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        copyRegion.imageSubresource.mipLevel = 0;
        copyRegion.imageSubresource.baseArrayLayer = 0;
        copyRegion.imageSubresource.layerCount = 1;
        copyRegion.imageExtent = imageExtent;

        vkCmdCopyBufferToImage(cmd, stagingBuffer, *image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &copyRegion);

        // Transition image from TRANSFER_DST to SHADER_READ_ONLY
        transition_image_layout(cmd, *image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
    });

    // Clean up temporary staging resources
    vmaDestroyBuffer(allocator, stagingBuffer, stagingAlloc);

    // --- 4. Finalize Handles ---
    *view = create_image_view(*image, imageInfo.format);

    // Register in your Bindless Array (Set 0, Binding 0, Index N)
    *descriptor_index = nextTextureSlot++;
    update_bindless_slot(*descriptor_index, *view, defaultSampler);
}

void graphics_init() {
    createInstance(window);
    createSurface(window);
    createDevice();

    createSwapchain(window);

    Renderer renderer(window);
}

void graphics_deinit() {
    vkDeviceWaitIdle(device);
}

void begin_frame() {

}

void end_frame() {

}

VKAPI_ATTR VkBool32 VKAPI_CALL debugCallback(
    VkDebugUtilsMessageSeverityFlagBitsEXT severity,
    VkDebugUtilsMessageTypeFlagsEXT type,
    const VkDebugUtilsMessengerCallbackDataEXT* callbackData,
    void* userData) {
    const char* severityStr = "UNKNOWN";
    if (severity & VK_DEBUG_UTILS_MESSAGE_SEVERITY_VERBOSE_BIT_EXT)
        severityStr = "VERBOSE";
    else if (severity & VK_DEBUG_UTILS_MESSAGE_SEVERITY_INFO_BIT_EXT)
        severityStr = "INFO";
    else if (severity & VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT)
        severityStr = "WARNING";
    else if (severity & VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT)
        severityStr = "ERROR";

    const char* typeStr = "";
    if (type & VK_DEBUG_UTILS_MESSAGE_TYPE_GENERAL_BIT_EXT)
        typeStr = "GENERAL";
    else if (type & VK_DEBUG_UTILS_MESSAGE_TYPE_VALIDATION_BIT_EXT)
        typeStr = "VALIDATION";
    else if (type & VK_DEBUG_UTILS_MESSAGE_TYPE_PERFORMANCE_BIT_EXT)
        typeStr = "PERFORMANCE";

    std::fprintf(stderr, "\n[VULKAN %s][%s]\n%s\n",
                 severityStr,
                 typeStr,
                 callbackData->pMessage);

    // Print associated objects (very useful on MoltenVK)
    for (uint32_t i = 0; i < callbackData->objectCount; i++) {
        const VkDebugUtilsObjectNameInfoEXT& obj =
            callbackData->pObjects[i];

        std::fprintf(stderr,
                     "    Object %u: type=%d handle=0x%llx name=%s\n",
                     i,
                     obj.objectType,
                     static_cast<unsigned long long>(obj.objectHandle),
                     obj.pObjectName ? obj.pObjectName : "null");
    }

    return VK_FALSE; // Do not abort Vulkan call
}

int createInstance(GLFWwindow *window) {
    volkInitialize();

    uint32_t extensions_count;
    const char** extensions = glfwGetRequiredInstanceExtensions(&extensions_count);
    std::vector<const char*> instance_exts;

    for (uint32_t i = 0; i < extensions_count; i++) {
        instance_exts.push_back(extensions[i]);
    }
    // Append portability enumeration (MANDATORY on macOS)
    instance_exts.push_back(VK_EXT_DEBUG_UTILS_EXTENSION_NAME);
    instance_exts.push_back(VK_KHR_GET_SURFACE_CAPABILITIES_2_EXTENSION_NAME);

    uint32_t availableCount = 0;
    vkEnumerateInstanceExtensionProperties(nullptr, &availableCount, nullptr);

    std::vector<VkExtensionProperties> available(availableCount);
    vkEnumerateInstanceExtensionProperties(nullptr, &availableCount, available.data());

    bool has_portability_extension = false;

    for (const VkExtensionProperties &property : available) {
        if (strcmp(property.extensionName, VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME) == 0) {
            has_portability_extension = true;
            instance_exts.push_back(VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME);
        }
    }

    constexpr VkApplicationInfo app{
        .sType = VK_STRUCTURE_TYPE_APPLICATION_INFO,
        .apiVersion = VK_API_VERSION_1_4,
    };

    const char* layers[] = {
        "VK_LAYER_KHRONOS_validation"
    };

    VkDebugUtilsMessengerCreateInfoEXT dbg{};
    dbg.sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_MESSENGER_CREATE_INFO_EXT;
    dbg.messageSeverity =
        VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT |
        VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT;
    dbg.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;
    dbg.pfnUserCallback = debugCallback;


    VkInstanceCreateInfo ici{
        .sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
        .pNext = &dbg,
        .pApplicationInfo = &app,
        .enabledLayerCount = 1,
        .ppEnabledLayerNames = layers,
        .enabledExtensionCount = static_cast<uint32_t>(instance_exts.size()),
        .ppEnabledExtensionNames = instance_exts.data(),
    };

    if (has_portability_extension) {
        ici.flags = VK_INSTANCE_CREATE_ENUMERATE_PORTABILITY_BIT_KHR;
    }

    VkResult res = vkCreateInstance(&ici, nullptr, &instance);
    if (res != VK_SUCCESS) {
        printf("vkCreateInstance failed: %d\n", res);
        return -1;
    }

    volkLoadInstance(instance);

    VkDebugUtilsMessengerEXT messenger;
    vkCreateDebugUtilsMessengerEXT(instance, &dbg, nullptr, &messenger);


    uint32_t layerCount = 0;
    vkEnumerateInstanceLayerProperties(&layerCount, nullptr);

    return 0;
}

void createSurface(GLFWwindow *window) {
    glfwCreateWindowSurface(instance, window, nullptr, &surface);
}

void createDevice() {
    uint32_t deviceCount = 0;
    uint32_t res = vkEnumeratePhysicalDevices(instance, &deviceCount, nullptr);
    if (res != VK_SUCCESS || deviceCount == 0) {
        printf("vkEnumeratePhysicalDevices failed or found no devices (%d)\n", res);
    }
    std::vector<VkPhysicalDevice> devices(deviceCount);
    vkEnumeratePhysicalDevices(instance, &deviceCount, devices.data());

    printf("Found %u physical device(s)\n", deviceCount);
    for (uint32_t i = 0; i < deviceCount; i++) {
        VkPhysicalDeviceProperties props{};
        vkGetPhysicalDeviceProperties(devices[i], &props);

        printf("Device %u:\n", i);
        printf("  Name: %s\n", props.deviceName);
        printf("  Type: %u\n", props.deviceType);
        printf("  API version: %u.%u.%u\n",
               VK_VERSION_MAJOR(props.apiVersion),
               VK_VERSION_MINOR(props.apiVersion),
               VK_VERSION_PATCH(props.apiVersion));
    }

    physicalDevice = devices[0];

    queueFamily = 0;
    uint32_t qCount = 0;
    vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &qCount, nullptr);
    std::vector<VkQueueFamilyProperties> qprops(qCount);
    vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &qCount, qprops.data());

    for (uint32_t i = 0; i < qCount; i++) {
        VkBool32 present = VK_FALSE;
        vkGetPhysicalDeviceSurfaceSupportKHR(physicalDevice, i, surface, &present);
        if ((qprops[i].queueFlags & VK_QUEUE_GRAPHICS_BIT) && present) {
            queueFamily = i;
            break;
        }
    }

    if (int r = glfwGetPhysicalDevicePresentationSupport(instance, physicalDevice, queueFamily) == GLFW_FALSE) {
        return;
    }

    std::println("{}", queueFamily);

    std::vector<const char*> devExts = {
        VK_KHR_SWAPCHAIN_EXTENSION_NAME,
        VK_EXT_EXTENDED_DYNAMIC_STATE_3_EXTENSION_NAME,
    };

    uint32_t extensionCount = 0;
    vkEnumerateDeviceExtensionProperties(physicalDevice, nullptr, &extensionCount, nullptr);

    std::vector<VkExtensionProperties> availableExtensions(extensionCount);
    vkEnumerateDeviceExtensionProperties(physicalDevice, nullptr, &extensionCount, availableExtensions.data());

    for (const auto& ext : availableExtensions) {
        if (strcmp(ext.extensionName, "VK_KHR_portability_subset") == 0) {
            devExts.push_back("VK_KHR_portability_subset");
            break;
        }
    }

    float prio = 1.0f;
    VkDeviceQueueCreateInfo qci{
        .sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
        .queueFamilyIndex = queueFamily,
        .queueCount = 1,
        .pQueuePriorities = &prio
    };

    VkPhysicalDeviceVulkan14Features enabledVk14Features{
    .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_4_FEATURES,
    .pNext = nullptr,
    .globalPriorityQuery = false,
    .shaderSubgroupRotate = false,
    .shaderSubgroupRotateClustered = false,
    .shaderFloatControls2 = false,
    .shaderExpectAssume = false,
    .rectangularLines = false,
    .bresenhamLines = false,
    .smoothLines = false,
    .stippledRectangularLines = false,
    .stippledBresenhamLines = false,
    .stippledSmoothLines = false,
    .vertexAttributeInstanceRateDivisor = false,
    .vertexAttributeInstanceRateZeroDivisor = false,
    .indexTypeUint8 = false,
    .dynamicRenderingLocalRead = false,
    .maintenance5 = false,
    .maintenance6 = false,
    .pipelineProtectedAccess = false,
    .pipelineRobustness = false,
    .hostImageCopy = false,
    .pushDescriptor = false,
    };

    VkPhysicalDeviceVulkan13Features enabledVk13Features{
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES,
        .pNext = &enabledVk14Features,
        .robustImageAccess = false,
        .inlineUniformBlock = false,
        .descriptorBindingInlineUniformBlockUpdateAfterBind = false,
        .pipelineCreationCacheControl = false,
        .privateData = false,
        .shaderDemoteToHelperInvocation = false,
        .shaderTerminateInvocation = false,
        .subgroupSizeControl = false,
        .computeFullSubgroups = false,
        .synchronization2 = true,
        .textureCompressionASTC_HDR = false,
        .shaderZeroInitializeWorkgroupMemory = false,
        .dynamicRendering = true,
        .shaderIntegerDotProduct = false,
        .maintenance4 = false
    };
    VkPhysicalDeviceVulkan12Features enabledVk12Features{
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES,
        .pNext = &enabledVk13Features,
        .descriptorIndexing = true,
        .shaderSampledImageArrayNonUniformIndexing = true,
        .descriptorBindingSampledImageUpdateAfterBind = true,
        .descriptorBindingPartiallyBound = true,
        .descriptorBindingVariableDescriptorCount = true,
        .runtimeDescriptorArray = true,
        .bufferDeviceAddress = true,
    };

    VkPhysicalDeviceVulkan11Features enabledVk11Features{
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES,
        .pNext = &enabledVk12Features,
        .shaderDrawParameters = true,
    };

    const VkPhysicalDeviceFeatures enabledVk10Features{
        .samplerAnisotropy = true,
    };

    const VkDeviceCreateInfo dci{
        .sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
        .pNext = &enabledVk11Features,
        .queueCreateInfoCount = 1,
        .pQueueCreateInfos = &qci,
        .enabledExtensionCount = (uint32_t)devExts.size(),
        .ppEnabledExtensionNames = devExts.data(),
        .pEnabledFeatures = &enabledVk10Features,
    };

    vkCreateDevice(physicalDevice, &dci, nullptr, &device);

    volkLoadDevice(device);

    vkGetDeviceQueue(device, queueFamily, 0, &graphics_queue);

    VmaAllocatorCreateInfo allocatorCreateInfo = {
        .flags = 0,
        .physicalDevice = physicalDevice,
        .device = device,
        .preferredLargeHeapBlockSize = 0,
        .pAllocationCallbacks = nullptr,
        .pDeviceMemoryCallbacks = nullptr,
        .pHeapSizeLimit = nullptr,
        .pVulkanFunctions = nullptr,
        .instance = instance,
        .vulkanApiVersion = VK_API_VERSION_1_4,
        .pTypeExternalMemoryHandleTypes = nullptr,
    };

    VmaVulkanFunctions vulkan_functions = {};
    res = vmaImportVulkanFunctionsFromVolk(&allocatorCreateInfo, &vulkan_functions);

    allocatorCreateInfo.pVulkanFunctions = &vulkan_functions;

    res = vmaCreateAllocator(&allocatorCreateInfo, &allocator);
}

void createSwapchain(GLFWwindow* window) {
    int fbWidth, fbHeight;
    glfwGetFramebufferSize(window, &fbWidth, &fbHeight);

    swapchain_extent = {
        static_cast<uint32_t>(fbWidth),
        static_cast<uint32_t>(fbHeight)
    };

    VkPhysicalDeviceSurfaceInfo2KHR surfaceInfo{
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SURFACE_INFO_2_KHR,
        .pNext = nullptr,
        .surface = surface,
    };

    VkSurfaceCapabilities2KHR surfCapabilities{
        .sType = VK_STRUCTURE_TYPE_SURFACE_CAPABILITIES_2_KHR,
    };

    vkGetPhysicalDeviceSurfaceCapabilities2KHR(physicalDevice, &surfaceInfo, &surfCapabilities);

    const VkSwapchainCreateInfoKHR sci{
        .sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
        .surface = surface,
        .minImageCount = surfCapabilities.surfaceCapabilities.minImageCount,
        .imageFormat = swapchain_format.format,
        .imageColorSpace = swapchain_format.colorSpace,
        .imageExtent = swapchain_extent,
        .imageArrayLayers = 1,
        .imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
        .imageSharingMode = VK_SHARING_MODE_EXCLUSIVE,
        .preTransform = VK_SURFACE_TRANSFORM_IDENTITY_BIT_KHR,
        .compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
        .presentMode = VK_PRESENT_MODE_FIFO_KHR,
        .clipped = VK_TRUE
    };

    vkCreateSwapchainKHR(device, &sci, nullptr, &swapchain);

    uint32_t imgCount;
    vkGetSwapchainImagesKHR(device, swapchain, &imgCount, nullptr);
    std::print("imgCount: {}", imgCount);
    images = std::vector<VkImage>(imgCount);
    vkGetSwapchainImagesKHR(device, swapchain, &imgCount, images.data());

    imageLayouts = std::vector<VkImageLayout>(
        imgCount,
        VK_IMAGE_LAYOUT_UNDEFINED
    );

    imageViews = std::vector<VkImageView>(imgCount);

    const VkSemaphoreCreateInfo seci{
        .sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO,
    };

    renderFinished.resize(imgCount);

    for (uint32_t i = 0; i < imgCount; i++) {
        VkImageViewCreateInfo ivci{};
        ivci.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
        ivci.image = images[i];
        ivci.viewType = VK_IMAGE_VIEW_TYPE_2D;
        ivci.format = swapchain_format.format;   // must match swapchain format
        ivci.components.r = VK_COMPONENT_SWIZZLE_IDENTITY;
        ivci.components.g = VK_COMPONENT_SWIZZLE_IDENTITY;
        ivci.components.b = VK_COMPONENT_SWIZZLE_IDENTITY;
        ivci.components.a = VK_COMPONENT_SWIZZLE_IDENTITY;
        ivci.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        ivci.subresourceRange.baseMipLevel = 0;
        ivci.subresourceRange.levelCount = 1;
        ivci.subresourceRange.baseArrayLayer = 0;
        ivci.subresourceRange.layerCount = 1;

        vkCreateImageView(device, &ivci, nullptr, &imageViews[i]);

        vkCreateSemaphore(device, &seci, nullptr, &renderFinished[i]);
    }
}
