/* SPDX-License-Identifier: GPL-2.0-or-later */

#include "BKE_pointcloud.h"

#include "BLI_task.hh"

#include "node_geometry_util.hh"

namespace blender::nodes::node_geo_points_cc {

static void node_declare(NodeDeclarationBuilder &b)
{
  b.add_input<decl::Int>(N_("Count"))
      .default_value(1)
      .description(N_("The number of points to create"))
      .min(0);
  b.add_input<decl::Vector>(N_("Position"))
      .supports_field()
      .default_value(float3(0.0f))
      .description(N_("The positions of the new points"));
  b.add_input<decl::Float>(N_("Radius"))
      .supports_field()
      .supports_field()
      .default_value(float(0.1f))
      .description(N_("The radii of the new points"));
  b.add_output<decl::Geometry>(N_("Geometry"));
}

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 UNUSED(mask),
                               ResourceScope &UNUSED(scope)) const
  {
    const bke::AttributeFieldInput *attribute_field_input =
        dynamic_cast<const bke::AttributeFieldInput *>(&field_input);
    if (attribute_field_input == nullptr) {
      return {};
    }
    if (attribute_field_input->attribute_name() != "Index" &&
        attribute_field_input->attribute_name() != "id") {
      return {};
    }
    Array<int> index(this->points_num_);
    threading::parallel_for(IndexRange(this->points_num_), 1, [&](const IndexRange range) {
      for (const int i : range) {
        index[i] = i;
      }
    });
    return VArray<int>::ForContainer(std::move(index));
  }
};

static void node_geo_exec(GeoNodeExecParams params)
{
  const int count = params.extract_input<int>("Count");
  if (count <= 0) {
    params.error_message_add(NodeWarningType::Warning, TIP_("Point count should be at least 1"));
    params.set_default_remaining_outputs();
    return;
  }

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

  PointCloud *new_point_cloud = BKE_pointcloud_new_nomain(count);
  GeometrySet geometry_set = GeometrySet::create_with_pointcloud(new_point_cloud);
  PointCloudComponent &points = geometry_set.get_component_for_write<PointCloudComponent>();
  OutputAttribute_Typed<float3> output_position = points.attribute_try_get_for_output_only<float3>(
      "position", ATTR_DOMAIN_POINT);
  OutputAttribute_Typed<float> output_radii = points.attribute_try_get_for_output_only<float>(
      "radius", ATTR_DOMAIN_POINT);

  // GeometryComponentFieldContext context{points, ATTR_DOMAIN_POINT};
  PointsFieldContext context{count};
  fn::FieldEvaluator evaluator{context, count};
  evaluator.add_with_destination(position_field, output_position.as_span());
  evaluator.add_with_destination(radius_field, output_radii.as_span());
  evaluator.evaluate();

  output_position.save();
  output_radii.save();
  params.set_output("Geometry", geometry_set);
}

}  // namespace blender::nodes::node_geo_points_cc

void register_node_type_geo_points()
{
  namespace file_ns = blender::nodes::node_geo_points_cc;
  static bNodeType ntype;
  geo_node_type_base(&ntype, GEO_NODE_POINTS, "Points", NODE_CLASS_GEOMETRY);
  ntype.geometry_node_execute = file_ns::node_geo_exec;
  ntype.declare = file_ns::node_declare;
  nodeRegisterType(&ntype);
}