/*
 * 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/render/VerticesRenderStep.h"

#include "src/core/SkSLTypeShared.h"
#include "src/core/SkVertState.h"
#include "src/core/SkVerticesPriv.h"
#include "src/gpu/graphite/DrawParams.h"
#include "src/gpu/graphite/DrawTypes.h"
#include "src/gpu/graphite/DrawWriter.h"
#include "src/gpu/graphite/PipelineData.h"
#include "src/gpu/graphite/render/CommonDepthStencilSettings.h"

namespace skgpu::graphite {

namespace {

static constexpr Attribute kPositionAttr =
        {"position", VertexAttribType::kFloat2, SkSLType::kFloat2};
static constexpr Attribute kTexCoordAttr =
        {"texCoords", VertexAttribType::kFloat2, SkSLType::kFloat2};
static constexpr Attribute kColorAttr =
        {"vertColor", VertexAttribType::kUByte4_norm, SkSLType::kHalf4};
static constexpr Attribute kSsboIndexAttr =
        {"ssboIndices", VertexAttribType::kUShort2, SkSLType::kUShort2};

static constexpr Attribute kAttributePositionOnly[] =
        {kPositionAttr, kSsboIndexAttr};
static constexpr Attribute kAttributeColor[] =
        {kPositionAttr, kColorAttr, kSsboIndexAttr};
static constexpr Attribute kAttributeTexCoords[] =
        {kPositionAttr, kTexCoordAttr, kSsboIndexAttr};
static constexpr Attribute kAttributeColorAndTexCoords[] =
        {kPositionAttr, kColorAttr, kTexCoordAttr, kSsboIndexAttr};

static constexpr SkSpan<const Attribute> kAttributes[4] = {
        kAttributePositionOnly,
        kAttributeColor,
        kAttributeTexCoords,
        kAttributeColorAndTexCoords,
    };

static constexpr Varying kVaryingColor[] =
        {{"color", SkSLType::kHalf4}};

static constexpr SkSpan<const Varying> kVaryings[2] = {
        /*none*/  {},
        /*color*/ kVaryingColor
    };

std::string variant_name(PrimitiveType type, bool hasColor, bool hasTexCoords) {
    SkASSERT(type == PrimitiveType::kTriangles || type == PrimitiveType::kTriangleStrip);
    std::string name = (type == PrimitiveType::kTriangles ? "tris" : "tristrips");
    if (hasColor) {
        name += "-color";
    }
    if (hasTexCoords) {
        name += "-texCoords";
    }
    return name;
}

}  // namespace

VerticesRenderStep::VerticesRenderStep(PrimitiveType type, bool hasColor, bool hasTexCoords)
        : RenderStep("VerticesRenderStep",
                     variant_name(type, hasColor, hasTexCoords),
                     hasColor ? Flags::kEmitsPrimitiveColor | Flags::kPerformsShading
                              : Flags::kPerformsShading,
                     /*uniforms=*/{{"localToDevice", SkSLType::kFloat4x4},
                                   {"depth", SkSLType::kFloat}},
                     type,
                     kDirectDepthGEqualPass,
                     /*vertexAttrs=*/  kAttributes[2*hasTexCoords + hasColor],
                     /*instanceAttrs=*/{},
                     /*varyings=*/     kVaryings[hasColor])
        , fHasColor(hasColor)
        , fHasTexCoords(hasTexCoords) {}

VerticesRenderStep::~VerticesRenderStep() {}

std::string VerticesRenderStep::vertexSkSL() const {
    if (fHasColor && fHasTexCoords) {
        return R"(
            color = half4(vertColor.bgr * vertColor.a, vertColor.a);
            float4 devPosition = localToDevice * float4(position, 0.0, 1.0);
            devPosition.z = depth;
            stepLocalCoords = texCoords;
        )";
    } else if (fHasTexCoords) {
        return R"(
            float4 devPosition = localToDevice * float4(position, 0.0, 1.0);
            devPosition.z = depth;
            stepLocalCoords = texCoords;
        )";
    } else if (fHasColor) {
        return R"(
            color = half4(vertColor.bgr * vertColor.a, vertColor.a);
            float4 devPosition = localToDevice * float4(position, 0.0, 1.0);
            devPosition.z = depth;
            stepLocalCoords = position;
        )";
    } else {
        return R"(
            float4 devPosition = localToDevice * float4(position, 0.0, 1.0);
            devPosition.z = depth;
            stepLocalCoords = position;
        )";
    }
}

const char* VerticesRenderStep::fragmentColorSkSL() const {
    if (fHasColor) {
        return "primitiveColor = color;\n";
    } else {
        return "";
    }
}

void VerticesRenderStep::writeVertices(DrawWriter* writer,
                                       const DrawParams& params,
                                       skvx::ushort2 ssboIndices) const {
    SkVerticesPriv info(params.geometry().vertices()->priv());
    const int vertexCount = info.vertexCount();
    const int indexCount = info.indexCount();
    const SkPoint* positions = info.positions();
    const uint16_t* indices = info.indices();
    const SkColor* colors = info.colors();
    const SkPoint* texCoords = info.texCoords();

    // This should always be the case if the Renderer was chosen appropriately, but the vertex
    // writing loop is set up in such a way that if the shader expects color or tex coords and they
    // are missing, it will just read 0s, so release builds are safe.
    SkASSERT(fHasColor == SkToBool(colors));
    SkASSERT(fHasTexCoords == SkToBool(texCoords));

    // TODO: We could access the writer's DrawBufferManager and upload the SkVertices index buffer
    // but that would require we manually manage the VertexWriter for interleaving the position,
    // color, and tex coord arrays together. This wouldn't be so bad if we let ::Vertices() take
    // a CPU index buffer that indexes into the accumulated vertex data (and handles offsetting for
    // merged drawIndexed calls), or if we could bind multiple attribute sources and copy the
    // position/color/texCoord data separately in bulk w/o using an Appender.
    DrawWriter::Vertices verts{*writer};
    verts.reserve(indices ? indexCount : vertexCount);

    VertState state(vertexCount, indices, indexCount);
    VertState::Proc vertProc = state.chooseProc(info.mode());
    while (vertProc(&state)) {
        verts.append(3) << positions[state.f0]
                        << VertexWriter::If(fHasColor, colors ? colors[state.f0]
                                                              : SK_ColorTRANSPARENT)
                        << VertexWriter::If(fHasTexCoords, texCoords ? texCoords[state.f0]
                                                                     : SkPoint{0.f, 0.f})
                        << ssboIndices
                        << positions[state.f1]
                        << VertexWriter::If(fHasColor, colors ? colors[state.f1]
                                                              : SK_ColorTRANSPARENT)
                        << VertexWriter::If(fHasTexCoords, texCoords ? texCoords[state.f1]
                                                                     : SkPoint{0.f, 0.f})
                        << ssboIndices
                        << positions[state.f2]
                        << VertexWriter::If(fHasColor, colors ? colors[state.f2]
                                                              : SK_ColorTRANSPARENT)
                        << VertexWriter::If(fHasTexCoords, texCoords ? texCoords[state.f2]
                                                                     : SkPoint{0.f, 0.f})
                        << ssboIndices;
    }
}

void VerticesRenderStep::writeUniformsAndTextures(const DrawParams& params,
                                                  PipelineDataGatherer* gatherer) const {
    // Vertices are transformed on the GPU. Store PaintDepth as a uniform to avoid copying the
    // same depth for each vertex.
    SkDEBUGCODE(UniformExpectationsValidator uev(gatherer, this->uniforms());)
    gatherer->write(params.transform().matrix());
    gatherer->write(params.order().depthAsFloat());
}

}  // namespace skgpu::graphite