/*
 * Copyright 2022 Rive
 */

#pragma once

#include "rive/renderer/gpu.hpp"
#include "rive/renderer.hpp"
#include <array>

namespace rive::gpu
{

// Copies an array of colors or stops for a gradient.
// Stores the data locally if there are 4 values or fewer.
// Spills onto the heap if there are >4 values.
template <typename T> class GradDataArray
{
public:
  static_assert(std::is_trivial_v<T> && std::is_standard_layout<T>());

  GradDataArray(const T data[], size_t count)
  {
    m_data = count <= m_localData.size() ? m_localData.data() : new T[count];
    memcpy(m_data, data, count * sizeof(T));
  }

  // Allocate without initializing (caller must fill in data).
  explicit GradDataArray(size_t count)
  {
    m_data = count <= m_localData.size() ? m_localData.data() : new T[count];
  }

  GradDataArray(GradDataArray&& other)
  {
    if (other.m_data == other.m_localData.data())
    {
      m_localData = other.m_localData;
      m_data = m_localData.data();
    }
    else
    {
      m_data = other.m_data;
      other.m_data = other.m_localData.data(); // Don't delete[] other.m_data.
    }
  }

  ~GradDataArray()
  {
    if (m_data != m_localData.data())
    {
      delete[] m_data;
    }
  }

  const T* get() const { return m_data; }
  const T operator[](size_t i) const { return m_data[i]; }
  T& operator[](size_t i) { return m_data[i]; }

private:
  std::array<T, 4> m_localData;
  T* m_data;
};

// RenderShader implementation for Rive's pixel local storage renderer.
class Gradient : public LITE_RTTI_OVERRIDE(RenderShader, Gradient)
{
public:
  static rcp<Gradient> MakeLinear(
      float sx,
      float sy,
      float ex,
      float ey,
      const ColorInt colors[], // [count]
      const float stops[],     // [count]
      size_t count);

  static rcp<Gradient> MakeRadial(
      float cx,
      float cy,
      float radius,
      const ColorInt colors[], // [count]
      const float stops[],     // [count]
      size_t count);

  static rcp<Gradient> MakeAngular(
      const ColorInt colors[], // [count]
      const float stops[],     // [count]
      size_t count);

  PaintType paintType() const { return m_paintType; }
  const float* coeffs() const { return m_coeffs.data(); }
  const ColorInt* colors() const { return m_colors.get(); }
  const float* stops() const { return m_stops.get(); }
  size_t count() const { return m_count; }
  bool isOpaque() const;

  // Get or create a modulated variant of this gradient.
  // Caches the last-used modulated gradient for efficient reuse when the same
  // modulation is requested multiple times (e.g., multiple draws in one
  // frame).
  rcp<Gradient> getModulated(float opacity, ColorInt color = 0xFFFFFFFF) const;

private:
  Gradient(
      PaintType paintType,
      GradDataArray<ColorInt>&& colors, // [count]
      GradDataArray<float>&& stops,     // [count]
      size_t count,
      float coeffX,
      float coeffY,
      float coeffZ)
      : m_paintType(paintType)
      , m_colors(std::move(colors))
      , m_stops(std::move(stops))
      , m_count(count)
      , m_coeffs{coeffX, coeffY, coeffZ}
  {
    assert(
        paintType == gpu::PaintType::linearGradient ||
        paintType == gpu::PaintType::radialGradient ||
        paintType == gpu::PaintType::angularGradient);
  }

  PaintType m_paintType; // Specifically, linearGradient or radialGradient.
  GradDataArray<ColorInt> m_colors;
  GradDataArray<float> m_stops;
  size_t m_count;
  std::array<float, 3> m_coeffs;
  mutable gpu::TriState m_isOpaque = gpu::TriState::unknown;

  // Single-entry cache for last-used modulated gradient
  mutable rcp<Gradient> m_lastModulatedGradient;
  mutable float
      m_lastModulatedOpacity = -1.0f; // -1 as sentinel (valid range is 0..1)
  mutable ColorInt m_lastModulatedColor = 0xFFFFFFFF;
};

} // namespace rive::gpu