#include "dual_arm_moveit_action_server/dual_arm_moveit_action_server.hpp"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <thread>

namespace
{
struct BusyReset
{
  std::atomic_bool & flag;
  explicit BusyReset(std::atomic_bool & flag_in) : flag(flag_in) {}
  ~BusyReset() { flag.store(false); }
};

struct FeedbackReset
{
  std::shared_ptr<dual_arm_moveit_action_server::DualArm::Feedback> & feedback;
  std::mutex & mutex;
  FeedbackReset(
    std::shared_ptr<dual_arm_moveit_action_server::DualArm::Feedback> & feedback_in,
    std::mutex & mutex_in)
  : feedback(feedback_in), mutex(mutex_in)
  {
  }
  ~FeedbackReset()
  {
    std::lock_guard<std::mutex> lock(mutex);
    feedback.reset();
  }
};

struct ExecTask
{
  std::shared_ptr<moveit::planning_interface::MoveGroupInterface> move_group;
  moveit_msgs::msg::RobotTrajectory trajectory;
  std::vector<std::string> joint_names;
  std::vector<double> target_positions;
  std::chrono::steady_clock::time_point deadline;
  bool done{false};
};

double compute_trajectory_duration(const moveit_msgs::msg::RobotTrajectory & trajectory)
{
  if (trajectory.joint_trajectory.points.empty()) {
    return 0.0;
  }
  const auto & last_point = trajectory.joint_trajectory.points.back();
  return static_cast<double>(last_point.time_from_start.sec) +
    static_cast<double>(last_point.time_from_start.nanosec) * 1e-9;
}
}  // namespace

namespace dual_arm_moveit_action_server
{
namespace
{
double percent_to_scaling(int percent)
{
  const double clamped = std::clamp(percent, 1, 100);
  return clamped / 100.0;
}
}  // namespace


DualArmMoveitActionServer::DualArmMoveitActionServer(const rclcpp::Node::SharedPtr & node)
: node_(node)
{
}

void DualArmMoveitActionServer::initialize()
{
  left_group_name_ = node_->declare_parameter<std::string>("left_group_name", "LeftArm");
  right_group_name_ = node_->declare_parameter<std::string>("right_group_name", "RightArm");
  planning_time_ = node_->declare_parameter<double>("planning_time", 5.0);
  acceleration_scaling_ = node_->declare_parameter<double>("acceleration_scaling", 0.5);
  cartesian_step_ = node_->declare_parameter<double>("cartesian_step", 0.005);
  cartesian_jump_threshold_ = node_->declare_parameter<double>("cartesian_jump_threshold", 0.1);
  cartesian_min_fraction_ = node_->declare_parameter<double>("cartesian_min_fraction", 0.5);

  try {
    left_move_group_ =
      std::make_shared<moveit::planning_interface::MoveGroupInterface>(node_, left_group_name_);
    right_move_group_ =
      std::make_shared<moveit::planning_interface::MoveGroupInterface>(node_, right_group_name_);
  } catch (const std::exception & e) {
    RCLCPP_ERROR(node_->get_logger(), "Failed to create MoveGroupInterface: %s", e.what());
    left_move_group_.reset();
    right_move_group_.reset();
    return;
  }

  left_move_group_->setPlanningTime(planning_time_);
  right_move_group_->setPlanningTime(planning_time_);

  display_pub_ = node_->create_publisher<moveit_msgs::msg::DisplayTrajectory>(
    "display_planned_path", rclcpp::SystemDefaultsQoS());

  action_server_ = rclcpp_action::create_server<DualArm>(
    node_,
    "DualArm",
    std::bind(&DualArmMoveitActionServer::handle_goal, this, std::placeholders::_1,
      std::placeholders::_2),
    std::bind(&DualArmMoveitActionServer::handle_cancel, this, std::placeholders::_1),
    std::bind(&DualArmMoveitActionServer::handle_accepted, this, std::placeholders::_1));

  RCLCPP_INFO(node_->get_logger(), "DualArm MoveIt action server is ready");
}

rclcpp_action::GoalResponse DualArmMoveitActionServer::handle_goal(
  const rclcpp_action::GoalUUID & /*uuid*/,
  std::shared_ptr<const DualArm::Goal> goal)
{
  if (!left_move_group_ || !right_move_group_) {
    RCLCPP_ERROR(node_->get_logger(), "MoveGroupInterface not initialized");
    return rclcpp_action::GoalResponse::REJECT;
  }

  bool expected = false;
  if (!planning_busy_.compare_exchange_strong(expected, true)) {
    RCLCPP_WARN(node_->get_logger(), "DualArm rejected: planner is busy");
    return rclcpp_action::GoalResponse::REJECT;
  }

  std::string error_msg;
  if (!validate_goal(*goal, &error_msg)) {
    RCLCPP_ERROR(node_->get_logger(), "DualArm rejected: %s", error_msg.c_str());
    planning_busy_.store(false);
    return rclcpp_action::GoalResponse::REJECT;
  }

  return rclcpp_action::GoalResponse::ACCEPT_AND_EXECUTE;
}

rclcpp_action::CancelResponse DualArmMoveitActionServer::handle_cancel(
  const std::shared_ptr<GoalHandleDualArm> /*goal_handle*/)
{
  RCLCPP_INFO(node_->get_logger(), "DualArm cancel requested");
  cancel_requested_.store(true);
  if (left_move_group_) {
    left_move_group_->stop();
  }
  if (right_move_group_) {
    right_move_group_->stop();
  }
  return rclcpp_action::CancelResponse::ACCEPT;
}

void DualArmMoveitActionServer::handle_accepted(const std::shared_ptr<GoalHandleDualArm> goal_handle)
{
  std::thread{&DualArmMoveitActionServer::execute, this, goal_handle}.detach();
}

void DualArmMoveitActionServer::execute(const std::shared_ptr<GoalHandleDualArm> goal_handle)
{
  cancel_requested_.store(false);

  BusyReset busy_reset(planning_busy_);
  FeedbackReset feedback_reset(feedback_state_, feedback_mutex_);
  auto result = std::make_shared<DualArm::Result>();
  auto feedback = std::make_shared<DualArm::Feedback>();
  result->success = false;
  result->final_progress = 0.0;
  result->message = "DualArm planning failed";

  const auto goal = goal_handle->get_goal();
  const auto params_list = normalize_goal_params(*goal);
  const size_t total = params_list.size();
  const auto should_cancel = [&]() {
      return cancel_requested_.load() || goal_handle->is_canceling();
    };
  const auto stop_execution = [&]() {
      if (left_move_group_) {
        left_move_group_->stop();
      }
      if (right_move_group_) {
        right_move_group_->stop();
      }
    };

  update_feedback(DualArm::Goal::STATUS_PLANNING, 0.0);
  publish_feedback_once(goal_handle);

  std::vector<ExecTask> exec_tasks;
  std::vector<std::string> exec_errors(total);
  exec_tasks.reserve(total);

  for (size_t i = 0; i < total; ++i) {
    if (should_cancel()) {
      result->message = "DualArm canceled";
      result->success = false;
      result->final_progress = 0.0;
      goal_handle->canceled(result);
      return;
    }

    const auto & params = params_list[i];
    auto move_group = (params.arm_id == interfaces::msg::ArmMotionParams::ARM_LEFT) ?
      left_move_group_ :
      right_move_group_;

    std::string plan_error;
    moveit_msgs::msg::RobotTrajectory trajectory;
    if (!plan_for_arm(params, *move_group, goal->velocity_scaling, &trajectory, &plan_error)) {
      result->message = plan_error.empty() ? "MoveIt planning failed" : plan_error;
      goal_handle->abort(result);
      return;
    }

    publish_display_trajectory(trajectory, *move_group);

    ExecTask task;
    task.move_group = move_group;
    task.trajectory = trajectory;
    task.joint_names = trajectory.joint_trajectory.joint_names;
    if (!trajectory.joint_trajectory.points.empty()) {
      task.target_positions = trajectory.joint_trajectory.points.back().positions;
    }
    const double duration = compute_trajectory_duration(trajectory);
    const double timeout_sec = std::max(2.0, duration + 2.0);
    task.deadline = std::chrono::steady_clock::now() +
      std::chrono::milliseconds(static_cast<int>(timeout_sec * 1000.0));

    const auto exec_code = move_group->asyncExecute(trajectory);
    if (exec_code != moveit::core::MoveItErrorCode::SUCCESS) {
      result->message = "MoveIt asyncExecute failed";
      goal_handle->abort(result);
      return;
    }
    exec_tasks.push_back(std::move(task));
  }

  while (rclcpp::ok()) {
    if (should_cancel()) {
      stop_execution();
      result->message = "DualArm canceled";
      result->success = false;
      result->final_progress = 0.0;
      goal_handle->canceled(result);
      return;
    }

    size_t done_count = 0;
    for (size_t i = 0; i < exec_tasks.size(); ++i) {
      auto & task = exec_tasks[i];
      if (task.done) {
        ++done_count;
        continue;
      }

      auto current_state = task.move_group->getCurrentState();
      if (current_state && !task.joint_names.empty() &&
          task.joint_names.size() == task.target_positions.size())
      {
        double max_error = 0.0;
        for (size_t j = 0; j < task.joint_names.size(); ++j) {
          const double current = current_state->getVariablePosition(task.joint_names[j]);
          const double error = std::abs(current - task.target_positions[j]);
          if (error > max_error) {
            max_error = error;
          }
        }
        if (max_error <= 0.01) {
          task.done = true;
          ++done_count;
          continue;
        }
      }

      if (std::chrono::steady_clock::now() >= task.deadline) {
        exec_errors[i] = "MoveIt execution timeout";
        stop_execution();
        result->message = exec_errors[i];
        result->success = false;
        result->final_progress = 0.0;
        goal_handle->abort(result);
        return;
      }
    }

    feedback->progress = total > 0 ? static_cast<double>(done_count) / total : 0.0;
    feedback->status = DualArm::Goal::STATUS_EXECUTING;
    fill_feedback_joints(*feedback);
    goal_handle->publish_feedback(feedback);

    if (done_count == total) {
      break;
    }

    std::this_thread::sleep_for(std::chrono::milliseconds(50));
  }

  for (const auto & task : exec_tasks) {
    if (task.move_group) {
      publish_display_trajectory(task.trajectory, *task.move_group);
    }
  }

  bool all_success = true;
  for (const auto & error : exec_errors) {
    if (!error.empty()) {
      all_success = false;
      break;
    }
  }

  result->success = all_success;
  result->final_progress = all_success ? 1.0 : 0.0;
  if (all_success) {
    result->message = "DualArm planned and executed successfully";
    goal_handle->succeed(result);
  } else {
    for (const auto & error : exec_errors) {
      if (!error.empty()) {
        result->message = error;
        break;
      }
    }
    goal_handle->abort(result);
  }
}

bool DualArmMoveitActionServer::validate_goal(const DualArm::Goal & goal, std::string * error_msg) const
{
  if (goal.velocity_scaling < 1 || goal.velocity_scaling > 100) {
    if (error_msg) {
      *error_msg = "velocity_scaling out of range [1, 100]";
    }
    return false;
  }

  if (goal.arm_motion_params.empty() || goal.arm_motion_params.size() > 2) {
    if (error_msg) {
      *error_msg = "arm_motion_params must contain 1 or 2 items";
    }
    return false;
  }

  bool has_left = false;
  bool has_right = false;
  for (const auto & params : goal.arm_motion_params) {
    if (params.arm_id == interfaces::msg::ArmMotionParams::ARM_LEFT) {
      has_left = true;
    } else if (params.arm_id == interfaces::msg::ArmMotionParams::ARM_RIGHT) {
      has_right = true;
    } else {
      if (error_msg) {
        *error_msg = "invalid arm_id (expected 0=left, 1=right)";
      }
      return false;
    }

    if (params.motion_type != interfaces::msg::ArmMotionParams::MOVEJ &&
        params.motion_type != interfaces::msg::ArmMotionParams::MOVEL)
    {
      if (error_msg) {
        *error_msg = "unsupported motion_type";
      }
      return false;
    }
  }

  if (goal.arm_motion_params.size() == 2 && (!has_left || !has_right)) {
    if (error_msg) {
      *error_msg = "dual-arm mode requires both left and right arm_id";
    }
    return false;
  }

  return true;
}

std::vector<interfaces::msg::ArmMotionParams> DualArmMoveitActionServer::normalize_goal_params(
  const DualArm::Goal & goal) const
{
  std::vector<interfaces::msg::ArmMotionParams> params;
  params.reserve(goal.arm_motion_params.size());

  if (goal.arm_motion_params.size() == 1) {
    params.push_back(goal.arm_motion_params.front());
    return params;
  }

  interfaces::msg::ArmMotionParams left;
  interfaces::msg::ArmMotionParams right;
  bool has_left = false;
  bool has_right = false;

  for (const auto & item : goal.arm_motion_params) {
    if (item.arm_id == interfaces::msg::ArmMotionParams::ARM_LEFT) {
      left = item;
      has_left = true;
    } else if (item.arm_id == interfaces::msg::ArmMotionParams::ARM_RIGHT) {
      right = item;
      has_right = true;
    }
  }

  if (has_left) {
    params.push_back(left);
  }
  if (has_right) {
    params.push_back(right);
  }

  return params;
}

bool DualArmMoveitActionServer::plan_for_arm(
  const interfaces::msg::ArmMotionParams & params,
  moveit::planning_interface::MoveGroupInterface & move_group,
  int velocity_scaling,
  moveit_msgs::msg::RobotTrajectory * out_trajectory,
  std::string * error_msg)
{
  if (!out_trajectory) {
    if (error_msg) {
      *error_msg = "Null trajectory output";
    }
    return false;
  }

  const double velocity_factor = percent_to_scaling(velocity_scaling);
  move_group.setStartStateToCurrentState();
  move_group.setMaxVelocityScalingFactor(velocity_factor);
  move_group.setMaxAccelerationScalingFactor(acceleration_scaling_);

  if (params.motion_type == interfaces::msg::ArmMotionParams::MOVEJ) {
    const auto robot_model = move_group.getRobotModel();
    const auto * jmg = robot_model ? robot_model->getJointModelGroup(move_group.getName()) : nullptr;
    if (!jmg) {
      if (error_msg) {
        *error_msg = "MoveGroup has no JointModelGroup";
      }
      return false;
    }
    const size_t expected = jmg->getVariableCount();
    if (params.target_joints.size() != expected) {
      if (error_msg) {
        *error_msg =
          "MOVEJ target_joints size mismatch (expected " + std::to_string(expected) + ")";
      }
      return false;
    }

    move_group.setJointValueTarget(params.target_joints);

    moveit::planning_interface::MoveGroupInterface::Plan plan;
    auto code = move_group.plan(plan);
    move_group.clearPoseTargets();

    if (code != moveit::core::MoveItErrorCode::SUCCESS) {
      if (error_msg) {
        *error_msg = "MOVEJ planning failed";
      }
      return false;
    }

    if (!time_parameterize_trajectory(move_group, velocity_factor, plan.trajectory_, error_msg)) {
      return false;
    }

    *out_trajectory = plan.trajectory_;
    return true;
  }

  if (params.motion_type == interfaces::msg::ArmMotionParams::MOVEL) {
    if (params.arm_id == interfaces::msg::ArmMotionParams::ARM_LEFT) {
      move_group.setPoseReferenceFrame("base_link_leftarm");
    } else if (params.arm_id == interfaces::msg::ArmMotionParams::ARM_RIGHT) {
      move_group.setPoseReferenceFrame("base_link_rightarm");
    }
    RCLCPP_INFO(
      node_->get_logger(),
      "MOVEL frames: pose_reference_frame=%s planning_frame=%s",
      move_group.getPoseReferenceFrame().c_str(),
      move_group.getPlanningFrame().c_str());
    std::vector<geometry_msgs::msg::Pose> waypoints;
    waypoints.push_back(params.target_pose);

    moveit_msgs::msg::RobotTrajectory trajectory;
    const double fraction = move_group.computeCartesianPath(
      waypoints, cartesian_step_, cartesian_jump_threshold_, trajectory, false);

    move_group.clearPoseTargets();

    if (fraction < cartesian_min_fraction_) {
      if (error_msg) {
        *error_msg = "MOVEL cartesian path fraction too low: " + std::to_string(fraction);
      }
      return false;
    }

    if (!time_parameterize_trajectory(move_group, velocity_factor, trajectory, error_msg)) {
      return false;
    }

    *out_trajectory = trajectory;
    return true;
  }

  if (error_msg) {
    *error_msg = "Unsupported motion_type";
  }
  return false;
}

bool DualArmMoveitActionServer::time_parameterize_trajectory(
  moveit::planning_interface::MoveGroupInterface & move_group,
  double velocity_scaling,
  moveit_msgs::msg::RobotTrajectory & trajectory,
  std::string * error_msg)
{
  const auto current_state = move_group.getCurrentState();
  if (!current_state) {
    if (error_msg) {
      *error_msg = "Failed to get current state for time parameterization";
    }
    return false;
  }

  robot_trajectory::RobotTrajectory robot_traj(
    move_group.getRobotModel(), move_group.getName());
  robot_traj.setRobotTrajectoryMsg(*current_state, trajectory);

  trajectory_processing::IterativeParabolicTimeParameterization time_param;
  if (!time_param.computeTimeStamps(robot_traj, velocity_scaling, acceleration_scaling_)) {
    if (error_msg) {
      *error_msg = "Time parameterization failed";
    }
    return false;
  }

  robot_traj.getRobotTrajectoryMsg(trajectory);
  return true;
}

void DualArmMoveitActionServer::publish_display_trajectory(
  const moveit_msgs::msg::RobotTrajectory & trajectory,
  moveit::planning_interface::MoveGroupInterface & move_group)
{
  if (!display_pub_) {
    return;
  }

  auto current_state = move_group.getCurrentState();
  if (!current_state) {
    RCLCPP_WARN(node_->get_logger(), "No current state available for display");
    return;
  }

  moveit_msgs::msg::DisplayTrajectory display_msg;
  display_msg.model_id = move_group.getRobotModel()->getName();

  moveit_msgs::msg::RobotState start_state_msg;
  moveit::core::robotStateToRobotStateMsg(*current_state, start_state_msg);
  display_msg.trajectory_start = start_state_msg;
  display_msg.trajectory.push_back(trajectory);

  display_pub_->publish(display_msg);
}

void DualArmMoveitActionServer::fill_feedback_joints(DualArm::Feedback & feedback)
{
  feedback.joints_left.clear();
  feedback.joints_right.clear();

  if (left_move_group_) {
    auto state = left_move_group_->getCurrentState();
    if (state) {
      const auto * jmg = state->getJointModelGroup(left_group_name_);
      if (jmg) {
        state->copyJointGroupPositions(jmg, feedback.joints_left);
      }
    }
  }

  if (right_move_group_) {
    auto state = right_move_group_->getCurrentState();
    if (state) {
      const auto * jmg = state->getJointModelGroup(right_group_name_);
      if (jmg) {
        state->copyJointGroupPositions(jmg, feedback.joints_right);
      }
    }
  }
}

void DualArmMoveitActionServer::publish_feedback_once(
  const std::shared_ptr<GoalHandleDualArm> & goal_handle)
{
  if (!goal_handle) {
    return;
  }

  std::shared_ptr<DualArm::Feedback> feedback_copy;
  {
    std::lock_guard<std::mutex> lock(feedback_mutex_);
    if (!feedback_state_) {
      return;
    }
    feedback_copy = std::make_shared<DualArm::Feedback>(*feedback_state_);
  }

  fill_feedback_joints(*feedback_copy);
  goal_handle->publish_feedback(feedback_copy);
}

void DualArmMoveitActionServer::update_feedback(uint8_t status, double progress)
{
  std::lock_guard<std::mutex> lock(feedback_mutex_);
  if (!feedback_state_) {
    feedback_state_ = std::make_shared<DualArm::Feedback>();
  }
  feedback_state_->status = status;
  feedback_state_->progress = progress;
}

}  // namespace dual_arm_moveit_action_server

int main(int argc, char ** argv)
{
  rclcpp::init(argc, argv);
  auto node = rclcpp::Node::make_shared("dual_arm_moveit_action_server");
  dual_arm_moveit_action_server::DualArmMoveitActionServer server(node);
  server.initialize();
  rclcpp::spin(node);
  rclcpp::shutdown();
  return 0;
}
