#include <algorithm>
#include <chrono>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>

#include "geometry_msgs/msg/pose_stamped.hpp"
#include "geometry_msgs/msg/twist.hpp"
#include "nav2_msgs/action/compute_path_to_pose.hpp"
#include "nav2_msgs/action/follow_path.hpp"
#include "nav_msgs/msg/occupancy_grid.hpp"
#include "nav_msgs/msg/odometry.hpp"
#include "nav_msgs/msg/path.hpp"
#include "obstacle_nav2/trajectory_guard.hpp"
#include "rclcpp/rclcpp.hpp"
#include "rclcpp_action/rclcpp_action.hpp"

using namespace std::chrono_literals;

namespace obstacle_nav2
{

class TrajectoryGuardNode : public rclcpp::Node
{
public:
  using ComputePathToPose = nav2_msgs::action::ComputePathToPose;
  using FollowPath = nav2_msgs::action::FollowPath;
  using ComputeGoalHandle = rclcpp_action::ClientGoalHandle<ComputePathToPose>;
  using FollowGoalHandle = rclcpp_action::ClientGoalHandle<FollowPath>;

  TrajectoryGuardNode()
  : Node("trajectory_guard_node")
  {
    declare_parameter("input_path_topic", "/trajectory_guard/input_path");
    declare_parameter("patched_path_topic", "/trajectory_guard/patched_path");
    declare_parameter("reject_path_topic", "/reject_path");
    declare_parameter("costmap_topic", "/local_costmap/costmap");
    declare_parameter("planner_costmap_topic", "/global_costmap/costmap");
    declare_parameter("odom_topic", "/odom_combined");
    declare_parameter("planner_action", "/compute_path_to_pose");
    declare_parameter("follow_action", "/follow_path");
    declare_parameter("planner_id", "GridBased");
    declare_parameter("controller_id", "FollowPath");
    declare_parameter("goal_checker_id", "");
    declare_parameter("check_period_sec", 0.5);
    declare_parameter("lookahead_distance", 2.0);
    declare_parameter("rejoin_min_distance", 1.0);
    declare_parameter("rejoin_max_distance", 4.0);
    declare_parameter("occupied_threshold", 50);
    declare_parameter("treat_unknown_as_occupied", true);
    declare_parameter("footprint_half_length", 0.14);
    declare_parameter("footprint_half_width", 0.085);
    declare_parameter("footprint_padding", 0.04);
    declare_parameter("footprint_sample_step", 0.05);
    declare_parameter("repair_cooldown_sec", 1.0);
    declare_parameter("blocked_retry_wait_sec", 0.8);
    declare_parameter("wait_for_planner_costmap_update", true);
    declare_parameter("min_repair_progress_indices", 5);
    declare_parameter("publish_zero_on_blocked", true);
    declare_parameter("cmd_vel_topic", "/cmd_vel");
    declare_parameter("execute_follow_path", true);

    loadParameters();
    last_repair_request_time_ =
      get_clock()->now() - rclcpp::Duration::from_seconds(repair_cooldown_sec_);

    path_sub_ = create_subscription<nav_msgs::msg::Path>(
      get_parameter("input_path_topic").as_string(), 1,
      [this](nav_msgs::msg::Path::SharedPtr msg) { onPath(msg); });
    costmap_sub_ = create_subscription<nav_msgs::msg::OccupancyGrid>(
      get_parameter("costmap_topic").as_string(), 1,
      [this](nav_msgs::msg::OccupancyGrid::SharedPtr msg) {
        latest_costmap_ = msg;
        ++local_costmap_generation_;
      });
    planner_costmap_sub_ = create_subscription<nav_msgs::msg::OccupancyGrid>(
      get_parameter("planner_costmap_topic").as_string(), 1,
      [this](nav_msgs::msg::OccupancyGrid::SharedPtr msg) {
        latest_planner_costmap_ = msg;
        ++planner_costmap_generation_;
      });
    odom_sub_ = create_subscription<nav_msgs::msg::Odometry>(
      get_parameter("odom_topic").as_string(), 10,
      [this](nav_msgs::msg::Odometry::SharedPtr msg) { latest_odom_ = msg; });

    patched_path_pub_ = create_publisher<nav_msgs::msg::Path>(
      get_parameter("patched_path_topic").as_string(), 1);
    reject_path_pub_ = create_publisher<nav_msgs::msg::Path>(
      get_parameter("reject_path_topic").as_string(), 1);
    cmd_vel_pub_ = create_publisher<geometry_msgs::msg::Twist>(
      get_parameter("cmd_vel_topic").as_string(), 1);

    planner_client_ = rclcpp_action::create_client<ComputePathToPose>(
      this, get_parameter("planner_action").as_string());
    follow_client_ = rclcpp_action::create_client<FollowPath>(
      this, get_parameter("follow_action").as_string());

    const auto period = std::chrono::duration<double>(
      std::max(0.1, get_parameter("check_period_sec").as_double()));
    timer_ = create_wall_timer(
      std::chrono::duration_cast<std::chrono::milliseconds>(period),
      [this]() { tick(); });

    RCLCPP_INFO(
      get_logger(), "trajectory guard ready check_period_sec=%.2f input=%s patched=%s",
      get_parameter("check_period_sec").as_double(),
      get_parameter("input_path_topic").as_string().c_str(),
      get_parameter("patched_path_topic").as_string().c_str());
  }

private:
  void loadParameters()
  {
    settings_.lookahead_distance = get_parameter("lookahead_distance").as_double();
    settings_.rejoin_min_distance = get_parameter("rejoin_min_distance").as_double();
    settings_.rejoin_max_distance = get_parameter("rejoin_max_distance").as_double();
    settings_.occupied_threshold = get_parameter("occupied_threshold").as_int();
    settings_.treat_unknown_as_occupied =
      get_parameter("treat_unknown_as_occupied").as_bool();
    settings_.footprint_half_length = get_parameter("footprint_half_length").as_double();
    settings_.footprint_half_width = get_parameter("footprint_half_width").as_double();
    settings_.footprint_padding = get_parameter("footprint_padding").as_double();
    settings_.footprint_sample_step = get_parameter("footprint_sample_step").as_double();
    repair_cooldown_sec_ = std::max(0.0, get_parameter("repair_cooldown_sec").as_double());
    blocked_retry_wait_sec_ =
      std::max(0.0, get_parameter("blocked_retry_wait_sec").as_double());
    wait_for_planner_costmap_update_ =
      get_parameter("wait_for_planner_costmap_update").as_bool();
    min_repair_progress_indices_ =
      std::max(0, static_cast<int>(get_parameter("min_repair_progress_indices").as_int()));
    publish_zero_on_blocked_ = get_parameter("publish_zero_on_blocked").as_bool();
    execute_follow_path_ = get_parameter("execute_follow_path").as_bool();
  }

  void onPath(const nav_msgs::msg::Path::SharedPtr msg)
  {
    active_path_ = *msg;
    sent_original_for_current_path_ = false;
    last_repair_nearest_index_.reset();
    last_repair_rejoin_index_.reset();
    retry_without_planner_costmap_update_ = false;
    waiting_for_planner_costmap_update_ = false;
    RCLCPP_INFO(
      get_logger(), "received replay path poses=%zu frame=%s",
      active_path_.poses.size(), active_path_.header.frame_id.c_str());
  }

  std::optional<geometry_msgs::msg::PoseStamped> currentPose() const
  {
    if (!latest_odom_) {
      return std::nullopt;
    }
    geometry_msgs::msg::PoseStamped pose;
    pose.header = latest_odom_->header;
    pose.pose = latest_odom_->pose.pose;
    return pose;
  }

  void tick()
  {
    if (active_path_.poses.empty() || !latest_costmap_) {
      return;
    }

    const auto pose = currentPose();
    if (!pose.has_value()) {
      return;
    }

    const auto nearest = nearestPathIndex(active_path_, *pose);
    const auto result = checkPathAhead(active_path_, nearest, *latest_costmap_, settings_);
    if (!result.blocked) {
      if (!sent_original_for_current_path_) {
        publishAndMaybeFollow(active_path_, "original_clear");
        sent_original_for_current_path_ = true;
      }
      return;
    }

    const auto now = get_clock()->now();
    if (repair_in_flight_ ||
      (now - last_repair_request_time_).seconds() < repair_cooldown_sec_)
    {
      return;
    }

    if (!waiting_for_planner_costmap_update_) {
      waiting_for_planner_costmap_update_ = true;
      waiting_planner_costmap_generation_ = planner_costmap_generation_;
      RCLCPP_WARN_THROTTLE(
        get_logger(), *get_clock(), 2000,
        "blocked replay path; waiting for next planner costmap frame before requesting bypass");
      publishStop();
      return;
    }

    if (planner_costmap_generation_ <= waiting_planner_costmap_generation_) {
      if (retry_without_planner_costmap_update_ &&
        (now - last_repair_request_time_).seconds() >= blocked_retry_wait_sec_)
      {
        waiting_for_planner_costmap_update_ = false;
      } else {
        RCLCPP_WARN_THROTTLE(
          get_logger(), *get_clock(), 2000,
          "blocked replay path; planner costmap has not updated yet");
        publishStop();
        return;
      }
    }

    waiting_for_planner_costmap_update_ = false;

    const bool costmap_updated_since_repair =
      !wait_for_planner_costmap_update_ ||
      planner_costmap_generation_ > last_repair_planner_costmap_generation_;
    const bool timeout_retry_allowed =
      retry_without_planner_costmap_update_ &&
      (now - last_repair_request_time_).seconds() >= blocked_retry_wait_sec_;
    if (!shouldRetryBlockedRepair(
        last_repair_nearest_index_, nearest, min_repair_progress_indices_,
        (now - last_repair_request_time_).seconds(), blocked_retry_wait_sec_,
        costmap_updated_since_repair) && !timeout_retry_allowed)
    {
      RCLCPP_WARN_THROTTLE(
        get_logger(), *get_clock(), 2000,
        "patched path still blocked near the last repair point; waiting for path progress or costmap update");
      publishStop();
      return;
    }

    retry_without_planner_costmap_update_ = false;

    std::size_t search_start = nearest;
    if (last_repair_rejoin_index_.has_value()) {
      const auto bump = static_cast<std::size_t>(std::max(1, min_repair_progress_indices_));
      const auto advanced = std::min(
        active_path_.poses.size() - 1, *last_repair_rejoin_index_ + bump);
      search_start = std::max(search_start, advanced);
    }

    const auto rejoin = findClearRejoinIndexWithPlannerCostmap(
      active_path_, search_start, *latest_costmap_, latest_planner_costmap_.get(), settings_);
    if (!rejoin.has_value()) {
      RCLCPP_WARN(get_logger(), "blocked replay path but no clear rejoin point found");
      publishStop();
      return;
    }

    RCLCPP_WARN(
      get_logger(), "blocked replay path index=%zu reason=%s rejoin_index=%zu",
      result.path_index, result.reason.c_str(), *rejoin);
    requestBypass(nearest, *rejoin);
  }

  void publishAndMaybeFollow(const nav_msgs::msg::Path & path, const std::string & label)
  {
    if (path.poses.empty()) {
      return;
    }

    nav_msgs::msg::Path stamped_path = path;
    stamped_path.header.stamp = now();
    for (auto & pose : stamped_path.poses) {
      pose.header.stamp = stamped_path.header.stamp;
      if (pose.header.frame_id.empty()) {
        pose.header.frame_id = stamped_path.header.frame_id;
      }
    }
    patched_path_pub_->publish(stamped_path);

    if (!execute_follow_path_) {
      RCLCPP_INFO(
        get_logger(), "published %s path poses=%zu without FollowPath execution",
        label.c_str(), stamped_path.poses.size());
      return;
    }

    sendFollowPath(stamped_path, label);
  }

  void sendFollowPath(const nav_msgs::msg::Path & path, const std::string & label)
  {
    if (!follow_client_->wait_for_action_server(100ms)) {
      RCLCPP_WARN(get_logger(), "FollowPath action server not ready");
      return;
    }

    if (active_follow_goal_handle_) {
      follow_client_->async_cancel_goal(active_follow_goal_handle_);
      active_follow_goal_handle_.reset();
    }

    FollowPath::Goal goal;
    goal.path = path;
    goal.controller_id = get_parameter("controller_id").as_string();
    goal.goal_checker_id = get_parameter("goal_checker_id").as_string();

    auto options = rclcpp_action::Client<FollowPath>::SendGoalOptions();
    options.goal_response_callback =
      [this, label](FollowGoalHandle::SharedPtr goal_handle) {
        if (!goal_handle) {
          RCLCPP_WARN(get_logger(), "FollowPath %s rejected", label.c_str());
          return;
        }
        active_follow_goal_handle_ = goal_handle;
        RCLCPP_INFO(get_logger(), "FollowPath %s accepted", label.c_str());
      };
    options.result_callback =
      [this, label](const FollowGoalHandle::WrappedResult & result) {
        active_follow_goal_handle_.reset();
        RCLCPP_INFO(
          get_logger(), "FollowPath %s finished code=%d",
          label.c_str(), static_cast<int>(result.code));
      };

    follow_client_->async_send_goal(goal, options);
    RCLCPP_INFO(get_logger(), "sent FollowPath %s poses=%zu", label.c_str(), path.poses.size());
  }

  void requestBypass(std::size_t nearest_index, std::size_t rejoin_index)
  {
    const auto start = currentPose();
    if (!start.has_value() || rejoin_index >= active_path_.poses.size()) {
      return;
    }

    if (!planner_client_->wait_for_action_server(100ms)) {
      RCLCPP_WARN(get_logger(), "ComputePathToPose action server not ready");
      publishStop();
      return;
    }

    repair_in_flight_ = true;
    last_repair_request_time_ = get_clock()->now();
    last_repair_nearest_index_ = nearest_index;
    last_repair_rejoin_index_ = rejoin_index;
    last_repair_planner_costmap_generation_ = planner_costmap_generation_;
    waiting_for_planner_costmap_update_ = false;

    ComputePathToPose::Goal goal;
    goal.start = *start;
    goal.goal = active_path_.poses[rejoin_index];
    goal.planner_id = get_parameter("planner_id").as_string();
    goal.use_start = true;

    auto options = rclcpp_action::Client<ComputePathToPose>::SendGoalOptions();
    options.goal_response_callback =
      [this](ComputeGoalHandle::SharedPtr goal_handle) {
        if (!goal_handle) {
          repair_in_flight_ = false;
          RCLCPP_WARN(get_logger(), "bypass planning goal rejected");
        }
      };
    options.result_callback =
      [this, rejoin_index](const ComputeGoalHandle::WrappedResult & result) {
        repair_in_flight_ = false;
        if (result.code != rclcpp_action::ResultCode::SUCCEEDED ||
          result.result->path.poses.empty())
        {
          RCLCPP_WARN(get_logger(), "bypass planning failed code=%d", static_cast<int>(result.code));
          publishStop();
          return;
        }

        const auto patched = stitchPaths(result.result->path, active_path_, rejoin_index);
        if (!patchedPathIsClear(patched)) {
          reject_path_pub_->publish(patched);
          retry_without_planner_costmap_update_ = true;
          waiting_for_planner_costmap_update_ = false;
          RCLCPP_WARN(
            get_logger(),
            "planned bypass still blocked in local costmap; stopping and retrying from a later rejoin point");
          publishStop();
          return;
        }

        active_path_ = patched;
        sent_original_for_current_path_ = true;
        waiting_for_planner_costmap_update_ = false;
        retry_without_planner_costmap_update_ = false;
        publishAndMaybeFollow(patched, "patched_bypass");
      };

    planner_client_->async_send_goal(goal, options);
    RCLCPP_INFO(get_logger(), "requested bypass to rejoin_index=%zu", rejoin_index);
  }

  bool patchedPathIsClear(const nav_msgs::msg::Path & patched) const
  {
    if (!latest_costmap_) {
      return false;
    }

    const auto pose = currentPose();
    if (!pose.has_value()) {
      return false;
    }

    const auto nearest = nearestPathIndex(patched, *pose);
    const auto result = checkPathAhead(patched, nearest, *latest_costmap_, settings_);
    return !result.blocked;
  }

  void publishStop()
  {
    if (!publish_zero_on_blocked_) {
      return;
    }
    geometry_msgs::msg::Twist stop;
    cmd_vel_pub_->publish(stop);
  }

  GuardSettings settings_;
  double repair_cooldown_sec_{1.0};
  double blocked_retry_wait_sec_{0.8};
  int min_repair_progress_indices_{5};
  bool publish_zero_on_blocked_{true};
  bool execute_follow_path_{true};
  bool wait_for_planner_costmap_update_{true};
  bool retry_without_planner_costmap_update_{false};
  bool repair_in_flight_{false};
  bool sent_original_for_current_path_{false};
  bool waiting_for_planner_costmap_update_{false};
  rclcpp::Time last_repair_request_time_{0, 0, RCL_ROS_TIME};
  std::optional<std::size_t> last_repair_nearest_index_;
  std::optional<std::size_t> last_repair_rejoin_index_;
  std::uint64_t local_costmap_generation_{0};
  std::uint64_t planner_costmap_generation_{0};
  std::uint64_t last_repair_planner_costmap_generation_{0};
  std::uint64_t waiting_planner_costmap_generation_{0};

  nav_msgs::msg::Path active_path_;
  nav_msgs::msg::OccupancyGrid::SharedPtr latest_costmap_;
  nav_msgs::msg::OccupancyGrid::SharedPtr latest_planner_costmap_;
  nav_msgs::msg::Odometry::SharedPtr latest_odom_;
  FollowGoalHandle::SharedPtr active_follow_goal_handle_;

  rclcpp::Subscription<nav_msgs::msg::Path>::SharedPtr path_sub_;
  rclcpp::Subscription<nav_msgs::msg::OccupancyGrid>::SharedPtr costmap_sub_;
  rclcpp::Subscription<nav_msgs::msg::OccupancyGrid>::SharedPtr planner_costmap_sub_;
  rclcpp::Subscription<nav_msgs::msg::Odometry>::SharedPtr odom_sub_;
  rclcpp::Publisher<nav_msgs::msg::Path>::SharedPtr patched_path_pub_;
  rclcpp::Publisher<nav_msgs::msg::Path>::SharedPtr reject_path_pub_;
  rclcpp::Publisher<geometry_msgs::msg::Twist>::SharedPtr cmd_vel_pub_;
  rclcpp_action::Client<ComputePathToPose>::SharedPtr planner_client_;
  rclcpp_action::Client<FollowPath>::SharedPtr follow_client_;
  rclcpp::TimerBase::SharedPtr timer_;
};

}  // namespace obstacle_nav2

int main(int argc, char ** argv)
{
  rclcpp::init(argc, argv);
  rclcpp::spin(std::make_shared<obstacle_nav2::TrajectoryGuardNode>());
  rclcpp::shutdown();
  return 0;
}
