/*
 * Copyright 2023 Google LLC
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */
#include "include/gpu/ganesh/vk/GrVkBackendSurface.h"

#include "include/core/SkRefCnt.h"
#include "include/core/SkTextureCompressionType.h"
#include "include/gpu/GpuTypes.h"
#include "include/gpu/GrBackendSurface.h"
#include "include/gpu/GrTypes.h"
#include "include/gpu/MutableTextureState.h"
#include "include/gpu/vk/GrVkTypes.h"
#include "include/gpu/vk/VulkanMutableTextureState.h"
#include "include/private/base/SkAssert.h"
#include "include/private/base/SkTo.h"
#include "include/private/gpu/ganesh/GrTypesPriv.h"
#include "src/gpu/ganesh/GrBackendSurfacePriv.h"
#include "src/gpu/ganesh/vk/GrVkTypesPriv.h"
#include "src/gpu/ganesh/vk/GrVkUtil.h"
#include "src/gpu/vk/VulkanMutableTextureStatePriv.h"
#include "src/gpu/vk/VulkanUtilsPriv.h"

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <string>

class GrVkBackendFormatData final : public GrBackendFormatData {
public:
    GrVkBackendFormatData(VkFormat format, const GrVkYcbcrConversionInfo& ycbcrInfo)
            : fFormat(format), fYcbcrConversionInfo(ycbcrInfo) {}

    VkFormat asVkFormat() const { return fFormat; }
    const GrVkYcbcrConversionInfo* getYcbcrConversionInfo() const { return &fYcbcrConversionInfo; }

private:
    SkTextureCompressionType compressionType() const override {
        switch (fFormat) {
            case VK_FORMAT_ETC2_R8G8B8_UNORM_BLOCK:
                return SkTextureCompressionType::kETC2_RGB8_UNORM;
            case VK_FORMAT_BC1_RGB_UNORM_BLOCK:
                return SkTextureCompressionType::kBC1_RGB8_UNORM;
            case VK_FORMAT_BC1_RGBA_UNORM_BLOCK:
                return SkTextureCompressionType::kBC1_RGBA8_UNORM;
            default:
                return SkTextureCompressionType::kNone;
        }
    }

    size_t bytesPerBlock() const override {
        return skgpu::VkFormatBytesPerBlock(fFormat);
    }

    int stencilBits() const override {
        return skgpu::VkFormatStencilBits(fFormat);
    }

    uint32_t channelMask() const override {
        return skgpu::VkFormatChannels(fFormat);
    }

    GrColorFormatDesc desc() const override {
        return GrVkFormatDesc(fFormat);
    }

    bool equal(const GrBackendFormatData* that) const override {
        SkASSERT(!that || that->type() == GrBackendApi::kVulkan);
        if (auto otherVk = static_cast<const GrVkBackendFormatData*>(that)) {
            return fFormat == otherVk->fFormat &&
                   fYcbcrConversionInfo == otherVk->fYcbcrConversionInfo;
        }
        return false;
    }

    std::string toString() const override {
#if defined(SK_DEBUG) || GR_TEST_UTILS
        return skgpu::VkFormatToStr(fFormat);
#else
        return "";
#endif
    }

    void copyTo(AnyFormatData& formatData) const override {
        formatData.emplace<GrVkBackendFormatData>(fFormat, fYcbcrConversionInfo);
    }

    void makeTexture2D() override {
        // If we have a ycbcr we remove it from the backend format and set the VkFormat to
        // R8G8B8A8_UNORM
        if (fYcbcrConversionInfo.isValid()) {
            fYcbcrConversionInfo = GrVkYcbcrConversionInfo();
            fFormat = VK_FORMAT_R8G8B8A8_UNORM;
        }
    }

#if defined(SK_DEBUG)
    GrBackendApi type() const override { return GrBackendApi::kVulkan; }
#endif

    VkFormat fFormat;
    GrVkYcbcrConversionInfo fYcbcrConversionInfo;
};

static const GrVkBackendFormatData* get_and_cast_data(const GrBackendFormat& format) {
    auto data = GrBackendSurfacePriv::GetBackendData(format);
    SkASSERT(!data || data->type() == GrBackendApi::kVulkan);
    return static_cast<const GrVkBackendFormatData*>(data);
}

namespace GrBackendFormats {

GrBackendFormat MakeVk(VkFormat format, bool willUseDRMFormatModifiers) {
    return GrBackendSurfacePriv::MakeGrBackendFormat(
            GrTextureType::k2D,
            GrBackendApi::kVulkan,
            GrVkBackendFormatData(format, GrVkYcbcrConversionInfo{}));
}

GrBackendFormat MakeVk(const GrVkYcbcrConversionInfo& ycbcrInfo, bool willUseDRMFormatModifiers) {
    SkASSERT(ycbcrInfo.isValid());
    GrTextureType textureType =
            ((ycbcrInfo.isValid() && ycbcrInfo.fExternalFormat) || willUseDRMFormatModifiers)
                    ? GrTextureType::kExternal
                    : GrTextureType::k2D;
    return GrBackendSurfacePriv::MakeGrBackendFormat(
            textureType,
            GrBackendApi::kVulkan,
            GrVkBackendFormatData(ycbcrInfo.fFormat, ycbcrInfo));
}

bool AsVkFormat(const GrBackendFormat& format, VkFormat* vkFormat) {
    SkASSERT(vkFormat);
    if (format.isValid() && format.backend() == GrBackendApi::kVulkan) {
        const GrVkBackendFormatData* data = get_and_cast_data(format);
        SkASSERT(data);
        *vkFormat = data->asVkFormat();
        return true;
    }
    return false;
}

const GrVkYcbcrConversionInfo* GetVkYcbcrConversionInfo(const GrBackendFormat& format) {
    if (format.isValid() && format.backend() == GrBackendApi::kVulkan) {
        const GrVkBackendFormatData* data = get_and_cast_data(format);
        SkASSERT(data);
        return data->getYcbcrConversionInfo();
    }
    return nullptr;
}

}  // namespace GrBackendFormats


class GrVkBackendTextureData final : public GrBackendTextureData {
public:
    GrVkBackendTextureData(const GrVkImageInfo& info,
                           sk_sp<skgpu::MutableTextureState> mutableState = nullptr)
            : fVkInfo(info) {
        if (mutableState) {
            fMutableState = std::move(mutableState);
        } else {
            fMutableState =
                    sk_make_sp<skgpu::MutableTextureState>(skgpu::MutableTextureStates::MakeVulkan(
                            info.fImageLayout, info.fCurrentQueueFamily));
        }
    }

    const GrVkImageInfo& info() const { return fVkInfo; }

    sk_sp<skgpu::MutableTextureState> getMutableState() const override {
        return fMutableState;
    }
    void setMutableState(const skgpu::MutableTextureState& state) override {
        fMutableState->set(state);
    }

    skgpu::MutableTextureState* mutableState() { return fMutableState.get(); }
    const skgpu::MutableTextureState* mutableState() const { return fMutableState.get(); }

private:
    void copyTo(AnyTextureData& textureData) const override {
        textureData.emplace<GrVkBackendTextureData>(fVkInfo, fMutableState);
    }

    bool isProtected() const override { return fVkInfo.fProtected == skgpu::Protected::kYes; }

    bool equal(const GrBackendTextureData* that) const override {
        SkASSERT(!that || that->type() == GrBackendApi::kVulkan);
        if (auto otherVk = static_cast<const GrVkBackendTextureData*>(that)) {
            // For our tests when checking equality we are assuming both backendTexture objects will
            // be using the same mutable state object.
            if (fMutableState != otherVk->fMutableState) {
                return false;
            }
            return GrVkImageInfoWithMutableState(fVkInfo, fMutableState.get()) ==
                   GrVkImageInfoWithMutableState(otherVk->fVkInfo, fMutableState.get());
        }
        return false;
    }

    bool isSameTexture(const GrBackendTextureData* that) const override {
        SkASSERT(!that || that->type() == GrBackendApi::kVulkan);
        if (auto otherVk = static_cast<const GrVkBackendTextureData*>(that)) {
            return fVkInfo.fImage == otherVk->fVkInfo.fImage;
        }
        return false;

    }

    GrBackendFormat getBackendFormat() const override {
        auto info = GrVkImageInfoWithMutableState(fVkInfo, fMutableState.get());
        bool usesDRMModifier = info.fImageTiling == VK_IMAGE_TILING_DRM_FORMAT_MODIFIER_EXT;
        if (info.fYcbcrConversionInfo.isValid()) {
            SkASSERT(info.fFormat == info.fYcbcrConversionInfo.fFormat);
            return GrBackendFormats::MakeVk(info.fYcbcrConversionInfo, usesDRMModifier);
        }
        return GrBackendFormats::MakeVk(info.fFormat, usesDRMModifier);
    }

#if defined(SK_DEBUG)
    GrBackendApi type() const override { return GrBackendApi::kVulkan; }
#endif

    GrVkImageInfo fVkInfo;
    sk_sp<skgpu::MutableTextureState> fMutableState;
};

static const GrVkBackendTextureData* get_and_cast_data(const GrBackendTexture& texture) {
    auto data = GrBackendSurfacePriv::GetBackendData(texture);
    SkASSERT(!data || data->type() == GrBackendApi::kVulkan);
    return static_cast<const GrVkBackendTextureData*>(data);
}

static GrVkBackendTextureData* get_and_cast_data(GrBackendTexture* texture) {
    auto data = GrBackendSurfacePriv::GetBackendData(texture);
    SkASSERT(!data || data->type() == GrBackendApi::kVulkan);
    return static_cast<GrVkBackendTextureData*>(data);
}

static GrTextureType vk_image_info_to_texture_type(const GrVkImageInfo& info) {
    if ((info.fYcbcrConversionInfo.isValid() && info.fYcbcrConversionInfo.fExternalFormat != 0) ||
        info.fImageTiling == VK_IMAGE_TILING_DRM_FORMAT_MODIFIER_EXT) {
        return GrTextureType::kExternal;
    }
    return GrTextureType::k2D;
}

static const VkImageUsageFlags kDefaultUsageFlags =
        VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT;

static const VkImageUsageFlags kDefaultRTUsageFlags =
        kDefaultUsageFlags | VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;

// We don't know if the backend texture is made renderable or not, so we default the usage flags
// to include color attachment as well.
static const VkImageUsageFlags kDefaultTexRTUsageFlags =
        kDefaultUsageFlags | VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;

static GrVkImageInfo apply_default_usage_flags(const GrVkImageInfo& info,
                                               VkImageUsageFlags defaultFlags) {
    if (info.fImageUsageFlags == 0) {
        GrVkImageInfo newInfo = info;
        newInfo.fImageUsageFlags = defaultFlags;
        return newInfo;
    }
    return info;
}

namespace GrBackendTextures {

GrBackendTexture MakeVk(int width,
                        int height,
                        const GrVkImageInfo& vkInfo,
                        std::string_view label) {
    return GrBackendSurfacePriv::MakeGrBackendTexture(
            width,
            height,
            label,
            skgpu::Mipmapped(vkInfo.fLevelCount > 1),
            GrBackendApi::kVulkan,
            vk_image_info_to_texture_type(vkInfo),
            GrVkBackendTextureData(apply_default_usage_flags(vkInfo, kDefaultTexRTUsageFlags)));
}

GrBackendTexture MakeVk(int width,
                        int height,
                        const GrVkImageInfo& vkInfo,
                        sk_sp<skgpu::MutableTextureState> mutableState) {
    return GrBackendSurfacePriv::MakeGrBackendTexture(
            width,
            height,
            /*label=*/{},
            skgpu::Mipmapped(vkInfo.fLevelCount > 1),
            GrBackendApi::kVulkan,
            vk_image_info_to_texture_type(vkInfo),
            GrVkBackendTextureData(apply_default_usage_flags(vkInfo, kDefaultTexRTUsageFlags),
                                   std::move(mutableState)));
}

bool GetVkImageInfo(const GrBackendTexture& tex, GrVkImageInfo* outInfo) {
    if (!tex.isValid() || tex.backend() != GrBackendApi::kVulkan) {
        return false;
    }
    const GrVkBackendTextureData* data = get_and_cast_data(tex);
    SkASSERT(data);
    *outInfo = GrVkImageInfoWithMutableState(data->info(), data->mutableState());
    return true;
}

void SetVkImageLayout(GrBackendTexture* tex, VkImageLayout layout) {
    if (tex && tex->isValid() && tex->backend() == GrBackendApi::kVulkan) {
        GrVkBackendTextureData* data = get_and_cast_data(tex);
        SkASSERT(data);
        skgpu::MutableTextureStates::SetVkImageLayout(data->mutableState(), layout);
    }
}

}  // namespace GrBackendTextures


class GrVkBackendRenderTargetData final : public GrBackendRenderTargetData {
public:
    GrVkBackendRenderTargetData(const GrVkImageInfo& info,
                                sk_sp<skgpu::MutableTextureState> mutableState = nullptr)
            : fVkInfo(info) {
        if (mutableState) {
            fMutableState = std::move(mutableState);
        } else {
            fMutableState =
                    sk_make_sp<skgpu::MutableTextureState>(skgpu::MutableTextureStates::MakeVulkan(
                            info.fImageLayout, info.fCurrentQueueFamily));
        }
    }

    const GrVkImageInfo& info() const { return fVkInfo; }

    sk_sp<skgpu::MutableTextureState> getMutableState() const override {
        return fMutableState;
    }
    void setMutableState(const skgpu::MutableTextureState& state) override {
        fMutableState->set(state);
    }

    skgpu::MutableTextureState* mutableState() { return fMutableState.get(); }
    const skgpu::MutableTextureState* mutableState() const { return fMutableState.get(); }

private:
    GrBackendFormat getBackendFormat() const override {
        auto info = GrVkImageInfoWithMutableState(fVkInfo, fMutableState.get());
        if (info.fYcbcrConversionInfo.isValid()) {
            SkASSERT(info.fFormat == info.fYcbcrConversionInfo.fFormat);
            return GrBackendFormats::MakeVk(info.fYcbcrConversionInfo);
        }
        return GrBackendFormats::MakeVk(info.fFormat);
    }

    bool isProtected() const override { return fVkInfo.fProtected == skgpu::Protected::kYes; }

    bool equal(const GrBackendRenderTargetData* that) const override {
        SkASSERT(!that || that->type() == GrBackendApi::kVulkan);
        if (auto otherVk = static_cast<const GrVkBackendRenderTargetData*>(that)) {
            // For our tests when checking equality we are assuming both objects will be using the
            // same mutable state object.
            if (fMutableState != otherVk->fMutableState) {
                return false;
            }
            return GrVkImageInfoWithMutableState(fVkInfo, fMutableState.get()) ==
                   GrVkImageInfoWithMutableState(otherVk->fVkInfo, fMutableState.get());
        }
        return false;
    }

    void copyTo(AnyRenderTargetData& rtData) const override {
        rtData.emplace<GrVkBackendRenderTargetData>(fVkInfo, fMutableState);
    }

#if defined(SK_DEBUG)
    GrBackendApi type() const override { return GrBackendApi::kVulkan; }
#endif

    GrVkImageInfo fVkInfo;
    sk_sp<skgpu::MutableTextureState> fMutableState;
};

static const GrVkBackendRenderTargetData* get_and_cast_data(const GrBackendRenderTarget& rt) {
    auto data = GrBackendSurfacePriv::GetBackendData(rt);
    SkASSERT(!data || data->type() == GrBackendApi::kVulkan);
    return static_cast<const GrVkBackendRenderTargetData*>(data);
}

static GrVkBackendRenderTargetData* get_and_cast_data(GrBackendRenderTarget* rt) {
    auto data = GrBackendSurfacePriv::GetBackendData(rt);
    SkASSERT(!data || data->type() == GrBackendApi::kVulkan);
    return static_cast<GrVkBackendRenderTargetData*>(data);
}

namespace GrBackendRenderTargets {

GrBackendRenderTarget MakeVk(int width, int height, const GrVkImageInfo& vkInfo) {
    return GrBackendSurfacePriv::MakeGrBackendRenderTarget(
            width,
            height,
            std::max(1U, vkInfo.fSampleCount),
            /*stencilBits=*/0,
            GrBackendApi::kVulkan,
            /*framebufferOnly=*/false,
            GrVkBackendRenderTargetData(apply_default_usage_flags(vkInfo, kDefaultRTUsageFlags)));
}

GrBackendRenderTarget MakeVk(int width,
                             int height,
                             const GrVkImageInfo& vkInfo,
                             sk_sp<skgpu::MutableTextureState> mutableState) {
    return GrBackendSurfacePriv::MakeGrBackendRenderTarget(
            width,
            height,
            std::max(1U, vkInfo.fSampleCount),
            /*stencilBits=*/0,
            GrBackendApi::kVulkan,
            /*framebufferOnly=*/false,
            GrVkBackendRenderTargetData(apply_default_usage_flags(vkInfo, kDefaultRTUsageFlags),
                                        std::move(mutableState)));
}

bool GetVkImageInfo(const GrBackendRenderTarget& rt, GrVkImageInfo* outInfo) {
    if (!rt.isValid() || rt.backend() != GrBackendApi::kVulkan) {
        return false;
    }
    const GrVkBackendRenderTargetData* data = get_and_cast_data(rt);
    SkASSERT(data);
    *outInfo = GrVkImageInfoWithMutableState(data->info(), data->mutableState());
    return true;
}

void SetVkImageLayout(GrBackendRenderTarget* rt, VkImageLayout layout) {
    if (rt && rt->isValid() && rt->backend() == GrBackendApi::kVulkan) {
        GrVkBackendRenderTargetData* data = get_and_cast_data(rt);
        skgpu::MutableTextureStates::SetVkImageLayout(data->mutableState(), layout);
    }
}

}  // namespace GrBackendRenderTargets