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

#include "BKE_curves.hh"

#include "UI_interface.h"
#include "UI_resources.h"

#include "node_geometry_util.hh"

namespace blender::nodes::node_geo_curve_split_curves_cc {

static void node_declare(NodeDeclarationBuilder &b)
{
  b.add_input<decl::Geometry>(N_("Curves"));
  b.add_input<decl::Bool>(N_("Selection")).hide_value().supports_field();
  b.add_input<decl::Bool>(N_("Remove Previous")).supports_field();
  b.add_input<decl::Bool>(N_("Remove Next")).default_value(true).supports_field();

  b.add_output<decl::Geometry>(N_("Curves"));
}

static void split_curves(GeometrySet &geometry_set,
                         const Field<bool> &selection_field,
                         const Field<bool> &previous_field,
                         const Field<bool> &next_field)
{
  const CurveComponent &component = *geometry_set.get_component_for_read<CurveComponent>();
  const Curves &curves_id = *component.get_for_read();
  bke::CurvesGeometry curves = bke::CurvesGeometry::wrap(curves_id.geometry);

  if (curves.curve_size == 0) {
    return;
  }

  const int domain_size = component.attribute_domain_size(ATTR_DOMAIN_POINT);
  GeometryComponentFieldContext context{component, ATTR_DOMAIN_POINT};
  fn::FieldEvaluator evaluator{context, domain_size};
  evaluator.add(selection_field);
  evaluator.add(previous_field);
  evaluator.add(next_field);
  evaluator.evaluate();
  const VArray<bool> selection = evaluator.get_evaluated<bool>(0);
  const VArray<bool> previous = evaluator.get_evaluated<bool>(1);
  const VArray<bool> next = evaluator.get_evaluated<bool>(2);

  if (selection.is_single() && !selection.get_internal_single()) {
    return;
  }

  /* new_offsets are curve offsets of for the output curve. old_offsets contains where each new
   * offset starts in the context of the old point domain to facilitate copying data from old
   * points to new points. */

  Vector<int> old_offsets;
  Vector<int> new_offsets;
  Vector<int> old_curve;

  int point = 0;
  int new_point = 0;
  for (const int i_curve : curves.curves_range()) {
    const IndexRange curve_index_range = curves.range_for_curve(i_curve);
    for (const int i_point : curve_index_range) {
      /* start point of each existing curve. Always a new curve*/
      if (i_point == 0) {
        old_curve.append(i_curve);
        old_offsets.append(point++);
        new_offsets.append(new_point++);
        continue;
      }

      const bool select_i = selection[i_point];
      const bool prev_i = previous[i_point];
      const bool select_i_minus_1 = selection[i_point - 1];
      const bool next_i_minus_1 = next[i_point - 1];

      /* previous was selected but there is a segment between current and previous, so create a new
       * curve starting on the last point of the previous curve*/
      if (select_i_minus_1 && !next_i_minus_1 && !prev_i) {
        old_curve.append(i_curve);
        old_offsets.append(point - 1);
        new_offsets.append(new_point++);

        old_curve.append(i_curve);
        old_offsets.append(point++);
        new_offsets.append(new_point++);
        continue;
      }

      /* If previous point was selected with next removed or current point was selected with
       * previous removed.*/
      if (select_i_minus_1 && next_i_minus_1 || select_i && prev_i) {
        old_curve.append(i_curve);
        old_offsets.append(point++);
        new_offsets.append(new_point++);
        continue;
      }

      /* No new curve, just move along the list. */
      point++;
      new_point++;
    }
  }
  /* If the very last point of the last curve was select but retained it's segments, add one more
   * point */
  const int last_point = curves.point_size - 1;
  if (selection[last_point] && !previous[last_point] && !next[last_point - 1]) {
    old_curve.append(curves.curve_size - 1);
    old_offsets.append(point - 1);
    new_offsets.append(new_point++);
  }

  /* Add an additional entry to the offsets for total or virtual "next item" */
  old_offsets.append(point);
  new_offsets.append(new_point);

  Map<AttributeIDRef, AttributeKind> gathered_attributes;
  geometry_set.gather_attributes_for_propagation(
      {GEO_COMPONENT_TYPE_CURVE}, GEO_COMPONENT_TYPE_CURVE, false, gathered_attributes);

  /* If point size changed, create new curves. */
  if (point != new_point) {
    Curves *new_curves_id = bke::curves_new_nomain_single(new_point, CURVE_TYPE_POLY);
    bke::CurvesGeometry &new_curves = bke::CurvesGeometry::wrap(new_curves_id->geometry);

    new_curves.resize(new_curves.point_size, new_offsets.size() - 1);
    new_curves.offsets().copy_from(new_offsets);
    new_curves.tag_topology_changed();

    CurveComponent temp_component;
    temp_component.replace(new_curves_id, GeometryOwnershipType::Editable);

    for (const Map<AttributeIDRef, AttributeKind>::Item entry : gathered_attributes.items()) {
      const AttributeIDRef attribute_id = entry.key;
      ReadAttributeLookup src_attribute = component.attribute_try_get_for_read(attribute_id);
      if (!src_attribute) {
        continue;
      }
      AttributeDomain out_domain = src_attribute.domain;
      const CustomDataType data_type = bke::cpp_type_to_custom_data_type(
          src_attribute.varray.type());
      OutputAttribute dst_attribute = temp_component.attribute_try_get_for_output_only(
          attribute_id, out_domain, data_type);
      if (!dst_attribute) {
        continue;
      }
      attribute_math::convert_to_static_type(data_type, [&](auto dummy) {
        using T = decltype(dummy);
        VArray_Span<T> src{src_attribute.varray.typed<T>()};
        MutableSpan<T> dst = dst_attribute.as_span<T>();
        switch (out_domain) {
          case ATTR_DOMAIN_CURVE:
            for (const int i : IndexRange(dst.size())) {
              dst[i] = src[old_curve[i]];
            }
            break;
          case ATTR_DOMAIN_POINT:
            for (const int i : IndexRange(new_offsets.size() - 1)) {
              const int copy_size = new_offsets[i + 1] - new_offsets[i];
              dst.slice(new_offsets[i], copy_size).copy_from(src.slice(old_offsets[i], copy_size));
            }
            break;
          default:
            break;
        }
      });
      dst_attribute.save();
    }
    geometry_set.replace_curves(new_curves_id);
  }
  else {
    /* Points Stayed the same, only copy data for curve domain. */
    CurveComponent &temp_component = geometry_set.get_component_for_write<CurveComponent>();
    Curves *curves_id_new = temp_component.get_for_write();
    bke::CurvesGeometry &curves_new = bke::CurvesGeometry::wrap(curves_id_new->geometry);

    curves_new.resize(new_point, new_offsets.size() - 1);
    curves_new.offsets().copy_from(new_offsets);
    curves_new.tag_topology_changed();

    for (const Map<AttributeIDRef, AttributeKind>::Item entry : gathered_attributes.items()) {
      const AttributeIDRef attribute_id = entry.key;
      ReadAttributeLookup src_attribute = component.attribute_try_get_for_read(attribute_id);
      if (!src_attribute || src_attribute.domain != ATTR_DOMAIN_CURVE) {
        continue;
      }
      AttributeDomain out_domain = src_attribute.domain;
      const CustomDataType data_type = bke::cpp_type_to_custom_data_type(
          src_attribute.varray.type());
      OutputAttribute dst_attribute = temp_component.attribute_try_get_for_output_only(
          attribute_id, out_domain, data_type);
      if (!dst_attribute) {
        continue;
      }
      attribute_math::convert_to_static_type(data_type, [&](auto dummy) {
        using T = decltype(dummy);
        VArray_Span<T> src{src_attribute.varray.typed<T>()};
        MutableSpan<T> dst = dst_attribute.as_span<T>();
        for (const int i : IndexRange(dst.size())) {
          dst[i] = src[old_curve[i]];
        }
      });
      dst_attribute.save();
    }
  }
}

static void node_geo_exec(GeoNodeExecParams params)
{
  GeometrySet geometry_set = params.extract_input<GeometrySet>("Curves");
  Field<bool> selection = params.extract_input<Field<bool>>("Selection");
  Field<bool> previous = params.extract_input<Field<bool>>("Remove Previous");
  Field<bool> next = params.extract_input<Field<bool>>("Remove Next");

  geometry_set.modify_geometry_sets([&](GeometrySet &geometry_set) {
    if (geometry_set.has(GeometryComponentType::GEO_COMPONENT_TYPE_CURVE)) {
      split_curves(geometry_set, selection, previous, next);
    }
  });

  params.set_output("Curves", geometry_set);
}
}  // namespace blender::nodes::node_geo_curve_split_curves_cc

void register_node_type_geo_curve_split_curves()
{
  namespace file_ns = blender::nodes::node_geo_curve_split_curves_cc;
  static bNodeType ntype;
  geo_node_type_base(&ntype, GEO_NODE_CURVE_SPLIT_CURVES, "Split Curves", NODE_CLASS_GEOMETRY);
  ntype.declare = file_ns::node_declare;
  ntype.geometry_node_execute = file_ns::node_geo_exec;
  nodeRegisterType(&ntype);
}