/* SPDX-FileCopyrightText: 2023 Blender Authors
 *
 * SPDX-License-Identifier: GPL-2.0-or-later */

#include "BKE_curves.hh"

#include "node_geometry_util.hh"

namespace blender::nodes::node_geo_offset_point_in_curve_cc {

static void node_declare(NodeDeclarationBuilder &b)
{
  b.add_input<decl::Int>("Point Index"_ustr)
      .default_input_type(NODE_DEFAULT_INPUT_INDEX_FIELD)
      .description("The index of the control point to evaluate. Defaults to the current index")
      .structure_type(StructureType::Field);
  b.add_input<decl::Int>("Offset"_ustr)
      .structure_type(StructureType::Field)
      .description("The number of control points along the curve to traverse");
  b.add_output<decl::Bool>("Is Valid Offset"_ustr)
      .structure_type(StructureType::Field)
      .propagate_references()
      .description(
          "Whether the input control point plus the offset is a valid index of the "
          "original curve");
  b.add_output<decl::Int>("Point Index"_ustr)
      .structure_type(StructureType::Field)
      .propagate_references()
      .description(
          "The index of the control point plus the offset within the entire "
          "curves data-block");
}

class ControlPointNeighborFieldInput final : public bke::GeometryFieldInput {
 private:
  const Field<int> index_;
  const Field<int> offset_;

 public:
  ControlPointNeighborFieldInput(Field<int> index, Field<int> offset)
      : GeometryFieldInput(CPPType::get<int>(), "Offset Point in Curve"),
        index_(std::move(index)),
        offset_(std::move(offset))
  {
  }

  GVArray get_varray_for_context(const bke::GeometryFieldContext &context,
                                 const IndexMask &mask) const final
  {
    const bke::CurvesGeometry *curves_ptr = context.curves_or_strokes();
    if (!curves_ptr) {
      return {};
    }
    const bke::CurvesGeometry &curves = *curves_ptr;

    const OffsetIndices points_by_curve = curves.points_by_curve();
    const VArray<bool> cyclic = curves.cyclic();
    const Array<int> parent_curves = curves.point_to_curve_map();

    fn::FieldEvaluator evaluator{context, &mask};
    evaluator.add(index_);
    evaluator.add(offset_);
    evaluator.evaluate();
    const VArray<int> indices = evaluator.get_evaluated<int>(0);
    const VArray<int> offsets = evaluator.get_evaluated<int>(1);

    Array<int> output(mask.min_array_size());
    mask.foreach_index(
        [&](const int i_selection) {
          const int point = std::clamp(indices[i_selection], 0, curves.points_num() - 1);
          const int curve = parent_curves[point];
          const IndexRange curve_points = points_by_curve[curve];
          const int shifted_point = point + offsets[i_selection];

          if (cyclic[curve]) {
            const int point_index_in_curve = shifted_point - curve_points.start();
            output[i_selection] = curve_points.start() +
                                  math::mod_periodic<int>(point_index_in_curve,
                                                          curve_points.size());
            return;
          }
          output[i_selection] = std::clamp(shifted_point, 0, curves.points_num() - 1);
        },
        exec_mode::grain_size(512));

    return VArray<int>::from_container(std::move(output));
  }

  void foreach_recursive_field(FunctionRef<void(const GField &)> fn) const override
  {
    fn(index_);
    fn(offset_);
  }
};

class OffsetValidFieldInput final : public bke::GeometryFieldInput {
 private:
  const Field<int> index_;
  const Field<int> offset_;

 public:
  OffsetValidFieldInput(Field<int> index, Field<int> offset)
      : GeometryFieldInput(CPPType::get<bool>(), "Offset Valid"),
        index_(std::move(index)),
        offset_(std::move(offset))
  {
  }

  GVArray get_varray_for_context(const bke::GeometryFieldContext &context,
                                 const IndexMask &mask) const final
  {
    const bke::CurvesGeometry *curves_ptr = context.curves_or_strokes();
    if (!curves_ptr) {
      return {};
    }
    const bke::CurvesGeometry &curves = *curves_ptr;

    const VArray<bool> cyclic = curves.cyclic();
    const OffsetIndices points_by_curve = curves.points_by_curve();
    const Array<int> parent_curves = curves.point_to_curve_map();

    fn::FieldEvaluator evaluator{context, &mask};
    evaluator.add(index_);
    evaluator.add(offset_);
    evaluator.evaluate();
    const VArray<int> indices = evaluator.get_evaluated<int>(0);
    const VArray<int> offsets = evaluator.get_evaluated<int>(1);

    Array<bool> output(mask.min_array_size());
    mask.foreach_index(
        [&](const int i_selection) {
          const int i_point = indices[i_selection];
          if (!curves.points_range().contains(i_point)) {
            output[i_selection] = false;
            return;
          }

          const int i_curve = parent_curves[i_point];
          const IndexRange curve_points = points_by_curve[i_curve];
          if (cyclic[i_curve]) {
            output[i_selection] = true;
            return;
          }
          output[i_selection] = curve_points.contains(i_point + offsets[i_selection]);
        },
        exec_mode::grain_size(512));
    return VArray<bool>::from_container(std::move(output));
  }

  void foreach_recursive_field(FunctionRef<void(const GField &)> fn) const override
  {
    fn(index_);
    fn(offset_);
  }
};

static void node_geo_exec(GeoNodeExecParams params)
{
  Field<int> index = params.extract_input<Field<int>>("Point Index"_ustr);
  Field<int> offset = params.extract_input<Field<int>>("Offset"_ustr);

  if (params.output_is_required("Point Index"_ustr)) {
    params.set_output("Point Index"_ustr,
                      Field<int>::from_input<ControlPointNeighborFieldInput>(index, offset));
  }
  if (params.output_is_required("Is Valid Offset"_ustr)) {
    params.set_output("Is Valid Offset"_ustr,
                      Field<bool>::from_input<OffsetValidFieldInput>(index, offset));
  }
}

static void node_register()
{
  static bke::bNodeType ntype;
  geo_node_type_base(
      &ntype, "GeometryNodeOffsetPointInCurve"_ustr, GEO_NODE_OFFSET_POINT_IN_CURVE);
  ntype.ui_name = "Offset Point in Curve";
  ntype.ui_description = "Offset a control point index within its curve";
  ntype.enum_name_legacy = "OFFSET_POINT_IN_CURVE";
  ntype.nclass = NODE_CLASS_INPUT;
  ntype.geometry_node_execute = node_geo_exec;
  ntype.declare = node_declare;
  ntype.default_width = bke::NodeWidth::_160;
  bke::node_register_type(ntype);
}
NOD_REGISTER_NODE(node_register)

}  // namespace blender::nodes::node_geo_offset_point_in_curve_cc
