/*
 * Copyright 2022 Google LLC
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#include "src/gpu/graphite/vk/VulkanTexture.h"

#include "include/gpu/MutableTextureState.h"
#include "include/gpu/graphite/vk/VulkanGraphiteTypes.h"
#include "include/gpu/vk/VulkanMutableTextureState.h"
#include "include/private/SkLog.h"
#include "src/core/SkCompressedDataUtils.h"
#include "src/core/SkMipmap.h"
#include "src/gpu/DataUtils.h"
#include "src/gpu/graphite/Sampler.h"
#include "src/gpu/graphite/TextureProxy.h"
#include "src/gpu/graphite/task/UploadTask.h"
#include "src/gpu/graphite/vk/VulkanCaps.h"
#include "src/gpu/graphite/vk/VulkanCommandBuffer.h"
#include "src/gpu/graphite/vk/VulkanDescriptorSet.h"
#include "src/gpu/graphite/vk/VulkanFramebuffer.h"
#include "src/gpu/graphite/vk/VulkanGraphiteUtils.h"
#include "src/gpu/graphite/vk/VulkanResourceProvider.h"
#include "src/gpu/graphite/vk/VulkanSharedContext.h"
#include "src/gpu/vk/VulkanMemory.h"
#include "src/gpu/vk/VulkanMutableTextureStatePriv.h"

using namespace skia_private;

namespace skgpu::graphite {

bool VulkanTexture::MakeVkImage(const VulkanSharedContext* sharedContext,
                                SkISize dimensions,
                                const TextureInfo& info,
                                CreatedImageInfo* outInfo) {
    SkASSERT(outInfo);
    const VulkanCaps& caps = sharedContext->vulkanCaps();

    if (dimensions.isEmpty()) {
        SKIA_LOG_E("Tried to create VkImage with empty dimensions.");
        return false;
    }
    if (dimensions.width() > caps.maxTextureSize() ||
        dimensions.height() > caps.maxTextureSize()) {
        SKIA_LOG_E("Tried to create VkImage with too large a size.");
        return false;
    }

    if ((info.isProtected() == Protected::kYes) != caps.protectedSupport()) {
        SKIA_LOG_E("Tried to create %s VkImage in %s Context.",
                    info.isProtected() == Protected::kYes ? "protected" : "unprotected",
                    caps.protectedSupport() ? "protected" : "unprotected");
        return false;
    }

    const auto& vkInfo = TextureInfoPriv::Get<VulkanTextureInfo>(info);

    bool isLinear = vkInfo.fImageTiling == VK_IMAGE_TILING_LINEAR;
    VkImageLayout initialLayout = isLinear ? VK_IMAGE_LAYOUT_PREINITIALIZED
                                           : VK_IMAGE_LAYOUT_UNDEFINED;

    // Create Image
    VkSampleCountFlagBits vkSamples = SampleCountToVkSampleCount(info.sampleCount());

    SkASSERT(!isLinear || vkSamples == VK_SAMPLE_COUNT_1_BIT);
    SkASSERT(info.isProtected() == Protected::kNo ||
             (caps.protectedSupport() && SkToBool(VK_IMAGE_CREATE_PROTECTED_BIT & vkInfo.fFlags)));

    uint32_t numMipLevels = 1;
    if (vkInfo.fMipmapped == Mipmapped::kYes) {
        numMipLevels = SkMipmap::ComputeLevelCount(dimensions) + 1;
    }

    uint32_t width = static_cast<uint32_t>(dimensions.fWidth);
    uint32_t height = static_cast<uint32_t>(dimensions.fHeight);

    const VkImageCreateInfo imageCreateInfo = {
        VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO, // sType
        nullptr,                             // pNext
        vkInfo.fFlags,                       // VkImageCreateFlags
        VK_IMAGE_TYPE_2D,                    // VkImageType
        vkInfo.fFormat,                      // VkFormat
        { width, height, 1 },                // VkExtent3D
        numMipLevels,                        // mipLevels
        1,                                   // arrayLayers
        vkSamples,                           // samples
        vkInfo.fImageTiling,                 // VkImageTiling
        vkInfo.fImageUsageFlags,             // VkImageUsageFlags
        vkInfo.fSharingMode,                 // VkSharingMode
        0,                                   // queueFamilyCount
        nullptr,                             // pQueueFamilyIndices
        initialLayout                        // initialLayout
    };

    auto device = sharedContext->device();

    VkImage image = VK_NULL_HANDLE;
    VkResult result;
    VULKAN_CALL_RESULT(
            sharedContext, result, CreateImage(device, &imageCreateInfo, nullptr, &image));
    if (result != VK_SUCCESS) {
        SKIA_LOG_E("Failed call to vkCreateImage with error: %d", result);
        return false;
    }

    auto allocator = sharedContext->memoryAllocator();
    bool forceDedicatedMemory = caps.shouldAlwaysUseDedicatedImageMemory();
    bool useLazyAllocation =
            SkToBool(vkInfo.fImageUsageFlags & VK_IMAGE_USAGE_TRANSIENT_ATTACHMENT_BIT);

    auto checkResult = [sharedContext](VkResult result) {
        return sharedContext->checkVkResult(result);
    };
    if (!skgpu::VulkanMemory::AllocImageMemory(allocator,
                                               image,
                                               info.isProtected(),
                                               forceDedicatedMemory,
                                               useLazyAllocation,
                                               checkResult,
                                               &outInfo->fMemoryAlloc)) {
        // If lazy memory allocation fails, fallback to attempting to use a regular allocation.
        if (useLazyAllocation &&
            skgpu::VulkanMemory::AllocImageMemory(allocator,
                                                  image,
                                                  info.isProtected(),
                                                  forceDedicatedMemory,
                                                  /*useLazyAllocation=*/false,
                                                  checkResult,
                                                  &outInfo->fMemoryAlloc)) {
            SKIA_LOG_W("Could not allocate lazy image memory; using non-lazy instead.");
            useLazyAllocation = false;
        } else {
            const char* protectednessStr =
                    info.isProtected() == Protected::kYes ? "protected" : "unprotected";
            const char* memoryTypeStr = forceDedicatedMemory ? "dedicated" : "shared";
            SKIA_LOG_E("Failed to allocate %s %s image memory.", protectednessStr, memoryTypeStr);

            VULKAN_CALL(sharedContext->interface(), DestroyImage(device, image, nullptr));
            return false;
        }
    }

    VULKAN_CALL_RESULT(
            sharedContext,
            result,
            BindImageMemory(
                    device, image, outInfo->fMemoryAlloc.fMemory, outInfo->fMemoryAlloc.fOffset));
    if (result != VK_SUCCESS) {
        skgpu::VulkanMemory::FreeImageMemory(allocator, outInfo->fMemoryAlloc);
        VULKAN_CALL(sharedContext->interface(), DestroyImage(device, image, nullptr));
        return false;
    }

    outInfo->fImage = image;
    outInfo->fMutableState = sk_make_sp<MutableTextureState>(
            skgpu::MutableTextureStates::MakeVulkan(initialLayout, VK_QUEUE_FAMILY_IGNORED));
    return true;
}

sk_sp<Texture> VulkanTexture::Make(const VulkanSharedContext* sharedContext,
                                   SkISize dimensions,
                                   const TextureInfo& info,
                                   sk_sp<VulkanYcbcrConversion> ycbcrConversion,
                                   std::string_view label) {
    CreatedImageInfo imageInfo;
    if (!MakeVkImage(sharedContext, dimensions, info, &imageInfo)) {
        return nullptr;
    }

    return sk_sp<Texture>(new VulkanTexture(sharedContext,
                                            dimensions,
                                            info,
                                            std::move(imageInfo.fMutableState),
                                            imageInfo.fImage,
                                            imageInfo.fMemoryAlloc,
                                            Ownership::kOwned,
                                            std::move(ycbcrConversion),
                                            label));
}

sk_sp<Texture> VulkanTexture::MakeWrapped(const VulkanSharedContext* sharedContext,
                                          SkISize dimensions,
                                          const TextureInfo& info,
                                          sk_sp<MutableTextureState> mutableState,
                                          VkImage image,
                                          const VulkanAlloc& alloc,
                                          sk_sp<VulkanYcbcrConversion> ycbcrConversion,
                                          std::string_view label) {
    return sk_sp<Texture>(new VulkanTexture(sharedContext,
                                            dimensions,
                                            info,
                                            std::move(mutableState),
                                            image,
                                            alloc,
                                            Ownership::kWrapped,
                                            std::move(ycbcrConversion),
                                            label));
}

VulkanTexture::~VulkanTexture() {}

void VulkanTexture::setImageLayoutAndQueueIndex(VulkanCommandBuffer* cmdBuffer,
                                                VkImageLayout newLayout,
                                                VkAccessFlags dstAccessMask,
                                                VkPipelineStageFlags dstStageMask,
                                                uint32_t newQueueFamilyIndex) const {
    auto sharedContext = static_cast<const VulkanSharedContext*>(this->sharedContext());
    VkImageLayout currentLayout = this->currentLayout();

    std::optional<VkImageMemoryBarrier> imageMemoryBarrier =
            CreateVulkanImageMemoryBarrier(*sharedContext,
                                           fImage,
                                           this->vulkanTextureInfo(),
                                           currentLayout,
                                           this->currentQueueFamilyIndex(),
                                           newLayout,
                                           dstAccessMask,
                                           newQueueFamilyIndex);
    if (!imageMemoryBarrier) {
        // If we don't get a barrier, that means a barrier doesn't need to be issued so we can skip
        // it. It does not mean creating the barrier failed.
        return;
    }

    VkPipelineStageFlags srcStageMask =
            VkImageLayoutToPipelineSrcStageFlags(currentLayout,
                                                 sharedContext->vulkanCaps());
    cmdBuffer->addImageMemoryBarrier(srcStageMask, dstStageMask, /*byRegion=*/false,
                                     &imageMemoryBarrier.value());

    skgpu::MutableTextureStates::SetVkImageLayout(this->mutableState(),
                                                  imageMemoryBarrier->newLayout);
    skgpu::MutableTextureStates::SetVkQueueFamilyIndex(this->mutableState(),
                                                       imageMemoryBarrier->dstQueueFamilyIndex);
}

namespace {

bool uses_lazy_memory(const VulkanAlloc& alloc) {
    return alloc.fFlags & VulkanAlloc::Flag::kLazilyAllocated_Flag;
}

#ifdef SK_DEBUG
bool has_transient_usage(const TextureInfo& info) {
    const auto& vkInfo = TextureInfoPriv::Get<VulkanTextureInfo>(info);
    return vkInfo.fImageUsageFlags & VK_IMAGE_USAGE_TRANSIENT_ATTACHMENT_BIT;
}
#endif

} // anonymous

VulkanTexture::VulkanTexture(const VulkanSharedContext* sharedContext,
                             SkISize dimensions,
                             const TextureInfo& info,
                             sk_sp<MutableTextureState> mutableState,
                             VkImage image,
                             const VulkanAlloc& alloc,
                             Ownership ownership,
                             sk_sp<VulkanYcbcrConversion> ycbcrConversion,
                             std::string_view label)
        : Texture(sharedContext,
                  dimensions,
                  info,
                  uses_lazy_memory(alloc),
                  std::move(mutableState),
                  ownership,
                  label)
        , fImage(image)
        , fMemoryAlloc(alloc)
        , fYcbcrConversion(std::move(ycbcrConversion)) {
    SkASSERT(!uses_lazy_memory(fMemoryAlloc) || has_transient_usage(info));
    // Update the newly-created underlying GPU object's label to match the Resource's
    this->synchronizeBackendLabel();
}

void VulkanTexture::freeGpuData() {
    // Need to delete any ImageViews first
    fImageViews.clear();

    // If the texture is wrapped we don't own this data
    if (this->ownership() != Ownership::kWrapped) {
        auto sharedContext = static_cast<const VulkanSharedContext*>(this->sharedContext());
        VULKAN_CALL(sharedContext->interface(),
                    DestroyImage(sharedContext->device(), fImage, nullptr));
        skgpu::VulkanMemory::FreeImageMemory(sharedContext->memoryAllocator(), fMemoryAlloc);
    }
}

void VulkanTexture::updateImageLayout(VkImageLayout newLayout) {
    skgpu::MutableTextureStates::SetVkImageLayout(this->mutableState(), newLayout);
}

VkImageLayout VulkanTexture::currentLayout() const {
    return skgpu::MutableTextureStates::GetVkImageLayout(this->mutableState());
}

uint32_t VulkanTexture::currentQueueFamilyIndex() const {
    return skgpu::MutableTextureStates::GetVkQueueFamilyIndex(this->mutableState());
}

const VulkanImageView* VulkanTexture::getImageView(VulkanImageView::Usage usage) const {
    for (int i = 0; i < fImageViews.size(); ++i) {
        if (fImageViews[i]->usage() == usage) {
            return fImageViews[i].get();
        }
    }

    auto sharedContext = static_cast<const VulkanSharedContext*>(this->sharedContext());
    const auto& vkTexInfo = this->vulkanTextureInfo();
    int miplevels = vkTexInfo.fMipmapped == Mipmapped::kYes
                    ? SkMipmap::ComputeLevelCount(this->dimensions()) + 1
                    : 1;
    auto imageView = VulkanImageView::Make(sharedContext,
                                           fImage,
                                           vkTexInfo.fFormat,
                                           usage,
                                           miplevels,
                                           fYcbcrConversion);
    return fImageViews.push_back(std::move(imageView)).get();
}

bool VulkanTexture::supportsInputAttachmentUsage() const {
    return (this->vulkanTextureInfo().fImageUsageFlags & VK_IMAGE_USAGE_INPUT_ATTACHMENT_BIT);
}

size_t VulkanTexture::onUpdateGpuMemorySize() {
    if (!uses_lazy_memory(fMemoryAlloc)) {
        // We don't expect non-transient textures to change their size over time.
        return this->gpuMemorySize();
    }

    auto sharedContext = static_cast<const VulkanSharedContext*>(this->sharedContext());
    VkDeviceSize committedMemory;
    VULKAN_CALL(sharedContext->interface(),
                GetDeviceMemoryCommitment(sharedContext->device(),
                                          fMemoryAlloc.fMemory,
                                          &committedMemory));
    return committedMemory;
}

sk_sp<VulkanDescriptorSet> VulkanTexture::getCachedSingleTextureDescriptorSet(
        const Sampler* sampler) const {
    SkASSERT(sampler);
    for (auto& cachedSet : fCachedSingleTextureDescSets) {
        if (cachedSet.first->uniqueID() == sampler->uniqueID()) {
            return cachedSet.second;
        }
    }
    return nullptr;
}

void VulkanTexture::addCachedSingleTextureDescriptorSet(sk_sp<VulkanDescriptorSet> set,
                                                        sk_sp<const Sampler> sampler) const {
    SkASSERT(set);
    SkASSERT(sampler);
    fCachedSingleTextureDescSets.push_back(std::make_pair(std::move(sampler), std::move(set)));
}

sk_sp<VulkanFramebuffer> VulkanTexture::getCachedFramebuffer(
        const RenderPassDesc& renderPassDesc,
        const VulkanTexture* msaaTexture,
        const VulkanTexture* depthStencilTexture) const {
    for (auto& cachedFB : fCachedFramebuffers) {
        if (cachedFB->compatible(renderPassDesc, msaaTexture, depthStencilTexture)) {
            return cachedFB;
        }
    }
    return nullptr;
}

void VulkanTexture::addCachedFramebuffer(sk_sp<VulkanFramebuffer> fb) {
    SkASSERT(fb);
    fCachedFramebuffers.push_back(std::move(fb));
}

bool VulkanTexture::canUploadOnHost() const {
    // Can't use host-image-copy if the usage flag is not set.
    if ((this->vulkanTextureInfo().fImageUsageFlags & VK_IMAGE_USAGE_HOST_TRANSFER_BIT) == 0) {
        return false;
    }

    // Can't use host-image-copy if the image is busy on the GPU. While UploadSource was responsible
    // for ensuring the TextureProxy was uniquely held (so other threads can't add work for that
    // proxy), we have to ensure the host upload won't conflict with any other use of the texture:
    //  - The texture could be used on the GPU by some now-deleted proxy, so we'll have to fall
    //    back to a regular upload via command buffer submission.
    //  - The texture needs to be non-shareable, otherwise multiple proxies could be instantiated
    //    with the same proxy.
    if (this->shareable() != Shareable::kNo || this->isBusyOnGPU()) {
        return false;
    }

    // For now, only use host-image-copy if the image has never been used. If needed in the future,
    // we could inspect the VkPhysicalDeviceHostImageCopyProperties::pCopySrcLayouts array to know
    // which layouts the image can be to be used with HIC. However, a better solution could be to
    // recreate the VkImage even if the existing one is busy on the GPU, since this function
    // entirely overwrites the texture anyway.
    if (this->currentLayout() != VK_IMAGE_LAYOUT_UNDEFINED) {
        return false;
    }

    return true;
}

bool VulkanTexture::uploadDataOnHost(const UploadSource& source) {
    // Should not have changed since canUploadOnHost() returned true given the usage possibilities
    // implied by these constraints (current thread is the only possible mutator).
    SkASSERT(source.formatXferFn().isIdentity());
    SkASSERT(source.view().proxy()->unique());
    SkASSERT(source.view().proxy()->texture() == this);
    SkASSERT(this->shareable() == Shareable::kNo);
    SkASSERT(!this->isBusyOnGPU());
    SkASSERT(this->vulkanTextureInfo().fImageUsageFlags & VK_IMAGE_USAGE_HOST_TRANSFER_BIT);
    SkASSERT(this->currentLayout() == VK_IMAGE_LAYOUT_UNDEFINED);

    auto sharedContext = static_cast<const VulkanSharedContext*>(this->sharedContext());
    SkSpan<const MipLevel> levels = source.levels();
    const unsigned int mipLevelCount = levels.size();

    const TextureInfo& textureInfo = this->textureInfo();
    const TextureFormat format = TextureInfoPriv::ViewFormat(textureInfo);
    const VkImageAspectFlags aspectFlags = GetVkImageAspectFlags(format);

    SkASSERT(this->currentLayout() == VK_IMAGE_LAYOUT_UNDEFINED);

    VkHostImageLayoutTransitionInfo transition = {};
    transition.sType = VK_STRUCTURE_TYPE_HOST_IMAGE_LAYOUT_TRANSITION_INFO;
    transition.image = fImage;
    transition.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED;
    transition.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    transition.subresourceRange.aspectMask = aspectFlags;
    transition.subresourceRange.levelCount = mipLevelCount;
    transition.subresourceRange.layerCount = 1;

    if (VULKAN_CALL(sharedContext->interface(),
                    TransitionImageLayout(sharedContext->device(), 1, &transition)) != VK_SUCCESS) {
        return false;
    }
    this->updateImageLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

    TArray<VkMemoryToImageCopy> copyRegions(mipLevelCount);

    // The assumption is either that we have no mipmaps, or that our rect is the entire texture,
    // which is guaranteed by UploadSource's creation.
    const SkIRect& dstRect = source.dstRect();
    SkASSERT(mipLevelCount == 1 || dstRect == SkIRect::MakeSize(this->dimensions()));

    // Copy data mip by mip.
    const int32_t offsetX = dstRect.x();
    const int32_t offsetY = dstRect.y();
    int32_t currentWidth = dstRect.width();
    int32_t currentHeight = dstRect.height();

    const int bpp = TextureFormatBytesPerBlock(TextureInfoPriv::ViewFormat(this->textureInfo()));
    for (unsigned int currentMipLevel = 0; currentMipLevel < mipLevelCount; currentMipLevel++) {
        VkMemoryToImageCopy copyRegion = {};
        copyRegion.sType = VK_STRUCTURE_TYPE_MEMORY_TO_IMAGE_COPY;
        copyRegion.pHostPointer = levels[currentMipLevel].fPixels;
        copyRegion.memoryRowLength = levels[currentMipLevel].fRowBytes / bpp;
        copyRegion.memoryImageHeight = 0;  // Tightly packed
        copyRegion.imageSubresource.aspectMask = aspectFlags;
        copyRegion.imageSubresource.mipLevel = currentMipLevel;
        copyRegion.imageSubresource.layerCount = 1;
        copyRegion.imageOffset.x = offsetX;
        copyRegion.imageOffset.y = offsetY;
        copyRegion.imageExtent.width = currentWidth;
        copyRegion.imageExtent.height = currentHeight;
        copyRegion.imageExtent.depth = 1;

        copyRegions.push_back(copyRegion);

        // Calculate the extent for the next mip. The offset does not need modification, since it's
        // zero if mipLevelCount > 1, asserted before the loop.
        currentWidth = std::max(1, currentWidth / 2);
        currentHeight = std::max(1, currentHeight / 2);
    }

    VkCopyMemoryToImageInfo copyInfo = {};
    copyInfo.sType = VK_STRUCTURE_TYPE_COPY_MEMORY_TO_IMAGE_INFO;
    copyInfo.dstImage = fImage;
    copyInfo.dstImageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    copyInfo.regionCount = mipLevelCount;
    copyInfo.pRegions = copyRegions.data();

    const VkResult result = VULKAN_CALL(sharedContext->interface(),
                                        CopyMemoryToImage(sharedContext->device(), &copyInfo));
    return result == VK_SUCCESS;
}

} // namespace skgpu::graphite