#include "rclcpp/rclcpp.hpp"
#include "sensor_msgs/msg/joint_state.hpp"
#include "trajectory_msgs/msg/joint_trajectory.hpp"
#include "control_msgs/action/follow_joint_trajectory.hpp"
#include "rclcpp_action/rclcpp_action.hpp"
#include <vector>
#include <string>
#include <memory>
#include <chrono>

using namespace std::chrono_literals;
using FollowJointTrajectory = control_msgs::action::FollowJointTrajectory;
using GoalHandleFollowJointTrajectory = rclcpp_action::ClientGoalHandle<FollowJointTrajectory>;

class ArmController : public rclcpp::Node
{
public:
  ArmController() : Node("arm_controller")
  {
    // 声明参数
    this->declare_parameter<std::vector<std::string>>("joint_names", 
      std::vector<std::string>{"shoulder_pan_joint", "shoulder_lift_joint", 
                              "elbow_joint", "wrist_1_joint", 
                              "wrist_2_joint", "wrist_3_joint"});
    
    // 获取关节名称
    joint_names_ = this->get_parameter("joint_names").as_string_array();
    
    // 创建关节状态发布者
    joint_state_pub_ = this->create_publisher<sensor_msgs::msg::JointState>(
      "joint_states", 10);
    
    // 创建轨迹动作客户端
    trajectory_client_ = rclcpp_action::create_client<FollowJointTrajectory>(
      this, "follow_joint_trajectory");
    
    // 创建定时器，定期发布关节状态
    timer_ = this->create_wall_timer(
      10ms, std::bind(&ArmController::publishJointStates, this));
    
    // 初始化关节状态
    initJointStates();
    
    RCLCPP_INFO(this->get_logger(), "Arm controller initialized");
  }
  
  // 执行关节轨迹
  void executeTrajectory(const std::vector<std::vector<double>>& waypoints, 
                         const std::vector<double>& time_points)
  {
    // 等待动作服务器可用
    if (!trajectory_client_->wait_for_action_server(10s)) {
      RCLCPP_ERROR(this->get_logger(), "Action server not available after waiting");
      return;
    }
    
    // 创建目标
    auto goal_msg = FollowJointTrajectory::Goal();
    goal_msg.trajectory.joint_names = joint_names_;
    
    // 添加轨迹点
    for (size_t i = 0; i < waypoints.size(); ++i) {
      trajectory_msgs::msg::JointTrajectoryPoint point;
      point.positions = waypoints[i];
      
      // 设置时间戳
      rclcpp::Duration time_from_start = rclcpp::Duration::from_seconds(time_points[i]);
      builtin_interfaces::msg::Duration duration_msg;
      rclcpp::Duration rclcpp_duration = rclcpp::Duration::from_seconds(time_points[i]);
      
      // 设置秒和纳秒部分
      duration_msg.sec = rclcpp_duration.seconds();
      duration_msg.nanosec = rclcpp_duration.nanoseconds() % 1000000000;
      
      point.time_from_start = duration_msg;
      
      goal_msg.trajectory.points.push_back(point);
    }
    
    // 发送目标
    auto send_goal_options = rclcpp_action::Client<FollowJointTrajectory>::SendGoalOptions();
    send_goal_options.feedback_callback = 
      [this](GoalHandleFollowJointTrajectory::SharedPtr, 
             const std::shared_ptr<const FollowJointTrajectory::Feedback> feedback) {
        RCLCPP_INFO(this->get_logger(), "Current position: %.2f, %.2f, %.2f",
                   feedback->actual.positions[0],
                   feedback->actual.positions[1],
                   feedback->actual.positions[2]);
      };
    
    send_goal_options.result_callback = 
      [this](const GoalHandleFollowJointTrajectory::WrappedResult& result) {
        switch (result.code) {
          case rclcpp_action::ResultCode::SUCCEEDED:
            RCLCPP_INFO(this->get_logger(), "Trajectory execution succeeded");
            break;
          case rclcpp_action::ResultCode::ABORTED:
            RCLCPP_ERROR(this->get_logger(), "Trajectory execution aborted");
            return;
          case rclcpp_action::ResultCode::CANCELED:
            RCLCPP_INFO(this->get_logger(), "Trajectory execution canceled");
            return;
          default:
            RCLCPP_ERROR(this->get_logger(), "Unknown result code");
            return;
        }
      };
    
    auto goal_handle_future = trajectory_client_->async_send_goal(goal_msg, send_goal_options);
    
    if (rclcpp::spin_until_future_complete(this->get_node_base_interface(), goal_handle_future) !=
        rclcpp::FutureReturnCode::SUCCESS) {
      RCLCPP_ERROR(this->get_logger(), "send goal call failed :(");
      return;
    }
    
    goal_handle_ = goal_handle_future.get();
    if (!goal_handle_) {
      RCLCPP_ERROR(this->get_logger(), "Goal was rejected by server");
      return;
    }
    
    RCLCPP_INFO(this->get_logger(), "Goal accepted by server, waiting for result");
  }
  
  // 简单的运动示例：将机械臂移动到预设位置
  void moveToHomePosition()
  {
    // 定义路径点（关节角度）
    std::vector<std::vector<double>> waypoints;
    std::vector<double> time_points;
    
    // 当前位置（假设）
    waypoints.push_back(current_positions_);
    time_points.push_back(0.0);
    
    // 目标位置（示例）
    std::vector<double> home_positions = {0.0, -1.57, 1.57, 0.0, 0.0, 0.0};
    waypoints.push_back(home_positions);
    time_points.push_back(3.0);  // 3秒内到达
    
    // 执行轨迹
    executeTrajectory(waypoints, time_points);
  }

private:
  void initJointStates()
  {
    // 初始化关节状态
    joint_state_msg_.header.frame_id = "base_link";
    joint_state_msg_.name = joint_names_;
    
    // 初始化关节位置（全部为0）
    current_positions_.resize(joint_names_.size(), 0.0);
    joint_state_msg_.position = current_positions_;
    
    // 初始化速度和加速度
    joint_state_msg_.velocity.resize(joint_names_.size(), 0.0);
    joint_state_msg_.effort.resize(joint_names_.size(), 0.0);
  }
  
  void publishJointStates()
  {
    // 更新时间戳
    joint_state_msg_.header.stamp = this->now();
    
    // 发布关节状态
    joint_state_pub_->publish(joint_state_msg_);
  }
  
  // 关节状态发布者
  rclcpp::Publisher<sensor_msgs::msg::JointState>::SharedPtr joint_state_pub_;
  
  // 轨迹动作客户端
  rclcpp_action::Client<FollowJointTrajectory>::SharedPtr trajectory_client_;
  GoalHandleFollowJointTrajectory::SharedPtr goal_handle_;
  
  // 定时器
  rclcpp::TimerBase::SharedPtr timer_;
  
  // 关节状态消息
  sensor_msgs::msg::JointState joint_state_msg_;
  
  // 关节名称和当前位置
  std::vector<std::string> joint_names_;
  std::vector<double> current_positions_;
};

int main(int argc, char * argv[])
{
  rclcpp::init(argc, argv);
  auto node = std::make_shared<ArmController>();
  
  // 示例：移动到预设位置
  node->moveToHomePosition();
  
  rclcpp::spin(node);
  rclcpp::shutdown();
  return 0;
}