/*
 * Copyright 2022 Rive
 */

// undef GENERATE_UNMULTIPLIED_PAINT_COLORS first because this file gets
// included multiple times with different defines in the Metal library.
#undef GENERATE_UNMULTIPLIED_PAINT_COLORS

#ifdef @NEVER_GENERATE_PREMULTIPLIED_PAINT_COLORS
// The specific fragment shader we're being compiled for expects un-multiplied
// paint colors all the time.
#define GENERATE_UNMULTIPLIED_PAINT_COLORS true
#elif defined(@ENABLE_ADVANCED_BLEND)
// If advanced blend is enabled, we generate unmultiplied paint colors in the
// shader. Otherwise we would have to just turn around and unmultiply them in
// order to run the blend equation.
#define GENERATE_UNMULTIPLIED_PAINT_COLORS @ENABLE_ADVANCED_BLEND
#else
// As long as advanced blend is not enabled, it's more efficient for the shader
// to generate premultiplied paint colors from the start.
#define GENERATE_UNMULTIPLIED_PAINT_COLORS false
#endif

// undef COVERAGE_TYPE first because this file gets included multiple times with
// different defines in the Metal library.
#undef COVERAGE_TYPE
#ifdef @ENABLE_FEATHER
#define COVERAGE_TYPE float4
#else
#define COVERAGE_TYPE half2
#endif

#ifdef @VERTEX
ATTR_BLOCK_BEGIN(Attrs)
#if defined(@DRAW_INTERIOR_TRIANGLES) || defined(@FEATHER_ATLAS_BLIT)
ATTR(0, packed_float3, @a_triangleVertex);
#else
ATTR(0,
     float4,
     @a_patchVertexData); // [localVertexID, outset, fillCoverage, vertexType]
ATTR(1, float4, @a_mirroredVertexData);
#endif
ATTR_BLOCK_END
#endif

VARYING_BLOCK_BEGIN
NO_PERSPECTIVE VARYING(0, float4, v_paint);

#ifdef @FEATHER_ATLAS_BLIT
NO_PERSPECTIVE VARYING(1, float2, v_atlasCoord);
#elif !defined(@RENDER_MODE_DEPTH_STENCIL)
#ifdef @DRAW_INTERIOR_TRIANGLES
@OPTIONALLY_FLAT VARYING(1, half, v_windingWeight);
#else
NO_PERSPECTIVE VARYING(2, COVERAGE_TYPE, v_coverages);
#endif //@DRAW_INTERIOR_TRIANGLES
@OPTIONALLY_FLAT VARYING(3, half, v_pathID);
#endif // !@RENDER_MODE_DEPTH_STENCIL

#ifdef @ENABLE_CLIPPING
#ifdef @FEATHER_ATLAS_BLIT
@OPTIONALLY_FLAT VARYING(4, half, v_clipID); // [clipID, outerClipID]
#else
@OPTIONALLY_FLAT VARYING(4, half2, v_clipIDs); // [clipID, outerClipID]
#endif
#endif // @ENABLE_CLIPPING
#if defined(@ENABLE_CLIP_RECT) && !defined(@RENDER_MODE_DEPTH_STENCIL)
NO_PERSPECTIVE VARYING(5, float4, v_clipRect);
#endif
#ifdef @ENABLE_ADVANCED_BLEND
@OPTIONALLY_FLAT VARYING(6, half, v_blendMode);
#endif
#ifdef @RENDER_MODE_CLOCKWISE_ATOMIC
FLAT VARYING(7, uint2, v_coveragePlacement);
VARYING(8, float2, v_coverageCoord);
#endif
#ifdef @ENABLE_MODULATED_IMAGE
NO_PERSPECTIVE VARYING(9, float3, v_image);
#endif

VARYING_BLOCK_END

#ifdef @VERTEX

#ifdef @EMULATE_DYNAMIC_COLOR_WRITE_DISABLE
// Emulation for VK_EXT_color_write_enable.
// 1 writes color normally; 0 suppresses it by outputting v_paint == 0 (which
// then gets discarded at the blend step).
// NOTE: This is intentionally declared inside "#ifdef @VERTEX" so it doesn't
// get needlessly added to fragment shaders.
PUSH_CONSTANT_BLOCK_BEGIN(PushConstants)
PUSH_CONSTANT(float, colorWriteEnable)
PUSH_CONSTANT_BLOCK_END(pushConstants)
#endif

VERTEX_MAIN(@drawVertexMain, Attrs, attrs, _vertexID, _instanceID)
{
#if defined(@DRAW_INTERIOR_TRIANGLES) || defined(@FEATHER_ATLAS_BLIT)
    ATTR_UNPACK(_vertexID, attrs, @a_triangleVertex, float3);
#else
    ATTR_UNPACK(_vertexID, attrs, @a_patchVertexData, float4);
    ATTR_UNPACK(_vertexID, attrs, @a_mirroredVertexData, float4);
#endif

    VARYING_INIT(v_paint, float4);
#if defined(@ENABLE_MODULATED_IMAGE)
    VARYING_INIT(v_image, float3);
#endif

#ifdef @FEATHER_ATLAS_BLIT
    VARYING_INIT(v_atlasCoord, float2);
#elif !defined(@RENDER_MODE_DEPTH_STENCIL)
#ifdef @DRAW_INTERIOR_TRIANGLES
    VARYING_INIT(v_windingWeight, half);
#else
    VARYING_INIT(v_coverages, COVERAGE_TYPE);
#endif //@DRAW_INTERIOR_TRIANGLES
    VARYING_INIT(v_pathID, half);
#endif // !@RENDER_MODE_DEPTH_STENCIL

#ifdef @ENABLE_CLIPPING
#ifdef @FEATHER_ATLAS_BLIT
    VARYING_INIT(v_clipID, half);
#else
    VARYING_INIT(v_clipIDs, half2);
#endif
#endif // @ENABLE_CLIPPING
#if defined(@ENABLE_CLIP_RECT) && !defined(@RENDER_MODE_DEPTH_STENCIL)
    VARYING_INIT(v_clipRect, float4);
#endif
#ifdef @ENABLE_ADVANCED_BLEND
    VARYING_INIT(v_blendMode, half);
#endif
#ifdef @RENDER_MODE_CLOCKWISE_ATOMIC
    VARYING_INIT(v_coveragePlacement, uint2);
    VARYING_INIT(v_coverageCoord, float2);
#endif

    bool shouldDiscardVertex = false;
    uint pathID;
    float2 vertexPosition;
#ifdef @RENDER_MODE_DEPTH_STENCIL
    ushort pathZIndex;
#endif

#ifdef @FEATHER_ATLAS_BLIT
    vertexPosition =
        unpack_atlas_coverage_vertex(@a_triangleVertex,
                                     pathID,
#ifdef @RENDER_MODE_DEPTH_STENCIL
                                     pathZIndex,
#endif
                                     v_atlasCoord VERTEX_CONTEXT_UNPACK);
#elif defined(@DRAW_INTERIOR_TRIANGLES)
    vertexPosition = unpack_interior_triangle_vertex(@a_triangleVertex,
                                                     pathID
#ifdef @RENDER_MODE_DEPTH_STENCIL
                                                     ,
                                                     pathZIndex
#else
                                                     ,
                                                     v_windingWeight
#endif
                                                         VERTEX_CONTEXT_UNPACK);
#else // !@DRAW_INTERIOR_TRIANGLES
    float4 coverages;
    shouldDiscardVertex =
        !unpack_tessellated_path_vertex(@a_patchVertexData,
                                        @a_mirroredVertexData,
                                        _instanceID,
                                        pathID,
                                        vertexPosition
#ifndef @RENDER_MODE_DEPTH_STENCIL
                                        ,
                                        coverages
#else
                                        ,
                                        pathZIndex
#endif
                                            VERTEX_CONTEXT_UNPACK);
#ifndef @RENDER_MODE_DEPTH_STENCIL
#ifdef @ENABLE_FEATHER
    v_coverages = coverages;
#else
    v_coverages.xy = cast_float2_to_half2(coverages.xy);
#endif
#endif
#endif // !DRAW_INTERIOR_TRIANGLES

    uint2 paintData = STORAGE_BUFFER_LOAD2(@paintBuffer, pathID);

#if !defined(@FEATHER_ATLAS_BLIT) && !defined(@RENDER_MODE_DEPTH_STENCIL)
    // Encode the integral pathID as a "half" that we know the hardware will see
    // as a unique value in the fragment shader.
    v_pathID = id_bits_to_f16(pathID, uniforms.pathIDGranularity);

    // Indicate even-odd fill rule by making pathID negative.
    if ((paintData.x & PAINT_FLAG_EVEN_ODD_FILL) != 0u)
        v_pathID = -v_pathID;
#endif // !@FEATHER_ATLAS_BLIT && !@RENDER_MODE_DEPTH_STENCIL

    uint paintType = paintData.x & 0xfu;
#ifdef @ENABLE_CLIPPING
    if (@ENABLE_CLIPPING)
    {
        uint clipIDBits =
            (paintType == CLIP_UPDATE_PAINT_TYPE ? paintData.y : paintData.x) >>
            16;
        half clipID = id_bits_to_f16(clipIDBits, uniforms.pathIDGranularity);
        // Negative clipID means to update the clip buffer instead of the color
        // buffer.
        if (paintType == CLIP_UPDATE_PAINT_TYPE)
            clipID = -clipID;
#ifdef @FEATHER_ATLAS_BLIT
        v_clipID = clipID;
#else
        v_clipIDs.x = clipID;
#endif
    }
#endif
#ifdef @ENABLE_ADVANCED_BLEND
    if (@ENABLE_ADVANCED_BLEND)
    {
        v_blendMode = float((paintData.x >> 4) & 0xfu);
    }
#endif

    // Paint matrices operate on the fragment shader's "_fragCoord", which is
    // bottom-up in GL.
    float2 fragCoord = vertexPosition;
#ifdef @FRAMEBUFFER_BOTTOM_UP
    fragCoord.y = float(uniforms.renderTargetHeight) - fragCoord.y;
#endif

#ifdef @ENABLE_CLIP_RECT
    if (@ENABLE_CLIP_RECT)
    {
        // clipRectInverseMatrix transforms from pixel coordinates to a space
        // where the clipRect is the normalized rectangle: [-1, -1, 1, 1].
        float2x2 clipRectInverseMatrix = make_float2x2(
            STORAGE_BUFFER_LOAD4(@paintAuxBuffer,
                                 pathID * PAINT_AUX_ENTRY_ELEMENT_COUNT + 2u));
        float4 clipRectInverseTranslate =
            STORAGE_BUFFER_LOAD4(@paintAuxBuffer,
                                 pathID * PAINT_AUX_ENTRY_ELEMENT_COUNT + 3u);
#ifndef @RENDER_MODE_DEPTH_STENCIL
        v_clipRect =
            find_clip_rect_coverage_distances(clipRectInverseMatrix,
                                              clipRectInverseTranslate.xy,
                                              fragCoord);
#else  // !@RENDER_MODE_DEPTH_STENCIL => @RENDER_MODE_DEPTH_STENCIL
        set_clip_rect_plane_distances(clipRectInverseMatrix,
                                      clipRectInverseTranslate.xy,
                                      fragCoord CLIP_CONTEXT_UNPACK);
#endif // @RENDER_MODE_DEPTH_STENCIL
    }
#endif // ENABLE_CLIP_RECT

    // Unpack the paint once we have a position.
    if (paintType == SOLID_COLOR_PAINT_TYPE)
    {
        half4 color = unpackUnorm4x8(paintData.y);
        if (GENERATE_UNMULTIPLIED_PAINT_COLORS)
        {
            // naga can't handle "if (!SpecConst)" when transpiling spv to wgsl.
            // Use this if -> else construct instead so we don't have to negate
            // a specialization constant.
        }
        else
        {
            color.rgb *= color.a;
        }
        v_paint = float4(color);
    }
#if defined(@ENABLE_CLIPPING) && !defined(@FEATHER_ATLAS_BLIT)
    else if (@ENABLE_CLIPPING && paintType == CLIP_UPDATE_PAINT_TYPE)
    {
        half outerClipID =
            id_bits_to_f16(paintData.x >> 16, uniforms.pathIDGranularity);
        v_clipIDs.y = outerClipID;
    }
#endif
    else
    {
        float2x2 paintMatrix = make_float2x2(
            STORAGE_BUFFER_LOAD4(@paintAuxBuffer,
                                 pathID * PAINT_AUX_ENTRY_ELEMENT_COUNT));
        float4 paintTranslate =
            STORAGE_BUFFER_LOAD4(@paintAuxBuffer,
                                 pathID * PAINT_AUX_ENTRY_ELEMENT_COUNT + 1u);
        float2 paintCoord = MUL(paintMatrix, fragCoord) + paintTranslate.xy;
        if (paintType == LINEAR_GRADIENT_PAINT_TYPE ||
            paintType == RADIAL_GRADIENT_PAINT_TYPE)
        {
            // v_paint.a contains "-row" of the gradient ramp at texel center,
            // in normalized space.
            v_paint.a = -uintBitsToFloat(paintData.y);
            // abs(v_paint.b) contains either:
            //   - 2 if the gradient ramp spans an entire row.
            //   - x0 of the gradient ramp in normalized space, if it's a simple
            //   2-texel ramp.
            float gradientSpan = paintTranslate.z;
            // gradientSpan is either ~1 (complex gradients span the whole width
            // of the texture minus 1px), or 1/GRAD_TEXTURE_WIDTH (simple
            // gradients span 1px).
            if (gradientSpan > .9)
            {
                // Complex ramps span an entire row. Set it to 2 to convey this.
                v_paint.b = 2.;
            }
            else
            {
                // This is a simple ramp.
                v_paint.b = paintTranslate.w;
            }
            if (paintType == LINEAR_GRADIENT_PAINT_TYPE)
            {
                // The paint is a linear gradient.
                v_paint.g = .0;
                v_paint.r = paintCoord.x;
            }
            else
            {
                // The paint is a radial gradient. Mark v_paint.b negative to
                // indicate this to the fragment shader. (v_paint.b can't be
                // zero because the gradient ramp is aligned on pixel centers,
                // so negating it will always produce a negative number.)
                v_paint.b = -v_paint.b;
                v_paint.rg = paintCoord.xy;
            }
        }
    }
#ifdef @EMULATE_DYNAMIC_COLOR_WRITE_DISABLE
    if (@EMULATE_DYNAMIC_COLOR_WRITE_DISABLE)
    {
        // Zeroing v_paint is all we need to disable color write; float4(0) gets
        // interpreted by the fragment shader as a fully transparent
        // SOLID_COLOR_PAINT_TYPE, and then discarded at the blend step.
        v_paint *= pushConstants.colorWriteEnable;
    }
#endif

#if defined(@ENABLE_MODULATED_IMAGE)
    if (@ENABLE_MODULATED_IMAGE && (paintData.x & PAINT_FLAG_HAS_IMAGE) != 0u)
    {
        float2x2 paintMatrix = make_float2x2(
            STORAGE_BUFFER_LOAD4(@paintAuxBuffer,
                                 pathID * PAINT_AUX_ENTRY_ELEMENT_COUNT + 4u));
        float4 paintTranslateAndLOD =
            STORAGE_BUFFER_LOAD4(@paintAuxBuffer,
                                 pathID * PAINT_AUX_ENTRY_ELEMENT_COUNT + 5u);
        float2 paintCoord =
            MUL(paintMatrix, fragCoord) + paintTranslateAndLOD.xy;

        v_image =
            float3(paintCoord.x, paintCoord.y, 1. + paintTranslateAndLOD.z);
    }
    else
    {
        v_image = float3(0.0, 0.0, 0.0);
    }
#endif

    float4 pos;
    if (!shouldDiscardVertex)
    {
        pos = RENDER_TARGET_COORD_TO_CLIP_COORD(vertexPosition);
#ifdef @POST_INVERT_Y
        pos.y = -pos.y;
#endif
#ifdef @RENDER_MODE_DEPTH_STENCIL
        pos.z = normalize_z_index(pathZIndex);
#elif defined(@RENDER_MODE_CLOCKWISE_ATOMIC)
        uint4 coverageData =
            STORAGE_BUFFER_LOAD4(@pathBuffer, pathID * 4u + 3u);
        v_coveragePlacement = coverageData.xy;
        v_coverageCoord = vertexPosition + uintBitsToFloat(coverageData.zw);
#endif
    }
    else
    {
        pos = float4(uniforms.vertexDiscardValue,
                     uniforms.vertexDiscardValue,
                     uniforms.vertexDiscardValue,
                     uniforms.vertexDiscardValue);
    }

    VARYING_PACK(v_paint);
#if defined(@ENABLE_MODULATED_IMAGE)
    VARYING_PACK(v_image);
#endif
#ifdef @FEATHER_ATLAS_BLIT
    VARYING_PACK(v_atlasCoord);
#elif !defined(@RENDER_MODE_DEPTH_STENCIL)
#ifdef @DRAW_INTERIOR_TRIANGLES
    VARYING_PACK(v_windingWeight);
#else
    VARYING_PACK(v_coverages);
#endif //@DRAW_INTERIOR_TRIANGLES
    VARYING_PACK(v_pathID);
#endif // !@RENDER_MODE_DEPTH_STENCIL

#ifdef @ENABLE_CLIPPING
#ifdef @FEATHER_ATLAS_BLIT
    VARYING_PACK(v_clipID);
#else
    VARYING_PACK(v_clipIDs);
#endif
#endif // @ENABLE_CLIPPING
#if defined(@ENABLE_CLIP_RECT) && !defined(@RENDER_MODE_DEPTH_STENCIL)
    VARYING_PACK(v_clipRect);
#endif
#ifdef @ENABLE_ADVANCED_BLEND
    VARYING_PACK(v_blendMode);
#endif
#ifdef @RENDER_MODE_CLOCKWISE_ATOMIC
    VARYING_PACK(v_coveragePlacement);
    VARYING_PACK(v_coverageCoord);
#endif
    EMIT_VERTEX(pos);
}
#endif

#ifdef @FRAGMENT

FRAG_STORAGE_BUFFER_BLOCK_BEGIN
FRAG_STORAGE_BUFFER_BLOCK_END

// Add a function here for fragments to unpack the paint since we're the ones
// who packed it in the vertex shader.
INLINE half4 find_paint_color(float4 paint,
#ifdef @ENABLE_MODULATED_IMAGE
                              float3 image,
#endif
                              float coverage FRAGMENT_CONTEXT_DECL)
{
    half4 color;
    if (paint.a >= .0) // Is the paint a solid color?
    {
        // The vertex shader will have premultiplied 'paint' (or not) based on
        // GENERATE_UNMULTIPLIED_PAINT_COLORS.
        color = cast_float4_to_half4(paint);
        if (GENERATE_UNMULTIPLIED_PAINT_COLORS)
            color.a *= coverage;
        else
            color *= coverage;
    }
    else // Paint is a gradient (linear or radial)?
    {
        float t =
            paint.b > .0 ? /*linear*/ paint.r : /*radial*/ length(paint.rg);
        t = clamp(t, .0, 1.);
        float span = abs(paint.b);
        float x = span > 1.
                      ? /*entire row*/ (1. - 1. / GRAD_TEXTURE_WIDTH) * t +
                            (.5 / GRAD_TEXTURE_WIDTH)
                      : /*two texels*/ (1. / GRAD_TEXTURE_WIDTH) * t + span;
        float row = -paint.a;
        // Our gradient texture is not mipmapped. Issue a texture-sample that
        // explicitly does not find derivatives for LOD computation.
        color =
            TEXTURE_SAMPLE_LOD(@gradTexture, gradSampler, float2(x, row), .0);
        color.a *= coverage;
        // Gradients are always unmultiplied so we don't lose color data while
        // doing the hardware filter.
        if (GENERATE_UNMULTIPLIED_PAINT_COLORS)
        {
            // naga can't handle "if (!SpecConst)" when transpiling spv to wgsl.
            // Use this if -> else construct instead so we don't have to
            // negate a specialization constant.
        }
        else
        {
            color.rgb *= color.a;
        }
    }

#if defined(@ENABLE_MODULATED_IMAGE)
    if (@ENABLE_MODULATED_IMAGE && image.z > 0.0)
    {
        half lod = image.z - 1.;
        half4 imageColor = TEXTURE_SAMPLE_DYNAMIC_LOD(@imageTexture,
                                                      imageSampler,
                                                      image.rg,
                                                      lod);

        // Images are always premultiplied so the (transparent) background color
        // doesn't bleed into the edges during the hardware filter.
        if (GENERATE_UNMULTIPLIED_PAINT_COLORS)
            imageColor = make_half4(unmultiply_rgb(imageColor), imageColor.a);

        color *= imageColor;
    }
#endif
    return color;
}

#if !defined(@DRAW_INTERIOR_TRIANGLES) && !defined(@FEATHER_ATLAS_BLIT)

// Add functions here for fragments to unpack and evaluate coverage since we're
// the ones who packed the coverage components in the vertex shader.
INLINE half find_stroke_coverage(COVERAGE_TYPE coverages TEXTURE_CONTEXT_DECL)
{
#ifdef @ENABLE_FEATHER
    if (@ENABLE_FEATHER && is_feathered_stroke(coverages))
        return eval_feathered_stroke(coverages TEXTURE_CONTEXT_FORWARD);
    else
#endif // @ENABLE_FEATHER
        return min(coverages.x, coverages.y);
}

INLINE half find_fill_coverage(COVERAGE_TYPE coverages TEXTURE_CONTEXT_DECL)
{
#if defined(@ENABLE_FEATHER)
    if (@ENABLE_FEATHER && is_feathered_fill(coverages))
        return eval_feathered_fill(coverages TEXTURE_CONTEXT_FORWARD);
    else
#endif // @ENABLE_FEATHER
        return coverages.x;
}

INLINE half find_frag_coverage(COVERAGE_TYPE coverages TEXTURE_CONTEXT_DECL)
{
    if (is_stroke(coverages))
        return find_stroke_coverage(coverages TEXTURE_CONTEXT_FORWARD);
    else // Fill. (Back-face culling handles the sign of coverages.x.)
        return find_fill_coverage(coverages TEXTURE_CONTEXT_FORWARD);
}

INLINE half apply_frag_coverage(half initialCoverage,
                                COVERAGE_TYPE coverages TEXTURE_CONTEXT_DECL)
{
    if (is_stroke(coverages))
    {
        half fragCoverage =
            find_stroke_coverage(coverages TEXTURE_CONTEXT_FORWARD);
        return max(fragCoverage, initialCoverage);
    }
    else // Fill. (Back-face culling handles the sign of coverages.x.)
    {
        half fragCoverage =
            find_fill_coverage(coverages TEXTURE_CONTEXT_FORWARD);
        return initialCoverage + fragCoverage;
    }
}

#endif // !@DRAW_INTERIOR_TRIANGLES && !@FEATHER_ATLAS_BLIT

#endif // @FRAGMENT