#include "vkdepth.h"
#include "vktex.h"
#include "main.h"

s32 vkdepth_findSupportedFormat(VkFormat *outFormat, u32 numCandidates, VkFormat *candidates, VkImageTiling tiling, VkFormatFeatureFlags features) {
    VkPhysicalDevice physicalDevice = vgmain_get_physical_device();

    for(u32 i = 0; i < numCandidates; i++) {
        VkFormat format = candidates[i];
        VkFormatProperties props;
        vkGetPhysicalDeviceFormatProperties(physicalDevice, format, &props);

        if (tiling == VK_IMAGE_TILING_LINEAR && (props.linearTilingFeatures & features) == features) {
            *outFormat = format;
            return 0;
        }
        else if (tiling == VK_IMAGE_TILING_OPTIMAL && (props.optimalTilingFeatures & features) == features) {
            *outFormat = format;
            return 0;
        }
    }

    return -1;
}

s32 vkdepth_findDepthFormat(VkFormat *outFormat) {
    VkFormat candidates[3] = {VK_FORMAT_D32_SFLOAT, VK_FORMAT_D32_SFLOAT_S8_UINT, VK_FORMAT_D24_UNORM_S8_UINT};
    return vkdepth_findSupportedFormat(
        outFormat,
        3,
        candidates,
        VK_IMAGE_TILING_OPTIMAL,
        VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT
        );
}

u8 vkdepth_hasStencilComponent(VkFormat format) {
    return format == VK_FORMAT_D32_SFLOAT_S8_UINT || format == VK_FORMAT_D24_UNORM_S8_UINT;
}

s32 vkdepth_createDepthImage(VkImage *depthImage, VkDeviceMemory *depthImageMemory, VkImageView *depthImageView, VkCommandPool commandPool) {
    VkExtent2D swapChainExtent = vgmain_get_swap_chain_extent();
    VkFormat depthFormat;
    if(vkdepth_findDepthFormat(&depthFormat) != 0) {
        printf("Could not find depth format!\n");
        return -1;
    }

    //I added VK_IMAGE_USAGE_TRANSFER_SRC_BIT for downloading. But will it hurt performance?
    vktex_create_image(depthImage, depthImageMemory, swapChainExtent.width, swapChainExtent.height, depthFormat, VK_IMAGE_TILING_OPTIMAL, VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
    vktex_createImageView(depthImageView, *depthImage, depthFormat, VK_IMAGE_ASPECT_DEPTH_BIT);

    vktex_transitionImageLayoutDepth(*depthImage, depthFormat, VK_IMAGE_LAYOUT_UNDEFINED,
                                VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
                                //VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
                                commandPool);

    return 0;
}

s32 vkdepth_get_depth_data(r32 *outData) {
    VkDevice device = vgmain_get_device();
    VkQueue graphicsQueue = vgmain_get_graphics_queue();
    VkImage depthImage = vgmain_get_depth_image();

    VkExtent2D extent = vgmain_get_swap_chain_extent();
    VkDeviceSize imageSize = extent.width * extent.height * 4;

    //TODO create this buffer one time
    VkBuffer stagingBuffer;
    VkDeviceMemory stagingBufferMemory;

    //create staging buffer as src for image
    vmem_createBuffer(&stagingBuffer, &stagingBufferMemory, imageSize, VK_BUFFER_USAGE_TRANSFER_DST_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

    VkCommandPool commandPool = vgmain_get_prim_command_pool();

    VkCommandBuffer commandBuffer = vkcmd_beginSingleTimeCommands(commandPool, device);

    VkBufferImageCopy region = {};
    region.bufferOffset = 0;
    region.bufferRowLength = 0;
    region.bufferImageHeight = 0;
    region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_DEPTH_BIT;
    region.imageSubresource.mipLevel = 0;
    region.imageSubresource.baseArrayLayer = 0;
    region.imageSubresource.layerCount = 1;
    region.imageOffset = (VkOffset3D){0, 0, 0};
    region.imageExtent = (VkExtent3D){
        extent.width,
        extent.height,
        1
    };

    //TODO does it work, transfering optimal layout to host coherent buffer?
    vkCmdCopyImageToBuffer(commandBuffer, depthImage,
                           VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
                           //VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
                           stagingBuffer, 1, &region);

    vkcmd_endSingleTimeCommands(commandBuffer, commandPool, device, graphicsQueue);

    //TODO do not copy, just use data as result
    void *data;
    vkMapMemory(device, stagingBufferMemory, 0, imageSize, 0, &data);
    //memcpy(outData, data, imageSize);

    u32 imgX = 0;
    u32 imgY = 0;
    r32 *dataf = (r32*)data;
    for(u32 y = 0; y < SCREEN_HEIGHT/2; y += 1, imgY += 2) {
        imgX = 0;
        for(u32 x = 0; x < 2 * SCREEN_WIDTH; x += 4, imgX += 2) {
            outData[y*2*SCREEN_WIDTH + x + 0] = dataf[imgY     * SCREEN_WIDTH + imgX + 0];
            outData[y*2*SCREEN_WIDTH + x + 1] = dataf[imgY     * SCREEN_WIDTH + imgX + 1];
            outData[y*2*SCREEN_WIDTH + x + 2] = dataf[(imgY+1) * SCREEN_WIDTH + imgX + 0];
            outData[y*2*SCREEN_WIDTH + x + 3] = dataf[(imgY+1) * SCREEN_WIDTH + imgX + 1];
        }
    }

    vkUnmapMemory(device, stagingBufferMemory);

    vkDestroyBuffer(device, stagingBuffer, nullptr);
    vkFreeMemory(device, stagingBufferMemory, nullptr);
}