/*
 * Copyright 2024 Rive
 */

#pragma once

#include "rive/math/simd.hpp"
#include "rive/enums.hpp"
#include <vector>

namespace rive::gpu
{

// Adding rectangles allows a grouping type to be specified - this determines
// whether a rectangle is allowed to share a group with another that it
// overlaps.
enum class GroupingType : bool
{
    // The current rectangle cannot be grouped with any others for any reason,
    // so no additional testing/logic is necessary. (this is the default for
    // most render modes)
    disjoint,

    // The current rectangle is potentially allowed to group with others it
    // overlaps, assuming the test of "overlap bits" succeeds.
    overlapAllowed,
};

// Tile within an intersection board that manages a set of rectangles and their
// groupIndex. From a given rectangle, finds the max groupIndex in the set of
// internal rectangles it intersects and, when needed, also finds the full set
// of overlap bits from any interesections. The size is <= 255 so we can store
// bounding box coordinates in 8 bits.
class IntersectionTile
{
public:
    // The width and height of an intersection tile (must fit in an unsigned
    // byte).
    constexpr static int TILE_DIM = 255;
    static_assert(TILE_DIM <= std::numeric_limits<uint8_t>::max());

    // Performance is better passing these two vectors as a struct than using
    // in/out parameters.
    struct FindResult
    {
        int16x8 maxGroupIndices = 0;
        uint16x8 overlapBits = 0;
    };

    void reset(int left,
               int top,
               int16_t baselineGroupIndex = 0,
               uint16_t baselineOverlapBits = 0);

    template <GroupingType Type>
    void addRectangle(int4 ltrb,
                      int16_t groupIndex,
                      uint16_t currentRectangleOverlapBits);

#ifdef WITH_RIVE_TOOLS
    // a testing-specific version that validates after doing the add
    template <GroupingType Type>
    void testingOnly_addRectangleAndValidate(
        int4 ltrb,
        int16_t groupIndex,
        uint16_t currentRectangleOverlapBits);
#endif

    // Accumulate local maximum intersecting group indices for the given
    // rectangle in each channel of a int16x8. "runningMaxGroupIndices" is a
    // running set of local maximums if the IntersectionBoard also ran this same
    // test on other tile(s) that the rectangle touched. The absolute maximum
    // group index that this rectangle intersects with will be
    // simd::reduce_max(result.maxGroupIndices), and the relevant overlap bits
    // will be simd::reduce_or(result.overlapBits)
    template <GroupingType Type>
    FindResult findMaxIntersectingGroupIndex(int4 ltrb,
                                             FindResult runningResult) const;

    // The following were exposed for unit testing

#ifdef WITH_RIVE_TOOLS
    size_t testingOnly_rectangleCount() const { return m_rectangleCount; }
    int16_t testingOnly_baselineGroupIndex() const
    {
        return m_baselineGroupIndex;
    }
    uint16_t testingOnly_baselineOverlapBits() const
    {
        return m_baselineOverlapBits;
    }

    int16_t testingOnly_maxGroupIndex() const { return m_maxGroupIndex; }
    uint16_t testingOnly_overlapBitsForMaxGroup() const
    {
        return m_overlapBitsForMaxGroup;
    }
#endif

private:
    // Set the baseline to equal the current max group index
    template <GroupingType Type>
    void updateBaselineToMaxGroupIndex(uint16_t additionalBaselineOverlapBits);

    int4 m_topLeft = {};
    int16_t m_baselineGroupIndex = 0;
    uint16_t m_baselineOverlapBits = 0;
    int16_t m_maxGroupIndex = 0;
    uint16_t m_overlapBitsForMaxGroup = 0;
    size_t m_rectangleCount = 0;

    // We may need a bias to get the tile coordinates to fit within a signed
    // 8-bit value (because base-level SSE does not have unsigned comparisons)
    constexpr static int TILE_EDGE_BIAS =
        (TILE_DIM > std::numeric_limits<int8_t>::max()) ? -128 : 0;

    // How many rectangles/groupIndices are in each chunk of data?
    constexpr static size_t CHUNK_SIZE = 8;

    // Chunk of 8 rectangles encoded as [L, T, TILE_DIM - R, TILE_DIM - B],
    // relative to m_left and m_top. The data is also transposed:
    // [L0..L7, T0..T7, -R0..R7, -B0..B7].
    std::vector<int8x32> m_edges;
    static_assert(sizeof(m_edges[0]) == CHUNK_SIZE * 4);

    // Chunk of 8 groupIndices corresponding to the above edges.
    std::vector<int16x8> m_groupIndices;
    static_assert(sizeof(m_groupIndices[0]) == CHUNK_SIZE * 2);

    // Chunk of 8 sets of overlap bits corresponding to the above edges.
    std::vector<uint16x8> m_overlapBits;
    static_assert(sizeof(m_overlapBits[0]) == CHUNK_SIZE * 2);
};

// Manages a set of rectangles and their groupIndex across a variable-sized
// viewport. Each time a rectangle is added, assigns and returns a groupIndex
// that is one larger than the max groupIndex in the set of existing rectangles
// it intersects.
class IntersectionBoard
{
public:
    constexpr static auto TILE_DIM = IntersectionTile::TILE_DIM;

    IntersectionBoard(GroupingType groupingType) : m_groupingType(groupingType)
    {}

    void resizeAndReset(uint32_t viewportWidth, uint32_t viewportHeight);

    // Adds a rectangle to the internal set and assigns it "layerCount"
    // contiguous group indices, beginning one larger than the max groupIndex in
    // the set of existing rectangles it intersects.
    //
    // Returns the first newly assigned groupIndex for the added rectangle.
    // If it does not intersect with any other rectangles, this groupIndex is 1.
    //
    // It is the caller's responsibility to not insert more rectangles than can
    // fit in a signed 16-bit integer. (The result is signed because SSE doesn't
    // have an unsigned max instruction.)
    int16_t addRectangle(
        int4 ltrb,
        uint16_t
            currentRectangleOverlapBits, // "what is allowed to overlap me?"
        uint16_t
            disallowedOverlapBitsMask, // "what am I not allowed to overlap?"
        int16_t layerCount);

    int16_t addRectangle(int4 ltrb, int16_t layerCount = 1)
    {
        return addRectangle(ltrb, 0, 0, layerCount);
    }

    GroupingType groupingType() const { return m_groupingType; }

private:
    template <GroupingType Type>
    int16_t addRectangle(int4 ltrb,
                         uint16_t currentRectangleOverlapBits,
                         uint16_t disallowedOverlapBitsMask,
                         int16_t layerCount);

    GroupingType m_groupingType{};
    int2 m_viewportSize{};
    int32_t m_cols = 0;
    int32_t m_rows = 0;
    std::vector<IntersectionTile> m_tiles;
};
} // namespace rive::gpu