#ifndef BRAIN_BT_REGISTRY_HPP_
#define BRAIN_BT_REGISTRY_HPP_

/**
 * @file bt_registry.hpp
 * @brief Behavior Tree registration utilities wrapping BehaviorTree.CPP factory with higher level helpers.
 *
 * Provides:
 *  - BTActionSpec / BTSequenceSpec: lightweight declarative descriptions of action nodes and sequences.
 *  - BTActionHandlers / RegisteredActionNode: injection of user supplied callbacks for start / running / halted phases.
 *  - BTRegistry: central registry to declare custom action node builders and sequence archetypes, and to build trees.
 *
 * Design goals:
 *  - Minimize boilerplate when adding custom stateful BT action nodes whose behavior is defined by lambdas.
 *  - Allow dynamic composition of sequences (converted to XML text consumed by BehaviorTreeFactory).
 *  - Keep user code decoupled from raw XML authoring while still producing valid BTCPP v4 format.
 *
 * Thread-safety:
 *  - Registration APIs (register_bt_node / register_bt_action / register_bt_sequence) are intended for single-threaded
 *    initialization before any trees are executed. They are NOT internally synchronized.
 *  - build_tree_from_sequence may be called concurrently if the underlying BehaviorTreeFactory is used in that way;
 *    BehaviorTree.CPP factory itself is generally not thread-safe for concurrent mutation, so calls should serialize.
 */

#include <behaviortree_cpp/bt_factory.h>
#include <behaviortree_cpp/action_node.h>
#include <rclcpp/rclcpp.hpp>
#include <functional>
#include <map>
#include <string>
#include <vector>
#include <sstream>
#include <unordered_map>
#include <fstream>

/**
 * @brief Specification for a single Behavior Tree action node used when programmatically composing XML.
 *
 * Each node is described by its type (registered with the factory), an instance name, and an optional
 * map of port name -> value pairs which become XML attributes in the emitted tree definition.
 */
struct BTActionSpec
{
  std::string type;                                      ///< Registered node type ID (XML tag)
  std::string name;                                      ///< Instance name attribute
  std::map<std::string, std::string> ports;              ///< Static input/output port bindings as string literals
};

/**
 * @brief A linear sequence of BTActionSpec items rendered as a <Sequence> subtree.
 */
using BTSequenceSpec = std::vector<BTActionSpec>;

/**
 * @brief Callback bundle implementing the lifecycle hooks of a BehaviorTree.CPP StatefulActionNode.
 *
 * All callbacks are optional. Missing callbacks default to immediate SUCCESS (for start/running)
 * or no-op (for halted).
 */
struct BTActionHandlers
{
  std::function<BT::NodeStatus(rclcpp::Node *, BT::TreeNode &)> on_start;   ///< Invoked once when the node transitions from IDLE to RUNNING
  std::function<BT::NodeStatus(rclcpp::Node *, BT::TreeNode &)> on_running; ///< Invoked each tick while RUNNING
  std::function<void(rclcpp::Node *, BT::TreeNode &)> on_halted;            ///< Invoked if the node is halted by the tree
};

/**
 * @brief Concrete BehaviorTree.CPP stateful action node whose behavior is delegated to BTActionHandlers.
 *
 * This class is registered dynamically by BTRegistry when register_bt_action is invoked. Users do not
 * normally instantiate it directly; instead they provide a handlers bundle associated with a type name.
 */
class RegisteredActionNode : public BT::StatefulActionNode
{
public:
  /**
   * @brief Construct with name, handler bundle, and ROS node context.
   * @param name Instance name inside the tree.
   * @param handlers Lifecycle callback handlers.
   * @param node ROS 2 node pointer passed back to callbacks.
   */
  RegisteredActionNode(
    const std::string & name, const BTActionHandlers & handlers,
    rclcpp::Node * node)
  : BT::StatefulActionNode(name, BT::NodeConfiguration()), name_(name), handlers_(handlers), node_(
      node) {}

  /**
   * @brief Declare static ports. Currently none (empty list) simplifying dynamic XML generation.
   */
  static BT::PortsList providedPorts() {return {};}

  /**
   * @brief First tick transition handler.
   * @return Status from user callback or SUCCESS if absent.
   */
  BT::NodeStatus onStart() override
  {
    if (handlers_.on_start) {return handlers_.on_start(node_, *this);}
    return BT::NodeStatus::SUCCESS;
  }
  /**
   * @brief Repeated tick while status remains RUNNING.
   * @return Status from user callback or SUCCESS if absent.
   */
  BT::NodeStatus onRunning() override
  {
    if (handlers_.on_running) {return handlers_.on_running(node_, *this);}
    return BT::NodeStatus::SUCCESS;
  }
  /**
   * @brief Invoked if the node is halted externally.
   */
  void onHalted() override
  {
    if (handlers_.on_halted) {handlers_.on_halted(node_, *this);}
  }

private:
  std::string name_;              ///< Instance name
  BTActionHandlers handlers_;     ///< Lifecycle callbacks
  rclcpp::Node * node_ {nullptr}; ///< Non-owning node pointer
};

/**
 * @brief Registry encapsulating BehaviorTreeFactory operations and declarative sequence handling.
 *
 * Responsibilities:
 *  - Register arbitrary custom node types (templated) mapping a type_name to a builder.
 *  - Register action node types backed by BTActionHandlers (auto-wrapped with RegisteredActionNode).
 *  - Store named sequences (lists of BTActionSpec) and build executable trees from them.
 */
class BTRegistry
{
public:
  /**
   * @brief Construct the registry with an external BehaviorTreeFactory and ROS node context.
   * @param factory Reference to an existing factory (not owned).
   * @param node ROS 2 node pointer for callbacks (not owned).
   */
  BTRegistry(BT::BehaviorTreeFactory & factory, rclcpp::Node * node)
  : factory_(factory), node_(node) {}

  /**
   * @brief Register a custom behavior tree node type using BehaviorTree.CPP's registerBuilder.
   * @tparam T Node class; must provide constructor (const std::string&, rclcpp::Node*).
   * @param type_name String identifier used in XML.
   * @thread_safety Not thread-safe with other registration calls.
   */
  template<typename T>
  void register_bt_node(const std::string & type_name)
  {
    factory_.registerBuilder<T>(
      type_name, [this](const std::string & name, const BT::NodeConfiguration & config) {
        (void)config;
        return std::make_unique<T>(name, node_);
      });
  }

  /**
   * @brief Register a stateful action node defined by handler callbacks.
   * @param type_name Type identifier used in XML.
   * @param handlers Lifecycle callbacks; may omit any for default behavior.
   * @throws std::runtime_error if a duplicate registration is attempted via internal factory behavior.
   * @thread_safety Not thread-safe with other registration calls.
   */
  void register_bt_action(const std::string & type_name, const BTActionHandlers & handlers)
  {
    bt_action_handlers_[type_name] = handlers;
    factory_.registerBuilder<RegisteredActionNode>(
      type_name, [this, type_name](const std::string & name, const BT::NodeConfiguration & config) {
        (void)config;
        auto it = bt_action_handlers_.find(type_name);
        if (it == bt_action_handlers_.end()) {
          throw std::runtime_error("Handlers not found for BT type: " + type_name);
        }
        return std::make_unique<RegisteredActionNode>(name, it->second, node_);
      });
  }

  /**
   * @brief Register (or overwrite) a named linear sequence specification.
   * @param seq_name Logical sequence identifier.
   * @param seq Ordered list of action specs.
   */
  void register_bt_sequence(const std::string & seq_name, const BTSequenceSpec & seq)
  {
    bt_sequences_[seq_name] = seq;
  }

  /**
   * @brief Build a BehaviorTree.CPP Tree instance from a previously registered sequence.
   * @param seq_name Sequence identifier.
   * @param out_tree Output tree (assigned on success).
   * @return true if the sequence exists and tree creation succeeded; false if sequence not found.
   * @note Fails early with log output if the sequence name is missing.
   */
  bool build_tree_from_sequence(const std::string & seq_name, BT::Tree & out_tree)
  {
    auto it = bt_sequences_.find(seq_name);
    if (it == bt_sequences_.end()) {
      if (node_) {
        RCLCPP_ERROR(node_->get_logger(), "BT sequence '%s' not found", seq_name.c_str());
      }
      return false;
    }
    const std::string xml_text = build_xml_from_sequence(it->second);
    RCLCPP_INFO(node_->get_logger(), "Building BT sequence '%s'", xml_text.c_str());
    out_tree = factory_.createTreeFromText(xml_text);
    return true;
  }

  /**
   * @brief Build a tree directly from raw XML text (already in BTCPP v4 format).
   * @param xml Raw XML document string.
   * @param out_tree Output tree.
   * @return true on success, false if creation throws.
   */
  bool build_tree_from_xml_text(const std::string & xml, BT::Tree & out_tree)
  {
    try {
      out_tree = factory_.createTreeFromText(xml);
      return true;
    } catch (const std::exception & e) {
      if (node_) {
        RCLCPP_ERROR(node_->get_logger(), "Failed to build BT from XML text: %s", e.what());
      }
      return false;
    }
  }

  /**
   * @brief Build a tree from an XML file on disk.
   *
   * Reads the file contents into memory (so that subsequent modifications to the file do not
   * affect the already built tree) and invokes build_tree_from_xml_text(). Any I/O or parsing
   * errors are logged.
   *
   * @param file_path Absolute or relative path to XML file.
   * @param out_tree Output tree.
   * @return true if the file was read and the tree successfully constructed; false otherwise.
   */
  bool build_tree_from_xml_file(const std::string & file_path, BT::Tree & out_tree)
  {
    std::ifstream ifs(file_path);
    if (!ifs.is_open()) {
      if (node_) {
        RCLCPP_ERROR(node_->get_logger(), "Could not open BT XML file: %s", file_path.c_str());
      }
      return false;
    }
    std::ostringstream buffer;
    buffer << ifs.rdbuf();
    if (buffer.str().empty()) {
      if (node_) {
        RCLCPP_ERROR(node_->get_logger(), "BT XML file empty: %s", file_path.c_str());
      }
      return false;
    }
    if (node_) {
      RCLCPP_INFO(node_->get_logger(), "Building BT from XML file: %s \n%s", file_path.c_str(), buffer.str().c_str());
    }
    return build_tree_from_xml_text(buffer.str(), out_tree);
  }

private:
  /**
   * @brief Convert a linear sequence spec into BTCPP v4 XML text.
   * @param seq Sequence specification.
   * @return Complete XML document string.
   */
  std::string build_xml_from_sequence(const BTSequenceSpec & seq) const
  {
    std::ostringstream oss;
    oss <<
      "\n    <root BTCPP_format=\"4\" >\n\n        <BehaviorTree ID=\"MainTree\">\n            <Sequence name=\"root_sequence\">\n";
    for (const auto & act : seq) {
      oss << "                <" << act.type << " name=\"" << act.name << "\"";
      for (const auto & kv : act.ports) {
        oss << " " << kv.first << "=\"" << kv.second << "\"";
      }
      oss << "/>\n";
    }
    oss << "            </Sequence>\n        </BehaviorTree>\n\n    </root>\n";
    return oss.str();
  }

  BT::BehaviorTreeFactory & factory_;                              ///< External factory reference
  rclcpp::Node * node_;                                            ///< Non-owning ROS node pointer
  std::unordered_map<std::string, BTSequenceSpec> bt_sequences_;   ///< Registered sequence specs
  std::unordered_map<std::string, BTActionHandlers> bt_action_handlers_; ///< Action handlers by type
};

#endif // BRAIN_BT_REGISTRY_HPP_
