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

#ifndef SkRuntimeEffectPriv_DEFINED
#define SkRuntimeEffectPriv_DEFINED

#include "include/core/SkRefCnt.h"
#include "include/core/SkString.h"
#include "include/effects/SkRuntimeEffect.h"
#include "include/private/SkSLSampleUsage.h"
#include "include/private/base/SkAssert.h"
#include "include/private/base/SkSpan_impl.h"
#include "include/private/base/SkTArray.h"
#include "src/shaders/SkShaderBase.h"

#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <vector>

#ifdef SK_ENABLE_SKSL
#include "include/sksl/SkSLVersion.h"

#ifdef SK_ENABLE_SKSL_IN_RASTER_PIPELINE
#include "src/sksl/codegen/SkSLRasterPipelineBuilder.h"
#endif

#ifdef SK_ENABLE_SKVM
#include "src/sksl/codegen/SkSLVMCodeGenerator.h"
#endif

class SkArenaAlloc;
class SkColorSpace;
class SkData;
class SkMatrix;
class SkReadBuffer;
class SkShader;
class SkWriteBuffer;
struct SkStageRec;

namespace SkSL {
class Context;
class Variable;
struct Program;
}

class SkCapabilities;
struct SkColorSpaceXformSteps;

class SkRuntimeEffectPriv {
public:
    struct UniformsCallbackContext {
        const SkColorSpace* fDstColorSpace;
    };

    // Private (experimental) API for creating runtime shaders with late-bound uniforms.
    // The callback must produce a uniform data blob of the correct size for the effect.
    // It is invoked at "draw" time (essentially, when a draw call is made against the canvas
    // using the resulting shader). There are no strong guarantees about timing.
    // Serializing the resulting shader will immediately invoke the callback (and record the
    // resulting uniforms).
    using UniformsCallback = std::function<sk_sp<const SkData>(const UniformsCallbackContext&)>;
    static sk_sp<SkShader> MakeDeferredShader(const SkRuntimeEffect* effect,
                                              UniformsCallback uniformsCallback,
                                              SkSpan<SkRuntimeEffect::ChildPtr> children,
                                              const SkMatrix* localMatrix = nullptr);

    // Helper function when creating an effect for a GrSkSLFP that verifies an effect will
    // implement the GrFragmentProcessor "constant output for constant input" optimization flag.
    static bool SupportsConstantOutputForConstantInput(const SkRuntimeEffect* effect) {
        // This optimization is only implemented for color filters without any children.
        if (!effect->allowColorFilter() || !effect->children().empty()) {
            return false;
        }
#if defined(SK_ENABLE_SKVM)
        return effect->getFilterColorProgram();
#else
        return true;
#endif
    }

    static uint32_t Hash(const SkRuntimeEffect& effect) {
        return effect.hash();
    }

    static const SkSL::Program& Program(const SkRuntimeEffect& effect) {
        return *effect.fBaseProgram;
    }

    static SkRuntimeEffect::Options ES3Options() {
        SkRuntimeEffect::Options options;
        options.maxVersionAllowed = SkSL::Version::k300;
        return options;
    }

    static void AllowPrivateAccess(SkRuntimeEffect::Options* options) {
        options->allowPrivateAccess = true;
    }

    static SkRuntimeEffect::Uniform VarAsUniform(const SkSL::Variable&,
                                                 const SkSL::Context&,
                                                 size_t* offset);

    // If there are layout(color) uniforms then this performs color space transformation on the
    // color values and returns a new SkData. Otherwise, the original data is returned.
    static sk_sp<const SkData> TransformUniforms(SkSpan<const SkRuntimeEffect::Uniform> uniforms,
                                                 sk_sp<const SkData> originalData,
                                                 const SkColorSpaceXformSteps&);
    static sk_sp<const SkData> TransformUniforms(SkSpan<const SkRuntimeEffect::Uniform> uniforms,
                                                 sk_sp<const SkData> originalData,
                                                 const SkColorSpace* dstCS);
    static SkSpan<const float> UniformsAsSpan(
        SkSpan<const SkRuntimeEffect::Uniform> uniforms,
        sk_sp<const SkData> originalData,
        bool alwaysCopyIntoAlloc,
        const SkColorSpace* destColorSpace,
        SkArenaAlloc* alloc);

    static bool CanDraw(const SkCapabilities*, const SkSL::Program*);
    static bool CanDraw(const SkCapabilities*, const SkRuntimeEffect*);

    static bool ReadChildEffects(SkReadBuffer& buffer,
                                 const SkRuntimeEffect* effect,
                                 skia_private::TArray<SkRuntimeEffect::ChildPtr>* children);
    static void WriteChildEffects(SkWriteBuffer &buffer,
                                  const std::vector<SkRuntimeEffect::ChildPtr> &children);

#ifdef SK_ENABLE_SKVM
    static std::vector<skvm::Val> MakeSkVMUniforms(skvm::Builder*,
                                                   skvm::Uniforms*,
                                                   size_t inputSize,
                                                   const SkData& inputs);
#endif

#if defined(SK_GRAPHITE)
static void AddChildrenToKey(SkSpan<const SkRuntimeEffect::ChildPtr> children,
                             SkSpan<const SkRuntimeEffect::Child> childInfo,
                             const skgpu::graphite::KeyContext& keyContext,
                             skgpu::graphite::PaintParamsKeyBuilder* builder,
                             skgpu::graphite::PipelineDataGatherer* gatherer);
#endif
};

// These internal APIs for creating runtime effects vary from the public API in two ways:
//
//     1) they're used in contexts where it's not useful to receive an error message;
//     2) they're cached.
//
// Users of the public SkRuntimeEffect::Make*() can of course cache however they like themselves;
// keeping these APIs private means users will not be forced into our cache or cache policy.

sk_sp<SkRuntimeEffect> SkMakeCachedRuntimeEffect(
        SkRuntimeEffect::Result (*make)(SkString sksl, const SkRuntimeEffect::Options&),
        SkString sksl);

inline sk_sp<SkRuntimeEffect> SkMakeCachedRuntimeEffect(
        SkRuntimeEffect::Result (*make)(SkString, const SkRuntimeEffect::Options&),
        const char* sksl) {
    return SkMakeCachedRuntimeEffect(make, SkString{sksl});
}

// Internal API that assumes (and asserts) that the shader code is valid, but does no internal
// caching. Used when the caller will cache the result in a static variable. Ownership is passed to
// the caller; the effect will be leaked if it the pointer is not stored or explicitly deleted.
inline SkRuntimeEffect* SkMakeRuntimeEffect(
        SkRuntimeEffect::Result (*make)(SkString, const SkRuntimeEffect::Options&),
        const char* sksl,
        SkRuntimeEffect::Options options = SkRuntimeEffect::Options{}) {
#if defined(SK_DEBUG)
    // Our SKSL snippets we embed in Skia should not have comments or excess indentation.
    // Removing them helps trim down code size and speeds up parsing
    if (SkStrContains(sksl, "//") || SkStrContains(sksl, "    ")) {
        SkDEBUGFAILF("Found SkSL snippet that can be minified: \n %s\n", sksl);
    }
#endif
    SkRuntimeEffectPriv::AllowPrivateAccess(&options);
    auto result = make(SkString{sksl}, options);
    if (!result.effect) {
        SK_ABORT("%s", result.errorText.c_str());
    }
    return result.effect.release();
}

#ifdef SK_ENABLE_SKSL_IN_RASTER_PIPELINE
class RuntimeEffectRPCallbacks : public SkSL::RP::Callbacks {
public:
    RuntimeEffectRPCallbacks(const SkStageRec& s,
                             const SkShaderBase::MatrixRec& m,
                             SkSpan<const SkRuntimeEffect::ChildPtr> c,
                             SkSpan<const SkSL::SampleUsage> u)
            : fStage(s), fMatrix(m), fChildren(c), fSampleUsages(u) {}

    bool appendShader(int index) override;
    bool appendColorFilter(int index) override;
    bool appendBlender(int index) override;

    // TODO: If an effect calls these intrinsics more than once, we could cache and re-use the steps
    // object(s), rather than re-creating them in the arena repeatedly.
    void toLinearSrgb(const void* color) override;

    void fromLinearSrgb(const void* color) override;

private:
    void applyColorSpaceXform(const SkColorSpaceXformSteps& tempXform, const void* color);

    const SkStageRec& fStage;
    const SkShaderBase::MatrixRec& fMatrix;
    SkSpan<const SkRuntimeEffect::ChildPtr> fChildren;
    SkSpan<const SkSL::SampleUsage> fSampleUsages;
};
#endif  // SK_ENABLE_SKSL_IN_RASTER_PIPELINE

#if defined(SK_ENABLE_SKVM)
class RuntimeEffectVMCallbacks : public SkSL::SkVMCallbacks {
public:
    RuntimeEffectVMCallbacks(skvm::Builder* builder,
                             skvm::Uniforms* uniforms,
                             SkArenaAlloc* alloc,
                             const std::vector<SkRuntimeEffect::ChildPtr>& children,
                             const SkShaderBase::MatrixRec& mRec,
                             skvm::Color inColor,
                             const SkColorInfo& colorInfo)
            : fBuilder(builder)
            , fUniforms(uniforms)
            , fAlloc(alloc)
            , fChildren(children)
            , fMRec(mRec)
            , fInColor(inColor)
            , fColorInfo(colorInfo) {}

    skvm::Color sampleShader(int ix, skvm::Coord coord) override;

    skvm::Color sampleColorFilter(int ix, skvm::Color color) override;

    skvm::Color sampleBlender(int ix, skvm::Color src, skvm::Color dst) override;

    skvm::Color toLinearSrgb(skvm::Color color) override;

    skvm::Color fromLinearSrgb(skvm::Color color) override;

    skvm::Builder* fBuilder;
    skvm::Uniforms* fUniforms;
    SkArenaAlloc* fAlloc;
    const std::vector<SkRuntimeEffect::ChildPtr>& fChildren;
    const SkShaderBase::MatrixRec& fMRec;
    const skvm::Color fInColor;
    const SkColorInfo& fColorInfo;
};
#endif  // defined(SK_ENABLE_SKVM)

#endif  // SK_ENABLE_SKSL

#endif  // SkRuntimeEffectPriv_DEFINED