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

#include "AS_asset_library.hh"
#include "AS_asset_representation.hh"

#include "BLI_listbase.h"
#include "BLI_string_utf8.h"

#include "DNA_space_types.h"

#include "BKE_asset.hh"
#include "BKE_context.hh"
#include "BKE_idprop.hh"
#include "BKE_lib_id.hh"
#include "BKE_main_invariants.hh"
#include "BKE_node_runtime.hh"
#include "BKE_node_tree_update.hh"
#include "BKE_screen.hh"

#include "UI_string_search.hh"

#include "NOD_socket.hh"
#include "NOD_socket_search_link.hh"

#include "BLT_translation.hh"

#include "WM_api.hh"

#include "DEG_depsgraph_build.hh"

#include "ED_asset.hh"
#include "ED_node.hh"
#include "ED_screen.hh"

#include "RNA_access.hh"
#include "RNA_define.hh"
#include "RNA_prototypes.hh"

#include "node_intern.hh"

namespace blender {

using nodes::SocketLinkOperation;

namespace ed::space_node {

struct LinkDragSearchStorage {
  bNode &from_node;
  bNodeSocket &from_socket;
  float2 cursor;
  Vector<SocketLinkOperation> search_link_ops;
  char search[256];
  bool update_items_tag = true;

  eNodeSocketInOut in_out() const
  {
    return static_cast<eNodeSocketInOut>(from_socket.in_out);
  }
};

static void link_drag_search_listen_fn(const wmRegionListenerParams *params, void *arg)
{
  LinkDragSearchStorage &storage = *static_cast<LinkDragSearchStorage *>(arg);
  const wmNotifier *wmn = params->notifier;

  switch (wmn->category) {
    case NC_ASSET:
      if (wmn->data == ND_ASSET_LIST_READING) {
        storage.update_items_tag = true;
      }
      break;
  }
}

static void add_reroute_node_fn(nodes::LinkSearchOpParams &params)
{
  bNode &reroute = params.add_node("NodeReroute"_ustr);
  if (params.socket.in_out == SOCK_IN) {
    bke::node_add_link(params.node_tree,
                       reroute,
                       *static_cast<bNodeSocket *>(reroute.outputs.first),
                       params.node,
                       params.socket);
  }
  else {
    bke::node_add_link(params.node_tree,
                       params.node,
                       params.socket,
                       reroute,
                       *static_cast<bNodeSocket *>(reroute.inputs.first));
  }
}

static void add_group_input_node_fn(nodes::LinkSearchOpParams &params)
{
  /* Add a group input based on the connected socket, and add a new group input node. */
  bNodeTreeInterfaceSocket *socket_iface = bke::node_interface::add_interface_socket_from_node(
      params.node_tree, params.node, params.socket, params.socket.name);
  params.node_tree.tree_interface.active_item_set(&socket_iface->item);

  bNode &group_input = params.add_node("NodeGroupInput"_ustr);

  /* This is necessary to create the new sockets in the other input nodes. */
  BKE_main_ensure_invariants(*CTX_data_main(&params.C), params.node_tree.id);

  /* Hide the new input in all other group input nodes, to avoid making them taller. */
  for (bNode *node : params.node_tree.all_nodes()) {
    if (node->is_group_input()) {
      bNodeSocket *new_group_input_socket = bke::node_find_socket(
          *node, SOCK_OUT, UString(socket_iface->identifier));
      if (new_group_input_socket) {
        new_group_input_socket->flag |= SOCK_HIDDEN;
      }
    }
  }

  /* Hide all existing inputs in the new group input node, to only display the new one. */
  for (bNodeSocket &socket : group_input.outputs) {
    socket.flag |= SOCK_HIDDEN;
  }

  bNodeSocket *socket = bke::node_find_socket(
      group_input, SOCK_OUT, UString(socket_iface->identifier));
  if (socket) {
    /* Unhide the socket for the new input in the new node and make a connection to it. */
    socket->flag &= ~SOCK_HIDDEN;
    bke::node_add_link(params.node_tree, group_input, *socket, params.node, params.socket);

    bke::node_socket_move_default_value(
        *CTX_data_main(&params.C), params.node_tree, params.socket, *socket);
  }
}

static void add_existing_group_input_fn(nodes::LinkSearchOpParams &params,
                                        const bNodeTreeInterfaceSocket &interface_socket)
{
  const eNodeSocketInOut in_out = eNodeSocketInOut(params.socket.in_out);
  NodeTreeInterfaceSocketFlag flag = NodeTreeInterfaceSocketFlag(0);
  SET_FLAG_FROM_TEST(flag, in_out & SOCK_IN, NODE_INTERFACE_SOCKET_INPUT);
  SET_FLAG_FROM_TEST(flag, in_out & SOCK_OUT, NODE_INTERFACE_SOCKET_OUTPUT);

  bNode &group_input = params.add_node("NodeGroupInput"_ustr);

  for (bNodeSocket &socket : group_input.outputs) {
    socket.flag |= SOCK_HIDDEN;
  }

  bNodeSocket *socket = bke::node_find_socket(
      group_input, SOCK_OUT, UString(interface_socket.identifier));
  if (socket != nullptr) {
    socket->flag &= ~SOCK_HIDDEN;
    bke::node_add_link(params.node_tree, group_input, *socket, params.node, params.socket);
  }
}

/**
 * \note This could use #search_link_ops_for_socket_templates, but we have to store the inputs and
 * outputs as IDProperties for assets because of asset indexing, so that's all we have without
 * loading the file.
 */
static void search_link_ops_for_asset_metadata(const bNodeTree &node_tree,
                                               const bNodeSocket &socket,
                                               const asset_system::AssetRepresentation &asset,
                                               Vector<SocketLinkOperation> &search_link_ops)
{
  const AssetMetaData &asset_data = asset.get_metadata();
  const StringRef asset_name = asset.get_name();
  const IDProperty *tree_type = BKE_asset_metadata_idprop_find(&asset_data, "type");
  if (tree_type == nullptr || IDP_int_get(tree_type) != node_tree.type) {
    return;
  }

  const bke::bNodeTreeType &node_tree_type = *node_tree.typeinfo;
  const eNodeSocketInOut in_out = socket.in_out == SOCK_OUT ? SOCK_IN : SOCK_OUT;

  const IDProperty *properties = BKE_asset_metadata_idprop_find(&asset_data, "properties");
  if (!properties || properties->type != IDP_GROUP) {
    return;
  }
  const IDProperty *sockets = IDP_GetPropertyFromGroup(
      properties, in_out == SOCK_IN ? "inputs" : "output_sockets");
  if (!sockets || sockets->type != IDP_GROUP) {
    return;
  }

  int weight = -1;
  Set<StringRef> socket_names;
  for (IDProperty &socket_property : sockets->data.group) {
    if (socket_property.type != IDP_GROUP) {
      continue;
    }
    const std::optional<int> type = IDP_group_lookup_int(socket_property, "type");
    if (!type) {
      continue;
    }
    eNodeSocketDatatype from = socket.type;
    eNodeSocketDatatype to = eNodeSocketDatatype(*type);
    if (socket.in_out == SOCK_OUT) {
      std::swap(from, to);
    }
    if (node_tree_type.validate_link && !node_tree_type.validate_link(from, to)) {
      continue;
    }
    const StringRefNull identifier = socket_property.name;
    const StringRefNull name =
        IDP_group_lookup_string(socket_property, "name").value_or(identifier);
    if (!socket_names.add(name)) {
      /* See comment in #search_link_ops_for_declarations. */
      continue;
    }

    search_link_ops.append(
        {asset_name + " " + UI_MENU_ARROW_SEP + name,
         [&asset, name, in_out](nodes::LinkSearchOpParams &params) {
           Main &bmain = *CTX_data_main(&params.C);

           bNodeTree *group = reinterpret_cast<bNodeTree *>(
               asset::asset_local_id_ensure_imported(bmain, asset));
           if (!group) {
             return;
           }
           bNode &node = params.add_node(params.node_tree.typeinfo->group_idname);
           node.id = &group->id;
           id_us_plus(node.id);
           BKE_ntree_update_tag_node_property(&params.node_tree, &node);
           DEG_relations_tag_update(&bmain);

           node.flag &= ~NODE_OPTIONS;
           node.width = group->default_group_node_width;

           /* Create the inputs and outputs on the new node. */
           nodes::update_node_declaration_and_sockets(params.node_tree, node);

           bNodeSocket *new_node_socket = bke::node_find_enabled_socket(node, in_out, name);
           if (new_node_socket != nullptr) {
             /* Rely on the way #node_add_link switches in/out if necessary. */
             bke::node_add_link(
                 params.node_tree, params.node, params.socket, node, *new_node_socket);
           }
         },
         weight});

    weight--;
  }
}

static void gather_search_link_ops_for_all_assets(const bContext &C,
                                                  const bNodeTree &node_tree,
                                                  const bNodeSocket &socket,
                                                  Vector<SocketLinkOperation> &search_link_ops)
{
  const AssetLibraryReference library_ref = asset_system::all_library_reference();
  asset::AssetFilterSettings filter_settings{};
  filter_settings.id_types = FILTER_ID_NT;

  asset::list::storage_fetch(&library_ref, &C);
  asset::list::iterate(library_ref, [&](asset_system::AssetRepresentation &asset) {
    if (!asset::filter_matches_asset(&filter_settings, asset)) {
      return true;
    }
    search_link_ops_for_asset_metadata(node_tree, socket, asset, search_link_ops);
    return true;
  });
}

/**
 * Call the callback to gather compatible socket connections for all node types, and the operations
 * that will actually make the connections. Also add some custom operations like connecting a group
 * output node.
 */
static void gather_socket_link_operations(const bContext &C,
                                          bNodeTree &node_tree,
                                          const bNodeSocket &socket,
                                          Vector<SocketLinkOperation> &search_link_ops)
{
  const SpaceNode &snode = *CTX_wm_space_node(&C);
  for (const bke::bNodeType *node_type : bke::node_types_get()) {
    const char *disabled_hint;
    if (node_type->poll && !node_type->poll(node_type, &node_tree, &disabled_hint)) {
      continue;
    }
    if (node_type->add_ui_poll && !node_type->add_ui_poll(&C)) {
      continue;
    }
    if (StringRefNull(node_type->ui_name).endswith("(Legacy)")) {
      continue;
    }
    if (node_type->gather_link_search_ops) {
      nodes::GatherLinkSearchOpParams params{
          *node_type, snode, node_tree, socket, search_link_ops};
      node_type->gather_link_search_ops(params);
    }
  }

  search_link_ops.append({IFACE_("Reroute"), add_reroute_node_fn});

  const bool is_node_group = !(node_tree.id.flag & ID_FLAG_EMBEDDED_DATA);

  if (is_node_group && socket.in_out == SOCK_IN) {
    search_link_ops.append({IFACE_("Group Input"), add_group_input_node_fn});

    int weight = -1;
    node_tree.tree_interface.foreach_item([&](const bNodeTreeInterfaceItem &item) {
      if (item.item_type != NodeTreeInterfaceItemType::Socket) {
        return true;
      }
      const bNodeTreeInterfaceSocket &interface_socket =
          reinterpret_cast<const bNodeTreeInterfaceSocket &>(item);
      if (!(interface_socket.flag & NODE_INTERFACE_SOCKET_INPUT)) {
        return true;
      }
      {
        const bke::bNodeSocketType *from_typeinfo = bke::node_socket_type_find(
            interface_socket.socket_type);
        const eNodeSocketDatatype from = from_typeinfo ? from_typeinfo->type : SOCK_CUSTOM;
        const eNodeSocketDatatype to = socket.typeinfo->type;
        if (node_tree.typeinfo->validate_link && !node_tree.typeinfo->validate_link(from, to)) {
          return true;
        }
      }
      search_link_ops.append({std::string(IFACE_("Group Input")) + " " + UI_MENU_ARROW_SEP +
                                  (interface_socket.name ? interface_socket.name : ""),
                              [interface_socket](nodes::LinkSearchOpParams &params) {
                                add_existing_group_input_fn(params, interface_socket);
                              },
                              weight});
      weight--;
      return true;
    });
  }

  gather_search_link_ops_for_all_assets(C, node_tree, socket, search_link_ops);
}

static void link_drag_search_update_fn(
    const bContext *C, void *arg, const char *str, ui::SearchItems *items, const bool is_first)
{
  LinkDragSearchStorage &storage = *static_cast<LinkDragSearchStorage *>(arg);
  if (storage.update_items_tag) {
    bNodeTree *node_tree = CTX_wm_space_node(C)->edittree;
    storage.search_link_ops.clear();
    gather_socket_link_operations(*C, *node_tree, storage.from_socket, storage.search_link_ops);
    storage.update_items_tag = false;
  }

  ui::string_search::StringSearch<SocketLinkOperation> search{
      string_search::MainWordsHeuristic::All};

  for (SocketLinkOperation &op : storage.search_link_ops) {
    search.add(op.name, &op, op.weight);
  }

  /* Don't filter when the menu is first opened, but still run the search
   * so the items are in the same order they will appear in while searching. */
  const char *string = is_first ? "" : str;
  const Vector<SocketLinkOperation *> filtered_items = search.query(string);

  for (SocketLinkOperation *item : filtered_items) {
    if (!search_item_add(items, item->name, item, ICON_NONE, 0, 0)) {
      break;
    }
  }
}

static bNode *get_new_linked_node(bNodeSocket &socket, const Span<bNode *> new_nodes)
{
  for (const bNodeLink *link : socket.directly_linked_links()) {
    if (new_nodes.contains(link->fromnode)) {
      return link->fromnode;
    }
    if (new_nodes.contains(link->tonode)) {
      return link->tonode;
    }
  }
  return nullptr;
}

static void link_drag_search_exec_fn(bContext *C, void *arg1, void *arg2)
{
  Main &bmain = *CTX_data_main(C);
  SpaceNode &snode = *CTX_wm_space_node(C);
  bNodeTree &node_tree = *snode.edittree;
  LinkDragSearchStorage &storage = *static_cast<LinkDragSearchStorage *>(arg1);
  SocketLinkOperation *item = static_cast<SocketLinkOperation *>(arg2);
  if (item == nullptr) {
    return;
  }

  node_deselect_all(node_tree);

  Vector<bNode *> new_nodes;
  nodes::LinkSearchOpParams params{
      *C, node_tree, storage.from_node, storage.from_socket, new_nodes};
  item->fn(params);
  if (new_nodes.is_empty()) {
    return;
  }

  /* Used to position the new nodes where the cursor is. */
  const float2 cursor_offset = (storage.cursor / UI_SCALE_FAC) + float2(0.0f, 20.0f);

  /* Used to position the new nodes so that the newly linked socket is aligned to the cursor. */
  float2 link_offset{};
  node_tree.ensure_topology_cache();
  if (bNode *new_directly_linked_node = get_new_linked_node(storage.from_socket, new_nodes)) {
    link_offset -= new_directly_linked_node->location;
    if (storage.in_out() == SOCK_IN) {
      link_offset.x -= new_directly_linked_node->width;
    }
  }

  const float2 offset_in_tree = cursor_offset + link_offset;
  for (bNode *new_node : new_nodes) {
    /* The node may have an initial offset already, so use +=. */
    new_node->location[0] += offset_in_tree.x;
    new_node->location[1] += offset_in_tree.y;
    bke::node_set_selected(*new_node, true);
  }
  bke::node_set_active(node_tree, *new_nodes[0]);

  /* Ideally it would be possible to tag the node tree in some way so it updates only after the
   * translate operation is finished, but normally moving nodes around doesn't cause updates. */
  BKE_main_ensure_invariants(bmain, node_tree.id);

  /* Start translation operator with the new node. */
  wmOperatorType *ot = WM_operatortype_find("NODE_OT_translate_attach_remove_on_cancel", true);
  BLI_assert(ot);
  PointerRNA ptr = WM_operator_properties_create_ptr(ot);
  WM_operator_name_call_ptr(C, ot, wm::OpCallContext::InvokeDefault, &ptr, nullptr);
  WM_operator_properties_free(&ptr);
}

static void link_drag_search_free_fn(void *arg)
{
  LinkDragSearchStorage *storage = static_cast<LinkDragSearchStorage *>(arg);
  delete storage;
}

static ui::Block *create_search_popup_block(bContext *C, ARegion *region, void *arg_op)
{
  LinkDragSearchStorage &storage = *static_cast<LinkDragSearchStorage *>(arg_op);

  ui::Block *block = block_begin(C, region, "_popup", ui::EmbossType::Emboss);
  block_flag_enable(block, ui::BLOCK_LOOP | ui::BLOCK_MOVEMOUSE_QUIT | ui::BLOCK_SEARCH_MENU);
  block_theme_style_set(block, ui::BLOCK_THEME_STYLE_POPUP);

  ui::Button *but = uiDefSearchBut(block,
                                   storage.search,
                                   ICON_VIEWZOOM,
                                   sizeof(storage.search),
                                   storage.in_out() == SOCK_OUT ? 10 : 10 - ui::searchbox_size_x(),
                                   0,
                                   ui::searchbox_size_x(),
                                   UI_UNIT_Y,
                                   "");
  button_func_search_set_sep_string(but, UI_MENU_ARROW_SEP);
  button_func_search_set_listen(but, link_drag_search_listen_fn);
  button_func_search_set(but,
                         nullptr,
                         link_drag_search_update_fn,
                         &storage,
                         false,
                         link_drag_search_free_fn,
                         link_drag_search_exec_fn,
                         nullptr);
  button_flag_enable(but, ui::BUT_ACTIVATE_ON_INIT);

  /* Fake button to hold space for the search items. */
  uiDefBut(block,
           ui::ButtonType::Label,
           "",
           storage.in_out() == SOCK_OUT ? 10 : 10 - ui::searchbox_size_x(),
           10 - ui::searchbox_size_y(),
           ui::searchbox_size_x(),
           ui::searchbox_size_y(),
           nullptr,
           0,
           0,
           std::nullopt);

  const int2 offset = {0, -UI_UNIT_Y};
  block_bounds_set_popup(block, 0.3f * U.widget_unit, offset);
  return block;
}

void invoke_node_link_drag_add_menu(bContext &C,
                                    bNode &node,
                                    bNodeSocket &socket,
                                    const float2 &cursor)
{
  LinkDragSearchStorage *storage = new LinkDragSearchStorage{node, socket, cursor};
  /* Use the "_ex" variant with `can_refresh` false to avoid a double free when closing Blender. */
  popup_block_invoke_ex(&C, create_search_popup_block, storage, nullptr, false);
}

static bool link_drag_operation_test_poll(bContext *C)
{
  if (!ED_operator_node_editable(C)) {
    return false;
  }
  PointerRNA socket_ptr = CTX_data_pointer_get_type(C, "socket", RNA_NodeSocket);
  if (!socket_ptr.data) {
    return false;
  }
  return true;
}

static wmOperatorStatus link_drag_operation_test_exec(bContext *C, wmOperator *op)
{
  SpaceNode &snode = *CTX_wm_space_node(C);
  if (!snode.edittree) {
    return OPERATOR_CANCELLED;
  }
  bNodeTree &ntree = *snode.edittree;
  PointerRNA socket_ptr = CTX_data_pointer_get_type(C, "socket", RNA_NodeSocket);
  if (!socket_ptr.data) {
    return OPERATOR_CANCELLED;
  }
  bNodeSocket &socket = *socket_ptr.data_as<bNodeSocket>();

  Vector<SocketLinkOperation> search_link_ops;
  gather_socket_link_operations(*C, ntree, socket, search_link_ops);

  if (RNA_boolean_get(op->ptr, "find_link_operations")) {
    IDProperty *idprops = IDP_EnsureProperties(&ntree.id);
    IDProperty *idp_link_ops_array = IDP_NewIDPArray("link_operation_names");
    IDP_ResizeIDPArray(idp_link_ops_array, search_link_ops.size());
    for (const int i : search_link_ops.index_range()) {
      const SocketLinkOperation &link_op = search_link_ops[i];
      IDProperty *idp_link_op_name = IDP_NewString(link_op.name.c_str(), "name");
      IDP_SetIndexArray(idp_link_ops_array, i, idp_link_op_name);
      /* IDP_SetIndexArray makes a shallow copy. */
      MEM_delete(idp_link_op_name);
    }
    IDP_ReplaceInGroup(idprops, idp_link_ops_array);
    return OPERATOR_FINISHED;
  }

  const int link_op_index = RNA_int_get(op->ptr, "link_operation_index");
  if (!search_link_ops.index_range().contains(link_op_index)) {
    BKE_report(op->reports, RPT_ERROR_INVALID_INPUT, "Link operation index out of range");
    return OPERATOR_CANCELLED;
  }

  const SocketLinkOperation &link_op = search_link_ops[link_op_index];
  Vector<bNode *> added_nodes;
  nodes::LinkSearchOpParams params{*C, ntree, socket.owner_node(), socket, added_nodes};
  link_op.fn(params);
  // const std::string msg = fmt::format(
  //     "Link operation \"{}\" added {} nodes", link_op.name, added_nodes.size());
  // BKE_report(op->reports, RPT_INFO, msg.c_str());

  /* Select only added nodes. */
  for (bNode &node : ntree.nodes) {
    bke::node_set_selected(node, false);
  }
  for (bNode *node : added_nodes) {
    bke::node_set_selected(*node, true);
  }

  return OPERATOR_FINISHED;
}

void NODE_OT_link_drag_operation_test(wmOperatorType *ot)
{
  ot->name = "Link Drag Operation Test";
  ot->idname = "NODE_OT_link_drag_operation_test";
  ot->description = "Run a node link-drag operation for testing";

  ot->poll = link_drag_operation_test_poll;
  ot->exec = link_drag_operation_test_exec;

  ot->flag = OPTYPE_INTERNAL;

  RNA_def_boolean(ot->srna,
                  "find_link_operations",
                  false,
                  "Find Link Operations",
                  "Write link operation names for the context socket the \"link_operation_names\" "
                  "property of the node tree");
  RNA_def_int(ot->srna,
              "link_operation_index",
              -1,
              -1,
              INT_MAX,
              "Link Operation Index",
              "Link operation to execute on the context socket",
              0,
              INT_MAX);
}

}  // namespace ed::space_node

}  // namespace blender
