#pragma once

#include <chrono>
#include <functional>
#include <memory>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>
#include <future>
#include <optional>

#include <rclcpp/rclcpp.hpp>
#include <rclcpp_action/rclcpp_action.hpp>

/**
 * @file action_registry.hpp
 * @brief Utilities for registering multiple ROS 2 action servers and clients via simple name lookups.
 *
 * This header defines two registry classes:
 *  - ActionServerRegistry: Register and own multiple action servers via a uniform handler bundle.
 *  - ActionClientRegistry: Register and trigger multiple action clients by name with reduced boilerplate.
 *
 * Design goals:
 *  - Reduce repetitive action server/client setup code inside nodes.
 *  - Provide a consistent pattern (name -> handlers) for adding new action types.
 *  - Ensure lifetime safety by storing shared pointers inside the registry objects.
 *
 * Notes:
 *  - Both registries are independent; either can be used alone.
 *  - The logical name should normally match the action topic (e.g. "move_to_position").
 *  - For servers, if only an execute functor is provided (no custom on_accepted), a worker thread is spawned.
 *
 * Thread-safety:
 *  - Registration methods are expected to be called during initialization (single-threaded).
 *  - Sending goals / cancelling may be invoked concurrently; underlying rclcpp action APIs are thread-safe.
 */

class ActionServerRegistry
{
public:
  /**
   * @brief Construct a server registry bound to a node.
   * @param node Owning rclcpp node (not owned).
   */
  explicit ActionServerRegistry(rclcpp::Node * node)
  : node_(node) {}

  /**
   * @brief Bundle of callbacks used to configure an action server.
   * @tparam ActionT Action type.
   *
   * Provide either (on_accepted) OR (execute). If on_accepted is empty and execute is set,
   * the registry spawns a detached thread invoking execute() so it does not block the executor.
   */
  template<typename ActionT>
  struct Handlers
  {
    using Goal = typename ActionT::Goal;
    using GoalHandle = rclcpp_action::ServerGoalHandle<ActionT>;
    /** @brief New goal request handler; return ACCEPT_AND_EXECUTE, REJECT, etc. */
    std::function<rclcpp_action::GoalResponse(
        const rclcpp_action::GoalUUID &,
        std::shared_ptr<const Goal>)> on_goal;
    /** @brief Cancellation request handler (optional). Return ACCEPT / REJECT. */
    std::function<rclcpp_action::CancelResponse(const std::shared_ptr<GoalHandle>)> on_cancel;
    /** @brief Handler invoked immediately after a goal is accepted (optional). */
    std::function<void(const std::shared_ptr<GoalHandle>)> on_accepted;  ///< Optional
    /** @brief Execute callback (optional, used if on_accepted is omitted). */
    std::function<void(const std::shared_ptr<GoalHandle>)> execute;      ///< Optional
  };

  /**
   * @brief Register (create and retain) an action server instance.
   * @tparam ActionT Action type.
   * @param name Action topic name.
   * @param handlers Callback bundle.
   * @thread_safety Call during initialization; not thread-safe with concurrent registration.
   */
  template<typename ActionT>
  void register_action_server(const std::string & name, const Handlers<ActionT> & handlers)
  {
    holders_.push_back(std::make_unique<Holder<ActionT>>(node_, name, handlers));
  }

private:
  rclcpp::Node * node_;
  struct IHolder { virtual ~IHolder() = default; };
  // Keep concrete holders alive via base pointer
  std::vector<std::unique_ptr<IHolder>> holders_;

  template<typename AT>
  struct Holder : IHolder
  {
    using GoalHandle = rclcpp_action::ServerGoalHandle<AT>;
    typename rclcpp_action::Server<AT>::SharedPtr server_;
    Holder(rclcpp::Node * node, const std::string & name, const Handlers<AT> & h)
    {
      auto goal_cb = [h](const rclcpp_action::GoalUUID & uuid,
          std::shared_ptr<const typename AT::Goal> goal) {
          return h.on_goal ? h.on_goal(
            uuid,
            goal) : rclcpp_action::GoalResponse::ACCEPT_AND_EXECUTE;
        };
      auto cancel_cb = [h](const std::shared_ptr<GoalHandle> gh) {
          return h.on_cancel ? h.on_cancel(gh) : rclcpp_action::CancelResponse::ACCEPT;
        };
      auto accepted_cb = [h](const std::shared_ptr<GoalHandle> gh) {
          if (h.on_accepted) {
            h.on_accepted(gh);
          } else if (h.execute) {
            std::thread{[h, gh]() {h.execute(gh);}}.detach();
          }
        };
      server_ = rclcpp_action::create_server<AT>(node, name, goal_cb, cancel_cb, accepted_cb);
    }
  };
};

/**
 * @brief Registry to manage and trigger multiple action clients by name.
 *
 * Provides simplified registration with goal construction and callback wiring, as well as
 * convenience methods to send, cancel, wait for server availability, and synchronously
 * wait for a result (with timeout and optional cancellation).
 */
class ActionClientRegistry
{
public:
  /**
   * @brief Construct a client registry bound to a node.
   * @param node Owning rclcpp node (not owned).
   */
  explicit ActionClientRegistry(rclcpp::Node * node)
  : node_(node) {}

  template<typename ActionT>
  using GoalHandle = rclcpp_action::ClientGoalHandle<ActionT>;

  /**
   * @brief Abstract base for stored client entries (type-erased).
   */
  struct EntryBase
  {
    virtual ~EntryBase() = default;
    /** @brief Send a goal using the entry's callbacks. */
    virtual void send(rclcpp::Logger logger) = 0;
    /** @brief Check if server is available within timeout. */
    virtual bool available(std::chrono::nanoseconds timeout) const = 0;
    /** @brief Attempt to cancel last goal if present. */
    virtual void cancel(rclcpp::Logger logger) = 0;
  };

  template<typename ActionT>
  struct Entry : EntryBase
  {
    using Goal = typename ActionT::Goal;
    typename rclcpp_action::Client<ActionT>::SharedPtr client;         ///< Underlying client instance
    std::weak_ptr<GoalHandle<ActionT>> last_goal_handle;               ///< Last goal handle (weak)
    std::function<Goal()> make_goal;                                   ///< Goal builder
    std::function<void(const std::shared_ptr<GoalHandle<ActionT>> &)> on_goal_response; ///< Optional goal response
    std::function<void(typename GoalHandle<ActionT>::SharedPtr,
      const std::shared_ptr<const typename ActionT::Feedback>)> on_feedback;                                                           ///< Optional feedback
    std::function<void(const typename GoalHandle<ActionT>::WrappedResult &)> on_result; ///< Optional result
    std::string name;                                                  ///< Action name/topic

    Entry(
      rclcpp::Node * node,
      const std::string & name,
      std::function<Goal()> mk,
      std::function<void(typename GoalHandle<ActionT>::SharedPtr,
      const std::shared_ptr<const typename ActionT::Feedback>)> fb,
      std::function<void(const typename GoalHandle<ActionT>::WrappedResult &)> rc,
      std::function<void(const std::shared_ptr<GoalHandle<ActionT>> &)> gr)
    : make_goal(std::move(mk)), on_goal_response(std::move(gr)), on_feedback(std::move(fb)),
      on_result(std::move(rc)), name(name)
    {
      client = rclcpp_action::create_client<ActionT>(node, name);
    }

    void send(rclcpp::Logger logger) override
    {
      if (!client) {return;}
      RCLCPP_INFO(logger, "[ActionClientRegistry] waiting for server '%s'", name.c_str());
      if (!client->wait_for_action_server(std::chrono::seconds(10))) {
        RCLCPP_ERROR(logger, "Action server '%s' not available", name.c_str());
        return;
      }
      auto goal_msg = make_goal ? make_goal() : Goal{};
      RCLCPP_INFO(logger, "[ActionClientRegistry] sending goal on '%s'", name.c_str());

      auto opts = typename rclcpp_action::Client<ActionT>::SendGoalOptions{};
      if (on_goal_response) {opts.goal_response_callback = on_goal_response;}
      if (on_feedback) {opts.feedback_callback = on_feedback;}
      if (on_result) {opts.result_callback = on_result;}
      auto future_handle = client->async_send_goal(goal_msg, opts);
      std::thread(
        [future_handle, this, logger]() mutable {
          if (future_handle.wait_for(std::chrono::seconds(0)) != std::future_status::ready) {
            future_handle.wait();
          }
          try {
            auto gh = future_handle.get();
            if (gh) {
              last_goal_handle = gh;
              RCLCPP_INFO(logger, "Stored goal handle for '%s'", name.c_str());
            }
          } catch (const std::exception & e) {
            RCLCPP_WARN(
              logger, "Failed to obtain goal handle for '%s': %s", name.c_str(),
              e.what());
          }
        }).detach();
    }

    bool available(std::chrono::nanoseconds timeout) const override
    {
      if (!client) {return false;}
      return client->wait_for_action_server(timeout);
    }

    void cancel(rclcpp::Logger logger) override
    {
      auto gh = last_goal_handle.lock();
      if (!gh) {
        RCLCPP_WARN(logger, "No active goal handle to cancel for '%s'", name.c_str());
        return;
      }
      RCLCPP_INFO(logger, "Cancelling goal on '%s'", name.c_str());
      client->async_cancel_goal(gh);
    }
  };

  /**
   * @brief Register an action client and associated callbacks.
   * @tparam ActionT Action type.
   * @param name Action topic name.
   * @param make_goal Goal construction functor.
   * @param on_feedback Optional feedback callback.
   * @param on_result Optional result callback.
   * @param on_goal_response Optional goal response callback.
   * @thread_safety Call during initialization or in a controlled context; not thread-safe with concurrent registration.
   */
  template<typename ActionT>
  void register_client(
    const std::string & name,
    std::function<typename ActionT::Goal()> make_goal,
    std::function<void(typename GoalHandle<ActionT>::SharedPtr, const std::shared_ptr<const typename ActionT::Feedback>)> on_feedback,
    std::function<void(const typename GoalHandle<ActionT>::WrappedResult &)> on_result,
    std::function<void(const std::shared_ptr<GoalHandle<ActionT>> &)> on_goal_response = nullptr)
  {
    entries_[name] =
      std::make_unique<Entry<ActionT>>(
      node_, name, std::move(make_goal), std::move(
        on_feedback), std::move(on_result), std::move(on_goal_response));
  }

  /**
   * @brief Send a goal on a registered client by name.
   * @param name Action client key.
   * @param logger Logger instance.
   */
  void send(const std::string & name, rclcpp::Logger logger)
  {
    auto it = entries_.find(name);
    if (it == entries_.end()) {
      RCLCPP_ERROR(logger, "No action client registered for '%s'", name.c_str());
      return;
    }
    RCLCPP_INFO(logger, "Sending action '%s'", name.c_str());
    it->second->send(logger);
  }

  /**
   * @brief Cancel the last goal sent on a registered client (if any).
   * @param name Action client key.
   * @param logger Logger instance.
   */
  void cancel(const std::string & name, rclcpp::Logger logger)
  {
    auto it = entries_.find(name);
    if (it == entries_.end()) {
      RCLCPP_ERROR(logger, "No action client registered for '%s'", name.c_str());
      return;
    }
    it->second->cancel(logger);
  }

  /**
   * @brief Wait for a server to become available on a registered client.
   * @param name Action client key.
   * @param timeout Maximum wait duration.
   * @return true if available within timeout.
   */
  bool wait_for_server(const std::string & name, std::chrono::nanoseconds timeout)
  {
    auto it = entries_.find(name);
    if (it == entries_.end()) {return false;}
    return it->second->available(timeout);
  }

  /**
   * @brief Check whether a client has been registered.
   * @param name Action client key.
   * @return true if present.
   */
  bool has_client(const std::string & name) const {return entries_.find(name) != entries_.end();}

  /**
   * @brief Lightweight availability check for a named client.
   * @param name Action client key.
   * @param timeout Wait duration.
   * @return true if server responds within timeout.
   */
  bool is_server_available(const std::string & name, std::chrono::nanoseconds timeout) const
  {
    auto it = entries_.find(name);
    if (it == entries_.end()) {return false;}
    return it->second->available(timeout);
  }

  /**
   * @brief Convenience: send a goal and block until result or timeout.
   * @tparam ActionT Action type.
   * @param name Action client key.
   * @param logger Logger instance.
   * @param timeout Overall result wait timeout.
   * @return WrappedResult if completed in time; std::nullopt on timeout or error.
   */
  template<typename ActionT>
  std::optional<typename rclcpp_action::ClientGoalHandle<ActionT>::WrappedResult> send_and_wait(
    const std::string & name,
    rclcpp::Logger logger,
    std::chrono::nanoseconds timeout = std::chrono::seconds(30))
  {
    auto it = entries_.find(name);
    if (it == entries_.end()) {
      RCLCPP_ERROR(logger, "No action client registered for '%s'", name.c_str());
      return std::nullopt;
    }
    auto * e = dynamic_cast<Entry<ActionT> *>(it->second.get());
    if (!e) {
      RCLCPP_ERROR(logger, "Action client type mismatch for '%s'", name.c_str());
      return std::nullopt;
    }
    if (!e->client) {return std::nullopt;}
    if (!e->client->wait_for_action_server(std::chrono::seconds(10))) {
      RCLCPP_ERROR(logger, "Action server '%s' not available", name.c_str());
      return std::nullopt;
    }
    auto goal = e->make_goal ? e->make_goal() : typename ActionT::Goal{};
    auto opts = typename rclcpp_action::Client<ActionT>::SendGoalOptions{};
    if (e->on_feedback) {opts.feedback_callback = e->on_feedback;}
    auto goal_future = e->client->async_send_goal(goal, opts);
    auto goal_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5);
    if (goal_future.wait_until(goal_deadline) != std::future_status::ready) {
      RCLCPP_ERROR(logger, "Timed out waiting for goal response on '%s'", name.c_str());
      return std::nullopt;
    }

    typename rclcpp_action::ClientGoalHandle<ActionT>::SharedPtr goal_handle;
    try {
      goal_handle = goal_future.get();
    } catch (const std::exception & ex) {
      RCLCPP_ERROR(logger, "Failed to get goal handle for '%s': %s", name.c_str(), ex.what());
      return std::nullopt;
    }
    if (!goal_handle) {
      RCLCPP_WARN(logger, "Goal rejected for '%s'", name.c_str());
      return std::nullopt;
    }
    e->last_goal_handle = goal_handle;
    if (e->on_goal_response) {
      try {
        e->on_goal_response(goal_handle);
      } catch (const std::exception & ex) {
        RCLCPP_WARN(logger, "Goal response callback threw for '%s': %s", name.c_str(), ex.what());
      }
    }

    auto result_future = e->client->async_get_result(goal_handle);
    auto deadline = std::chrono::steady_clock::now() + timeout;
    if (result_future.wait_until(deadline) != std::future_status::ready) {
      RCLCPP_ERROR(logger, "Timed out waiting for result on '%s', sending cancel", name.c_str());
      try {
        e->client->async_cancel_goal(goal_handle);
      } catch (const std::exception & ex) {
        RCLCPP_WARN(logger, "Failed cancelling goal for '%s': %s", name.c_str(), ex.what());
      }
      return std::nullopt;
    }
    typename rclcpp_action::ClientGoalHandle<ActionT>::WrappedResult wrapped_result;
    try {
      wrapped_result = result_future.get();
    } catch (const std::exception & ex) {
      RCLCPP_ERROR(logger, "Failed to get result for '%s': %s", name.c_str(), ex.what());
      return std::nullopt;
    }
    if (e->on_result) {
      try {
        e->on_result(wrapped_result);
      } catch (const std::exception & ex) {
        RCLCPP_WARN(logger, "Result callback threw for '%s': %s", name.c_str(), ex.what());
      }
    }
    return wrapped_result;
  }

private:
  rclcpp::Node * node_;
  std::unordered_map<std::string, std::unique_ptr<EntryBase>> entries_;
};
