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

#include "BKE_pointcloud.hh"

#include "DNA_pointcloud_types.h"

#include "node_geometry_util.hh"

namespace blender::nodes::node_geo_points_cc {

static void node_declare(NodeDeclarationBuilder &b)
{
  b.add_input<decl::Int>("Count"_ustr)
      .default_value(1)
      .min(0)
      .description("The number of points to create");
  b.add_input<decl::Vector>("Position"_ustr)
      .subtype(PROP_TRANSLATION)
      .default_value(float3(0.0f))
      .structure_type(StructureType::Field)
      .description("The positions of the new points");
  b.add_input<decl::Float>("Radius"_ustr)
      .min(0.0f)
      .default_value(0.1f)
      .subtype(PROP_DISTANCE)
      .structure_type(StructureType::Field)
      .description("The radii of the new points");
  b.add_output<decl::Geometry>("Points"_ustr, "Geometry"_ustr);
}

class PointsFieldContext : public FieldContext {
 private:
  int points_num_;

 public:
  PointsFieldContext(const int points_num) : points_num_(points_num) {}

  int64_t points_num() const
  {
    return points_num_;
  }

  GVArray get_varray_for_input(const FieldInput &field_input,
                               const IndexMask &mask,
                               ResourceScope & /*scope*/) const override
  {
    const bke::IDAttributeFieldInput *id_field_input =
        dynamic_cast<const bke::IDAttributeFieldInput *>(&field_input);

    const fn::IndexFieldInput *index_field_input = dynamic_cast<const fn::IndexFieldInput *>(
        &field_input);

    if (id_field_input == nullptr && index_field_input == nullptr) {
      return {};
    }

    return fn::IndexFieldInput::get_index_varray(mask);
  }
};

static void node_geo_exec(GeoNodeExecParams params)
{
  const int count = params.extract_input<int>("Count"_ustr);
  if (count <= 0) {
    params.set_default_remaining_outputs();
    return;
  }

  Field<float3> position_field = params.extract_input<Field<float3>>("Position"_ustr);
  Field<float> radius_field = params.extract_input<Field<float>>("Radius"_ustr);

  PointCloud *points = BKE_pointcloud_new_nomain(count);
  MutableAttributeAccessor attributes = points->attributes_for_write();

  PointsFieldContext context{count};
  fn::FieldEvaluator evaluator{context, count};
  evaluator.add_with_destination(position_field, points->positions_for_write());
  if (radius_field.depends_on_input()) {
    AttributeWriter<float> output_radii = attributes.lookup_or_add_for_write<float>(
        "radius", AttrDomain::Point);
    evaluator.add_with_destination(radius_field, output_radii.varray);
    evaluator.evaluate();
    output_radii.finish();
  }
  else {
    evaluator.evaluate();
    const float radius = fn::evaluate_constant_field(radius_field);
    attributes.add<float>("radius", AttrDomain::Point, bke::AttributeInitValue(radius));
  }

  params.set_output("Geometry"_ustr, GeometrySet::from_pointcloud(points));
}

static void node_register()
{
  static bke::bNodeType ntype;
  geo_node_type_base(&ntype, "GeometryNodePoints"_ustr, GEO_NODE_POINTS);
  ntype.ui_name = "Points";
  ntype.ui_description = "Generate a point cloud with positions and radii defined by fields";
  ntype.enum_name_legacy = "POINTS";
  ntype.nclass = NODE_CLASS_GEOMETRY;
  ntype.geometry_node_execute = node_geo_exec;
  ntype.declare = node_declare;
  bke::node_register_type(ntype);
}
NOD_REGISTER_NODE(node_register)

}  // namespace blender::nodes::node_geo_points_cc
