// VulkanWindow: physical device selection, swapchain, MSAA target and framebuffers.
#include "vulkan_window.h"

#include "host_util.h"

#include <algorithm>
#include <cstdlib>
#include <cstring>
#include <stdexcept>

namespace mt_host {

std::uint32_t VulkanWindow::find_memory_type(std::uint32_t type_bits, VkMemoryPropertyFlags flags) const {
    VkPhysicalDeviceMemoryProperties properties{};
    vkGetPhysicalDeviceMemoryProperties(physical_, &properties);
    for (std::uint32_t i = 0; i < properties.memoryTypeCount; ++i)
        if ((type_bits & (1u << i)) && (properties.memoryTypes[i].propertyFlags & flags) == flags)
            return i;
    throw std::runtime_error("no matching Vulkan memory type");
}

void VulkanWindow::select_device() {
    std::uint32_t count = 0;
    check(vkEnumeratePhysicalDevices(instance_, &count, nullptr), "vkEnumeratePhysicalDevices count");
    if (!count) throw std::runtime_error("no Vulkan physical device; set VK_ICD_FILENAMES for MoltenVK");
    std::vector<VkPhysicalDevice> candidates(count);
    check(vkEnumeratePhysicalDevices(instance_, &count, candidates.data()), "vkEnumeratePhysicalDevices");
    for (auto candidate : candidates) {
        std::uint32_t family_count = 0;
        vkGetPhysicalDeviceQueueFamilyProperties(candidate, &family_count, nullptr);
        std::vector<VkQueueFamilyProperties> families(family_count);
        vkGetPhysicalDeviceQueueFamilyProperties(candidate, &family_count, families.data());
        for (std::uint32_t i = 0; i < family_count; ++i) {
            VkBool32 present = VK_FALSE;
            check(vkGetPhysicalDeviceSurfaceSupportKHR(candidate, i, surface_, &present), "surface support");
            if ((families[i].queueFlags & VK_QUEUE_GRAPHICS_BIT) && present) {
                physical_ = candidate; queue_family_ = i; break;
            }
        }
        if (physical_) break;
    }
    if (!physical_) throw std::runtime_error("no graphics/present queue family");
    VkPhysicalDeviceProperties properties{};
    vkGetPhysicalDeviceProperties(physical_, &properties);
    device_name_ = properties.deviceName;
    timestamp_period_ns_ = properties.limits.timestampPeriod;
    VkPhysicalDeviceFeatures supported_features{};
    vkGetPhysicalDeviceFeatures(physical_, &supported_features);
    VkPhysicalDeviceFeatures enabled_features{};
    if (supported_features.samplerAnisotropy) {
        enabled_features.samplerAnisotropy = VK_TRUE;
        max_anisotropy_ = std::min(16.0f, properties.limits.maxSamplerAnisotropy);
    }
    // Multisampled 3D edges; MT_MSAA=1 keeps the single-sample path, 2/4/8 request a count.
    {
        int wanted = 4;
        if (const char* env = std::getenv("MT_MSAA")) wanted = std::max(1, std::atoi(env));
        const VkSampleCountFlags supported =
            properties.limits.framebufferColorSampleCounts & properties.limits.framebufferDepthSampleCounts;
        samples_ = VK_SAMPLE_COUNT_1_BIT;
        for (VkSampleCountFlagBits c : {VK_SAMPLE_COUNT_8_BIT, VK_SAMPLE_COUNT_4_BIT, VK_SAMPLE_COUNT_2_BIT})
            if (int(c) <= wanted && (supported & c)) { samples_ = c; break; }
    }
    std::uint32_t extension_count = 0;
    check(vkEnumerateDeviceExtensionProperties(physical_, nullptr, &extension_count, nullptr), "device extensions count");
    std::vector<VkExtensionProperties> available(extension_count);
    check(vkEnumerateDeviceExtensionProperties(physical_, nullptr, &extension_count, available.data()), "device extensions");
    std::vector<const char*> extensions{VK_KHR_SWAPCHAIN_EXTENSION_NAME};
    constexpr const char* portability_subset = "VK_KHR_portability_subset";
    for (const auto& extension : available)
        if (std::strcmp(extension.extensionName, portability_subset) == 0)
            extensions.push_back(portability_subset);
    const float priority = 1.0f;
    VkDeviceQueueCreateInfo queue_info{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO};
    queue_info.queueFamilyIndex = queue_family_; queue_info.queueCount = 1; queue_info.pQueuePriorities = &priority;
    VkDeviceCreateInfo device_info{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO};
    device_info.queueCreateInfoCount = 1; device_info.pQueueCreateInfos = &queue_info;
    device_info.enabledExtensionCount = static_cast<std::uint32_t>(extensions.size());
    device_info.ppEnabledExtensionNames = extensions.data();
    device_info.pEnabledFeatures = &enabled_features;
    check(vkCreateDevice(physical_, &device_info, nullptr, &device_), "vkCreateDevice");
    vkGetDeviceQueue(device_, queue_family_, 0, &queue_);
}

void VulkanWindow::create_swapchain() {
    VkSurfaceCapabilitiesKHR capabilities{};
    check(vkGetPhysicalDeviceSurfaceCapabilitiesKHR(physical_, surface_, &capabilities), "surface capabilities");
    std::uint32_t count = 0;
    check(vkGetPhysicalDeviceSurfaceFormatsKHR(physical_, surface_, &count, nullptr), "surface formats count");
    if (!count) throw std::runtime_error("surface has no formats");
    std::vector<VkSurfaceFormatKHR> formats(count);
    check(vkGetPhysicalDeviceSurfaceFormatsKHR(physical_, surface_, &count, formats.data()), "surface formats");
    VkSurfaceFormatKHR format = formats.front();
    for (auto candidate : formats) if (candidate.format == VK_FORMAT_B8G8R8A8_UNORM) format = candidate;
    swapchain_format_ = format.format;
    int width = 0, height = 0;
    SDL_GetWindowSizeInPixels(window_, &width, &height);
    if (capabilities.currentExtent.width != UINT32_MAX) extent_ = capabilities.currentExtent;
    else {
        extent_.width = std::clamp(std::uint32_t(width), capabilities.minImageExtent.width, capabilities.maxImageExtent.width);
        extent_.height = std::clamp(std::uint32_t(height), capabilities.minImageExtent.height, capabilities.maxImageExtent.height);
    }

    present_mode_ = VK_PRESENT_MODE_FIFO_KHR;
    if (!vsync_) {
        std::uint32_t mode_count = 0;
        check(vkGetPhysicalDeviceSurfacePresentModesKHR(physical_, surface_, &mode_count, nullptr), "present modes count");
        std::vector<VkPresentModeKHR> modes(mode_count);
        if (mode_count)
            check(vkGetPhysicalDeviceSurfacePresentModesKHR(physical_, surface_, &mode_count, modes.data()), "present modes");
        if (std::find(modes.begin(), modes.end(), VK_PRESENT_MODE_IMMEDIATE_KHR) != modes.end())
            present_mode_ = VK_PRESENT_MODE_IMMEDIATE_KHR;
        else if (std::find(modes.begin(), modes.end(), VK_PRESENT_MODE_MAILBOX_KHR) != modes.end())
            present_mode_ = VK_PRESENT_MODE_MAILBOX_KHR;
    }

    VkSwapchainCreateInfoKHR info{VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR};
    info.surface = surface_;
    info.minImageCount = std::min(capabilities.minImageCount + 1, capabilities.maxImageCount ? capabilities.maxImageCount : UINT32_MAX);
    info.imageFormat = format.format; info.imageColorSpace = format.colorSpace; info.imageExtent = extent_;
    info.imageArrayLayers = 1; info.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT |
        (capabilities.supportedUsageFlags & VK_IMAGE_USAGE_TRANSFER_SRC_BIT);
    info.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
    // Android surfaces can report a quarter-turn transform even when SDL's window is
    // already landscape. Rendering in SDL window coordinates and requesting that transform
    // again rotates the entire game image on presentation.
    info.preTransform = (capabilities.supportedTransforms & VK_SURFACE_TRANSFORM_IDENTITY_BIT_KHR)
        ? VK_SURFACE_TRANSFORM_IDENTITY_BIT_KHR : capabilities.currentTransform;
#ifdef __ANDROID__
    SDL_Log("swapchain: window=%dx%d extent=%ux%u currentTransform=%u preTransform=%u",
        width, height, extent_.width, extent_.height,
        unsigned(capabilities.currentTransform), unsigned(info.preTransform));
#endif
    info.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR; info.presentMode = present_mode_;
    info.clipped = VK_TRUE;
    check(vkCreateSwapchainKHR(device_, &info, nullptr, &swapchain_), "vkCreateSwapchainKHR");
    check(vkGetSwapchainImagesKHR(device_, swapchain_, &count, nullptr), "swapchain images count");
    std::vector<VkImage> images(count);
    check(vkGetSwapchainImagesKHR(device_, swapchain_, &count, images.data()), "swapchain images");
    swapchain_images_ = images;
    for (auto image : images) {
        VkImageViewCreateInfo view{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO};
        view.image = image; view.viewType = VK_IMAGE_VIEW_TYPE_2D; view.format = swapchain_format_;
        view.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        view.subresourceRange.levelCount = 1; view.subresourceRange.layerCount = 1;
        VkImageView handle = VK_NULL_HANDLE;
        check(vkCreateImageView(device_, &view, nullptr, &handle), "vkCreateImageView");
        image_views_.push_back(handle);
    }
    VkImageCreateInfo depth_info{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO};
    depth_info.imageType = VK_IMAGE_TYPE_2D;
    depth_info.format = VK_FORMAT_D32_SFLOAT;
    depth_info.extent = {extent_.width, extent_.height, 1};
    depth_info.mipLevels = 1; depth_info.arrayLayers = 1;
    depth_info.samples = samples_;
    depth_info.tiling = VK_IMAGE_TILING_OPTIMAL;
    depth_info.usage = VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT;
    depth_info.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    check(vkCreateImage(device_, &depth_info, nullptr, &depth_image_), "vkCreateImage depth");
    VkMemoryRequirements depth_reqs{};
    vkGetImageMemoryRequirements(device_, depth_image_, &depth_reqs);
    VkMemoryAllocateInfo depth_alloc{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO};
    depth_alloc.allocationSize = depth_reqs.size;
    depth_alloc.memoryTypeIndex = find_memory_type(depth_reqs.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
    check(vkAllocateMemory(device_, &depth_alloc, nullptr, &depth_memory_), "vkAllocateMemory depth");
    check(vkBindImageMemory(device_, depth_image_, depth_memory_, 0), "vkBindImageMemory depth");
    VkImageViewCreateInfo depth_view{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO};
    depth_view.image = depth_image_; depth_view.viewType = VK_IMAGE_VIEW_TYPE_2D;
    depth_view.format = VK_FORMAT_D32_SFLOAT;
    depth_view.subresourceRange.aspectMask = VK_IMAGE_ASPECT_DEPTH_BIT;
    depth_view.subresourceRange.levelCount = 1; depth_view.subresourceRange.layerCount = 1;
    check(vkCreateImageView(device_, &depth_view, nullptr, &depth_view_), "vkCreateImageView depth");
    if (samples_ != VK_SAMPLE_COUNT_1_BIT) create_msaa_color();
}

void VulkanWindow::recreate_swapchain() {
    int w = 0, h = 0;
    SDL_GetWindowSizeInPixels(window_, &w, &h);
    if (w <= 0 || h <= 0) return;
    vkDeviceWaitIdle(device_);
    for (auto& entry : pipelines_) vkDestroyPipeline(device_, entry.second, nullptr);
    pipelines_.clear();
    for (auto fb : framebuffers_) vkDestroyFramebuffer(device_, fb, nullptr);
    framebuffers_.clear();
    if (depth_view_) { vkDestroyImageView(device_, depth_view_, nullptr); depth_view_ = VK_NULL_HANDLE; }
    if (depth_image_) { vkDestroyImage(device_, depth_image_, nullptr); depth_image_ = VK_NULL_HANDLE; }
    if (depth_memory_) { vkFreeMemory(device_, depth_memory_, nullptr); depth_memory_ = VK_NULL_HANDLE; }
    destroy_msaa_color();
    for (auto view : image_views_) vkDestroyImageView(device_, view, nullptr);
    image_views_.clear();
    if (swapchain_) { vkDestroySwapchainKHR(device_, swapchain_, nullptr); swapchain_ = VK_NULL_HANDLE; }
    create_swapchain();
    create_framebuffers();
    get_pipeline(0, 2, 0);
}

void VulkanWindow::create_msaa_color() {
    VkImageCreateInfo info{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO};
    info.imageType = VK_IMAGE_TYPE_2D;
    info.format = swapchain_format_;
    info.extent = {extent_.width, extent_.height, 1};
    info.mipLevels = 1; info.arrayLayers = 1;
    info.samples = samples_;
    info.tiling = VK_IMAGE_TILING_OPTIMAL;
    info.usage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_TRANSIENT_ATTACHMENT_BIT;
    info.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    check(vkCreateImage(device_, &info, nullptr, &msaa_image_), "vkCreateImage msaa color");
    VkMemoryRequirements reqs{};
    vkGetImageMemoryRequirements(device_, msaa_image_, &reqs);
    VkMemoryAllocateInfo alloc{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO};
    alloc.allocationSize = reqs.size;
    try {
        alloc.memoryTypeIndex = find_memory_type(reqs.memoryTypeBits, VK_MEMORY_PROPERTY_LAZILY_ALLOCATED_BIT);
    } catch (const std::runtime_error&) {
        alloc.memoryTypeIndex = find_memory_type(reqs.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
    }
    check(vkAllocateMemory(device_, &alloc, nullptr, &msaa_memory_), "vkAllocateMemory msaa color");
    check(vkBindImageMemory(device_, msaa_image_, msaa_memory_, 0), "vkBindImageMemory msaa color");
    VkImageViewCreateInfo view{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO};
    view.image = msaa_image_; view.viewType = VK_IMAGE_VIEW_TYPE_2D; view.format = swapchain_format_;
    view.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    view.subresourceRange.levelCount = 1; view.subresourceRange.layerCount = 1;
    check(vkCreateImageView(device_, &view, nullptr, &msaa_view_), "vkCreateImageView msaa color");
}

void VulkanWindow::destroy_msaa_color() {
    if (msaa_view_) { vkDestroyImageView(device_, msaa_view_, nullptr); msaa_view_ = VK_NULL_HANDLE; }
    if (msaa_image_) { vkDestroyImage(device_, msaa_image_, nullptr); msaa_image_ = VK_NULL_HANDLE; }
    if (msaa_memory_) { vkFreeMemory(device_, msaa_memory_, nullptr); msaa_memory_ = VK_NULL_HANDLE; }
}

void VulkanWindow::create_framebuffers() {
    for (auto view : image_views_) {
        // Single-sample: {swapchain, depth}. MSAA: {msaa colour, depth, swapchain resolve}.
        VkImageView fb_attachments[3] = {view, depth_view_, VK_NULL_HANDLE};
        std::uint32_t count = 2;
        if (samples_ != VK_SAMPLE_COUNT_1_BIT) {
            fb_attachments[0] = msaa_view_;
            fb_attachments[2] = view;
            count = 3;
        }
        VkFramebufferCreateInfo framebuffer{VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO};
        framebuffer.renderPass = pass_; framebuffer.attachmentCount = count; framebuffer.pAttachments = fb_attachments;
        framebuffer.width = extent_.width; framebuffer.height = extent_.height; framebuffer.layers = 1;
        VkFramebuffer handle = VK_NULL_HANDLE;
        check(vkCreateFramebuffer(device_, &framebuffer, nullptr, &handle), "vkCreateFramebuffer");
        framebuffers_.push_back(handle);
    }
}

} // namespace mt_host