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

#include "BLI_assert.h"
#include "BLI_listbase.h"
#include "BLI_math_vector_types.hh"
#include "BLI_memory_utils.hh"
#include "BLI_set.hh"
#include "BLI_string.h"
#include "BLI_string_ref.hh"

#include "DNA_layer_types.h"
#include "DNA_node_types.h"
#include "DNA_scene_types.h"
#include "DNA_space_types.h"

#include "BKE_compositor.hh"
#include "BKE_context.hh"
#include "BKE_image.hh"
#include "BKE_lib_id.hh"
#include "BKE_node.hh"
#include "BKE_node_runtime.hh"
#include "BKE_scene.hh"

#include "RE_engine.h"

#include "RNA_access.hh"

#include "UI_interface.hh"
#include "UI_interface_layout.hh"
#include "UI_resources.hh"

#include "NOD_node_extra_info.hh"

#include "COM_algorithm_extract_alpha.hh"
#include "COM_node_operation.hh"
#include "COM_utilities.hh"

#include "node_composite_util.hh"

namespace blender::nodes::node_composite_render_layer_cc {

static void node_init(const bContext *context, PointerRNA *node_pointer)
{
  Scene *scene = CTX_data_scene(context);
  bNode *node = node_pointer->data_as<bNode>();
  node->flag |= NODE_PREVIEW;

  node->id = &scene->id;
  id_us_plus(node->id);
}

/* Default declaration for contextless static declarations, when no scene is assigned, or when the
 * engine has no extra passes. */
static void declare_default(NodeDeclarationBuilder &b)
{
  b.add_output<decl::Color>("Image"_ustr).structure_type(StructureType::Dynamic);
  b.add_output<decl::Float>("Alpha"_ustr).structure_type(StructureType::Dynamic);
}

/* Declares an already existing output. */
static BaseSocketDeclarationBuilder &declare_existing_output(NodeDeclarationBuilder &b,
                                                             const bNodeSocket *output)
{
  if (output->type == SOCK_VECTOR) {
    const int dimensions = output->default_value_typed<bNodeSocketValueVector>()->dimensions;
    return b.add_output<decl::Vector>(output->identifier_ustr())
        .dimensions(dimensions)
        .structure_type(StructureType::Dynamic);
  }
  return b.add_output(output->type, output->identifier_ustr())
      .structure_type(StructureType::Dynamic);
}

/* Declares the already existing outputs. This is done in cases where the scene references an
 * engine that is not registered or a view layer that does not exist. Which gives the user the
 * opportunity to register the engine or update the view layer while maintaining sockets and out
 * going links. */
static void declare_existing(NodeDeclarationBuilder &b)
{
  const bNode *node = b.node_or_null();
  for (const bNodeSocket &output : node->outputs) {
    declare_existing_output(b, &output);
  }
}

/* Declares an output that matches the type of the given pass. */
static void declare_pass_callback(void *user_data,
                                  Scene * /*scene*/,
                                  ViewLayer * /*view_layer*/,
                                  const char *pass_name,
                                  int channels_count,
                                  const char * /*channel_id*/,
                                  eNodeSocketDatatype socket_type)
{
  NodeDeclarationBuilder &b = *static_cast<NodeDeclarationBuilder *>(user_data);

  /* The combined pass is aliased as Image. */
  const UString name = StringRef(pass_name) == RE_PASSNAME_COMBINED ? "Image"_ustr :
                                                                      UString(pass_name);
  if (socket_type == SOCK_VECTOR) {
    b.add_output<decl::Vector>(UString(name))
        .dimensions(channels_count)
        .structure_type(StructureType::Dynamic);
  }
  else {
    b.add_output(socket_type, name).structure_type(StructureType::Dynamic);
  }

  /* The Alpha pass is generated based on the combined pass. */
  if (StringRef(pass_name) == RE_PASSNAME_COMBINED) {
    b.add_output<decl::Float>("Alpha"_ustr).structure_type(StructureType::Dynamic);
  }
}

static void declare_extra_passes(NodeDeclarationBuilder &b,
                                 const Scene *scene,
                                 const ViewLayer *view_layer)
{
  if ((scene->r.mode & R_EDGE_FRS) &&
      (view_layer->freestyle_config.flags & FREESTYLE_AS_RENDER_PASS))
  {
    b.add_output<decl::Color>(RE_PASSNAME_FREESTYLE ""_ustr)
        .structure_type(StructureType::Dynamic);
  }

  if (view_layer->grease_pencil_flags & GREASE_PENCIL_AS_SEPARATE_PASS) {
    b.add_output<decl::Color>(RE_PASSNAME_GREASE_PENCIL ""_ustr)
        .structure_type(StructureType::Dynamic);
  }
}

/* Declares outputs that are linked and existed in the previous state of the node but no longer
 * exist in the new state. The outputs are set as unavailable, so they are not accessible to the
 * user. This is useful to retain links if the user changed the render engine and thus the passes
 * changed. */
static void declare_old_linked_outputs(NodeDeclarationBuilder &b)
{
  Set<UString> added_outputs_identifiers;
  for (const SocketDeclaration *output_declaration : b.declaration().sockets(SOCK_OUT)) {
    added_outputs_identifiers.add_new(output_declaration->identifier);
  }
  const bNodeTree *node_tree = b.tree_or_null();
  const bNode *node = b.node_or_null();
  node_tree->ensure_topology_cache();
  for (const bNodeSocket *output : node->output_sockets()) {
    if (added_outputs_identifiers.contains(output->identifier_ustr())) {
      continue;
    }
    if (!output->is_directly_linked()) {
      continue;
    }
    declare_existing_output(b, output).available(false);
  }
}

static void node_declare(NodeDeclarationBuilder &b)
{
  const bNode *node = b.node_or_null();
  if (!node) {
    declare_default(b);
    return;
  }

  const bNodeTree *node_tree = b.tree_or_null();
  if (!node_tree) {
    declare_default(b);
    return;
  }

  BLI_SCOPED_DEFER([&]() { declare_old_linked_outputs(b); });

  Scene *scene = reinterpret_cast<Scene *>(node->id);
  if (!scene) {
    declare_default(b);
    return;
  }

  if (!RE_engines_is_registered(scene->r.engine)) {
    declare_existing(b);
    return;
  }

  RenderEngineType *engine_type = RE_engines_find(scene->r.engine);
  if (!engine_type->update_render_passes) {
    declare_default(b);
    return;
  }

  ViewLayer *view_layer = static_cast<ViewLayer *>(
      BLI_findlink(&scene->view_layers, node->custom1));
  if (!view_layer) {
    declare_existing(b);
    return;
  }

  RenderEngine *engine = RE_engine_create(engine_type);
  RE_engine_update_render_passes(engine, scene, view_layer, declare_pass_callback, &b);
  RE_engine_free(engine);

  declare_extra_passes(b, scene, view_layer);
}

static void node_draw_buttons(ui::Layout &layout, bContext *context, PointerRNA *node_pointer)
{
  template_id(&layout, context, node_pointer, "scene", nullptr, nullptr, nullptr);

  const bNode *node = node_pointer->data_as<bNode>();
  if (!node->id) {
    return;
  }

  ui::Layout &column = layout.column(false);
  ui::Layout &row = column.row(true);
  row.prop(node_pointer, "layer", ui::ITEM_R_SPLIT_EMPTY_NAME, "", ICON_NONE);

  PropertyRNA *layer_property = RNA_struct_find_property(node_pointer, "layer");
  const char *layer_name;
  if (!RNA_property_enum_identifier(context,
                                    node_pointer,
                                    layer_property,
                                    RNA_property_enum_get(node_pointer, layer_property),
                                    &layer_name))
  {
    return;
  }

  char scene_name[MAX_ID_NAME - 2];
  PointerRNA scene_pointer = RNA_pointer_get(node_pointer, "scene");
  RNA_string_get(&scene_pointer, "name", scene_name);

  PointerRNA render_operator = row.op(
      "RENDER_OT_render", "", ICON_RENDER_STILL, wm::OpCallContext::InvokeDefault, UI_ITEM_NONE);
  RNA_string_set(&render_operator, "layer", layer_name);
  RNA_string_set(&render_operator, "scene", scene_name);
}

static void node_extra_info(NodeExtraInfoParams &parameters)
{
  SpaceNode *space_node = CTX_wm_space_node(&parameters.C);
  if (space_node->node_tree_sub_type != SNODE_COMPOSITOR_SCENE) {
    NodeExtraInfoRow row;
    row.text = RPT_("Node Unsupported");
    row.tooltip = TIP_("The Render Layers node is only supported for scene compositing");
    row.icon = ICON_ERROR;
    parameters.rows.append(std::move(row));
  }

  /* EEVEE supports passes. */
  const Scene *scene = CTX_data_scene(&parameters.C);
  if (StringRef(scene->r.engine) == RE_engine_id_BLENDER_EEVEE) {
    return;
  }

  if (!bke::compositor::is_viewport_compositor_used(parameters.C)) {
    return;
  }

  bool is_any_pass_used = false;
  for (const bNodeSocket *output : parameters.node.output_sockets()) {
    /* Combined pass is always available. */
    if (STR_ELEM(output->name, "Image", "Alpha", "Grease Pencil")) {
      continue;
    }
    if (output->is_logically_linked()) {
      is_any_pass_used = true;
      break;
    }
  }

  if (!is_any_pass_used) {
    return;
  }

  NodeExtraInfoRow row;
  row.text = RPT_("Passes Not Supported");
  row.tooltip = TIP_("Render passes in the Viewport compositor are only supported in EEVEE");
  row.icon = ICON_ERROR;
  parameters.rows.append(std::move(row));
}

using namespace blender::compositor;

class RenderLayerOperation : public NodeOperation {
 public:
  using NodeOperation::NodeOperation;

  void execute() override
  {
    const Scene *scene = this->get_scene();
    const int view_layer = this->get_viewer_layer_index();

    for (const bNodeSocket *output : this->node().output_sockets()) {
      if (!is_socket_available(output)) {
        continue;
      }

      Result &result = this->get_result(output->identifier);
      if (!result.should_compute()) {
        continue;
      }

      if (StringRef(output->identifier) == "Alpha") {
        Result combined_pass = this->context().get_pass(scene, view_layer, RE_PASSNAME_COMBINED);
        extract_alpha(this->context(), combined_pass, result);
        combined_pass.release();
        continue;
      }

      Result pass = this->context().get_pass(scene, view_layer, output->identifier);
      result.set_type(pass.type());
      result.set_precision(pass.precision());
      result.share_data(pass);
      pass.release();
    }
  }

  int get_viewer_layer_index()
  {
    return this->node().custom1;
  }

  const Scene *get_scene()
  {
    return reinterpret_cast<const Scene *>(this->node().id);
  }
};

static NodeOperation *get_compositor_operation(Context &context, const bNode &node)
{
  return new RenderLayerOperation(context, node);
}

static void node_register()
{
  static bke::bNodeType ntype;

  cmp_node_type_base(&ntype, "CompositorNodeRLayers"_ustr, CMP_NODE_R_LAYERS);
  ntype.ui_name = "Render Layers";
  ntype.ui_description = "Input render passes from a scene render";
  ntype.enum_name_legacy = "R_LAYERS";
  ntype.nclass = NODE_CLASS_INPUT;
  ntype.flag |= NODE_PREVIEW;
  ntype.initfunc_api = node_init;
  ntype.declare = node_declare;
  ntype.draw_buttons = node_draw_buttons;
  ntype.get_compositor_operation = get_compositor_operation;
  ntype.get_extra_info = node_extra_info;
  ntype.default_width = bke::NodeWidth::_240;

  bke::node_register_type(ntype);
}
NOD_REGISTER_NODE(node_register)

}  // namespace blender::nodes::node_composite_render_layer_cc
