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

#include "BLI_devirtualize_parameters.hh"
#include "BLI_length_parameterize.hh"

#include "BKE_attribute_math.hh"
#include "BKE_curves.hh"

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

#include "node_geometry_util.hh"

namespace blender::nodes::node_geo_curve_sample_cc {

NODE_STORAGE_FUNCS(NodeGeometryCurveSample)

static void node_declare(NodeDeclarationBuilder &b)
{
  b.add_input<decl::Geometry>(N_("Curve"))
      .only_realized_data()
      .supported_type(GEO_COMPONENT_TYPE_CURVE);
  b.add_input<decl::Float>(N_("Factor"))
      .min(0.0f)
      .max(1.0f)
      .subtype(PROP_FACTOR)
      .supports_field()
      .make_available([](bNode &node) { node_storage(node).mode = GEO_NODE_CURVE_SAMPLE_FACTOR; });
  b.add_input<decl::Float>(N_("Length"))
      .min(0.0f)
      .subtype(PROP_DISTANCE)
      .supports_field()
      .make_available([](bNode &node) { node_storage(node).mode = GEO_NODE_CURVE_SAMPLE_LENGTH; });
  
  b.add_input<decl::Int>(N_("Spline")).supports_field().min(0);

  b.add_input<decl::Float>(N_("Value"), "Value_Float").hide_value().supports_field();
  b.add_input<decl::Int>(N_("Value"), "Value_Int").hide_value().supports_field();
  b.add_input<decl::Vector>(N_("Value"), "Value_Vector").hide_value().supports_field();
  b.add_input<decl::Color>(N_("Value"), "Value_Color").hide_value().supports_field();
  b.add_input<decl::Bool>(N_("Value"), "Value_Bool").hide_value().supports_field();

  b.add_output<decl::Float>(N_("Value"), "Value_Float").field_source();
  b.add_output<decl::Int>(N_("Value"), "Value_Int").field_source();
  b.add_output<decl::Vector>(N_("Value"), "Value_Vector").field_source();
  b.add_output<decl::Color>(N_("Value"), "Value_Color").field_source();
  b.add_output<decl::Bool>(N_("Value"), "Value_Bool").field_source();
}

static void node_layout(uiLayout *layout, bContext *UNUSED(C), PointerRNA *ptr)
{
  uiItemR(layout, ptr, "data_type", 0, "", ICON_NONE);
  uiItemR(layout, ptr, "mode", UI_ITEM_R_EXPAND, 0, ICON_NONE);
}

static void node_type_init(bNodeTree *UNUSED(tree), bNode *node)
{
  NodeGeometryCurveSample *data = MEM_cnew<NodeGeometryCurveSample>(__func__);
  data->data_type = CD_PROP_FLOAT3;
  data->mode = GEO_NODE_CURVE_SAMPLE_LENGTH;
  node->storage = data;
}

static void node_update(bNodeTree *ntree, bNode *node)
{
  const NodeGeometryCurveSample &storage = node_storage(*node);
  const GeometryNodeCurveSampleMode mode = (GeometryNodeCurveSampleMode)storage.mode;
  const eCustomDataType data_type = (eCustomDataType)storage.data_type;

  bNodeSocket *sock_factor = ((bNodeSocket *)node->inputs.first)->next;
  bNodeSocket *sock_length = sock_factor->next;

  nodeSetSocketAvailability(ntree, sock_factor, mode == GEO_NODE_CURVE_SAMPLE_FACTOR);
  nodeSetSocketAvailability(ntree, sock_length, mode == GEO_NODE_CURVE_SAMPLE_LENGTH);
  
  bNodeSocket *sock_in_float = sock_length->next->next;
  bNodeSocket *sock_in_int = sock_in_float->next;
  bNodeSocket *sock_in_vector = sock_in_int->next;
  bNodeSocket *sock_in_color = sock_in_vector->next;
  bNodeSocket *sock_in_bool = sock_in_color->next;

  bNodeSocket *sock_out_float = static_cast<bNodeSocket *>(node->outputs.first);
  bNodeSocket *sock_out_int = sock_out_float->next;
  bNodeSocket *sock_out_vector = sock_out_int->next;
  bNodeSocket *sock_out_color = sock_out_vector->next;
  bNodeSocket *sock_out_bool = sock_out_color->next;

  nodeSetSocketAvailability(ntree, sock_in_float, data_type == CD_PROP_FLOAT);
  nodeSetSocketAvailability(ntree, sock_in_int, data_type == CD_PROP_INT32);
  nodeSetSocketAvailability(ntree, sock_in_vector, data_type == CD_PROP_FLOAT3);
  nodeSetSocketAvailability(ntree, sock_in_color, data_type == CD_PROP_COLOR);
  nodeSetSocketAvailability(ntree, sock_in_bool, data_type == CD_PROP_BOOL);

  nodeSetSocketAvailability(ntree, sock_out_float, data_type == CD_PROP_FLOAT);
  nodeSetSocketAvailability(ntree, sock_out_int, data_type == CD_PROP_INT32);
  nodeSetSocketAvailability(ntree, sock_out_vector, data_type == CD_PROP_FLOAT3);
  nodeSetSocketAvailability(ntree, sock_out_color, data_type == CD_PROP_COLOR);
  nodeSetSocketAvailability(ntree, sock_out_bool, data_type == CD_PROP_BOOL);
}

class SampleCurveFieldContext : public FieldContext {
 private:
  const int sampling_count_;
 
  Array<int> sample_curve_indices_;
  Array<float> sample_lengths_;
  
  Span<float3> evaluated_positions_;
  Span<float3> evaluated_tangents_;
  Span<float3> evaluated_normals_;
  
  Array<IndexRange> curves_;
  Array<Span<float>> curves_lengths_;
  VArray<bool> cyclic_;

 public:
  
  SampleCurveFieldContext(const CurveComponent &component, Array<int> sample_curve_indices, Array<float> sample_lengths) :
      sampling_count_(sample_curve_indices.size()), sample_curve_indices_(sample_curve_indices), sample_lengths_(sample_lengths)
  {
    const Curves &curves_id = *component.get_for_read();
    const bke::CurvesGeometry &curves = bke::CurvesGeometry::wrap(curves_id.geometry);
    
    evaluated_positions_ = curves.evaluated_positions();
    evaluated_tangents_ = curves.evaluated_tangents();
    evaluated_normals_ = curves.evaluated_normals();
    
    BLI_assert(sampling_count_ == sample_lengths_.size());
    
    curves_.reinitialize(curves.curves_num());
    curves_lengths_.reinitialize(curves.curves_num());
    cyclic_ = curves.cyclic();
    
    curves.ensure_evaluated_lengths();
    for (const int64_t index : curves_lengths_.index_range()) {
      curves_lengths_[index] = curves.evaluated_lengths_for_curve(index, cyclic_[index]);
      curves_[index] = curves.evaluated_points_for_curve(index);
    }
  }
 
  int64_t points_num() const
  {
    return sampling_count_;
  }

  GVArray get_varray_for_input(const FieldInput &field_input,
                               const IndexMask mask,
                               ResourceScope &UNUSED(scope)) const
  {
    if (dynamic_cast<const fn::IndexFieldInput *>(&field_input)) {
      return fn::IndexFieldInput::get_index_varray(mask);
    }
    
    const bke::AttributeFieldInput *attribute_field_input = dynamic_cast<const bke::AttributeFieldInput *>(&field_input);
    if (attribute_field_input) {
      if (attribute_field_input->attribute_name() == "position") {
        Array<float3> positions(sampling_count_);
        this->sample_date(mask, GSpan(evaluated_positions_), GMutableSpan(positions.as_mutable_span()));
        return VArray<float3>::ForContainer(std::move(positions));
      }
    }
    
    if (dynamic_cast<const bke::NormalFieldInput *>(&field_input)) {
        Array<float3> normals(sampling_count_);
        this->sample_date(mask, GSpan(evaluated_tangents_), GMutableSpan(normals.as_mutable_span()));
        printf("Semple Tangents\n");
        return VArray<float3>::ForContainer(std::move(normals));
    }
    /*
    if (dynamic_cast<const blender::nodes::node_geo_input_tangent_cc::TangentFieldInput *>(&field_input)) {
        Array<float3> tangents(sampling_count_);
        sample_date(mask, GSpan(evaluated_normals_), GMutableSpan(tangents.as_mutable_span()));
        printf("Semple Normals\n");
        return VArray<float3>::ForContainer(std::move(tangents));
    }
    */
    return {};
  }
  
 private:
  void sample_date(const IndexMask mask, GSpan inputs, GMutableSpan outputs) const {
    MultiValueMap<int, int64_t> indices_per_curve;
    for (const int64_t i : mask) {
      indices_per_curve.add(sample_curve_indices_[i], i);
    }
    
    Array<int> indices;
    Array<float> factors;
    
    for (const int curve_i : indices_per_curve.keys()) {
      const IndexMask mask(indices_per_curve.lookup(curve_i));
      indices.reinitialize(mask.size());
      factors.reinitialize(mask.size());

      mask.to_best_mask_type([&](const auto mask) {
        for (const int64_t i : IndexRange(mask.size())) {
          length_parameterize::sample_at_length(
              curves_lengths_[curve_i],
              sample_lengths_[mask[i]],
              indices[i],
              factors[i]);
        }
      });

      attribute_math::convert_to_static_type(inputs.type(), [&](auto dummy) {
        using T = decltype(dummy);
        Span<T> inputs_((T *)(inputs.data()), inputs.size());
        length_parameterize::interpolate_to_masked<T>(
            inputs_.slice(curves_[curve_i]), indices, factors, mask, outputs.typed<T>());
      });
    }
  }
};

class SampleCurveFieldInput : public GeometryFieldInput {
 private:
  GeometrySet geometry_set_;
  
  GField input_field_;
  Field<int> input_curves_indices_field_;
  Field<float> input_distances_field_;

 public:
  SampleCurveFieldInput(GeometrySet geometry_set, GField input_field, Field<int> input_curves_indices_field, Field<float> input_distances_field) :
        GeometryFieldInput(input_field.cpp_type(), "Sample Curve"),
        geometry_set_(std::move(geometry_set)),
        input_field_(std::move(input_field)),
        input_curves_indices_field_(std::move(input_curves_indices_field)),
        input_distances_field_(std::move(input_distances_field))
  {
  }

  GVArray get_varray_for_context(const GeometryComponent &component,
                                 const eAttrDomain domain,
                                 IndexMask mask) const final
  {
    if (!geometry_set_.has_curves()) {
      return {};
    }
    const Curves &curves_id = *geometry_set_.get_curves_for_read();
    const bke::CurvesGeometry &curves = bke::CurvesGeometry::wrap(curves_id.geometry);
    if (curves.points_num() == 0) {
      return {};
    }
    
    const int64_t domain_size = component.attribute_domain_size(domain);
    
    Array<int> indices(domain_size);
    Array<float> lengths(domain_size);
    
    const GeometryComponentFieldContext field_context{component, domain};
    FieldEvaluator evaluator(field_context, domain_size);
    evaluator.add_with_destination(input_curves_indices_field_, indices.as_mutable_span());
    evaluator.add_with_destination(input_distances_field_, lengths.as_mutable_span());
    evaluator.evaluate();
    
    SampleCurveFieldContext sample_component(*geometry_set_.get_component_for_read<CurveComponent>(),
                                             std::move(indices), std::move(lengths));
    GArray output(input_field_.cpp_type(), domain_size);
    FieldEvaluator sampler(sample_component, domain_size);
    sampler.add_with_destination(input_field_, output.as_mutable_span());
    sampler.evaluate();
    return GVArray::ForGArray(std::move(output));
  }
};

static StringRefNull identifier_suffix(eCustomDataType data_type)
{
  switch (data_type) {
    case CD_PROP_BOOL:
      return "Bool";
    case CD_PROP_FLOAT:
      return "Float";
    case CD_PROP_INT32:
      return "Int";
    case CD_PROP_COLOR:
      return "Color";
    case CD_PROP_FLOAT3:
      return "Vector";
    default:
      BLI_assert_unreachable();
      return "";
  }
}

static void node_geo_exec(GeoNodeExecParams params)
{
  GeometrySet geometry_set = params.extract_input<GeometrySet>("Curve");

  const NodeGeometryCurveSample &storage = node_storage(params.node());
  const GeometryNodeCurveSampleMode mode = (GeometryNodeCurveSampleMode)storage.mode;
  const eCustomDataType data_type = (eCustomDataType)storage.data_type;
  
  Field<int> curves_field = params.extract_input<Field<int>>("Spline");
  Field<float> distance_field = params.extract_input<Field<float>>(mode == GEO_NODE_CURVE_SAMPLE_FACTOR ? "Factor" : "Length");

  attribute_math::convert_to_static_type(data_type, [&](auto dummy) {
    using T = decltype(dummy);
    static const std::string identifier = "Value_" + identifier_suffix(data_type);
    Field<T> value_field = params.extract_input<Field<T>>(identifier);
    Field<T> output_field{std::make_shared<SampleCurveFieldInput>(std::move(geometry_set),
                                                                  std::move(value_field),
                                                                  std::move(curves_field),
                                                                  std::move(distance_field))};
    params.set_output(identifier, std::move(output_field));
  });
}

}  // namespace blender::nodes::node_geo_curve_sample_cc

void register_node_type_geo_curve_sample()
{
  namespace file_ns = blender::nodes::node_geo_curve_sample_cc;

  static bNodeType ntype;

  geo_node_type_base(&ntype, GEO_NODE_SAMPLE_CURVE, "Sample Curve", NODE_CLASS_GEOMETRY);
  ntype.geometry_node_execute = file_ns::node_geo_exec;
  ntype.declare = file_ns::node_declare;
  node_type_init(&ntype, file_ns::node_type_init);
  node_type_update(&ntype, file_ns::node_update);
  node_type_storage(
      &ntype, "NodeGeometryCurveSample", node_free_standard_storage, node_copy_standard_storage);
  ntype.draw_buttons = file_ns::node_layout;

  nodeRegisterType(&ntype);
}