#include "dual_arm_action_server/dual_arm_action_server.hpp"
#include "arm_action_define.h"
#include <chrono>
#include <future>
#include <thread>

namespace dual_arm_action_server
{

DualArmActionServer::DualArmActionServer(const rclcpp::Node::SharedPtr & node)
: node_(node),
  left_arm_joint_start_index_(9),
  right_arm_joint_start_index_(15),
  driver_(),
  executing_left_(false),
  executing_right_(false)
{
}

void DualArmActionServer::initialize()
{
  node_->declare_parameter<std::string>("left_arm_ip");
  node_->declare_parameter<std::string>("right_arm_ip");
  node_->declare_parameter<int>("arm_port");
  node_->declare_parameter<std::string>("left_push_ip");
  node_->declare_parameter<std::string>("right_push_ip");
  node_->declare_parameter<int>("push_port");
  node_->declare_parameter<int>("push_cycle");
  node_->declare_parameter<std::vector<std::string>>("joint_names", std::vector<std::string>());
  node_->declare_parameter<int>("left_arm_joint_start_index", 9);
  node_->declare_parameter<int>("right_arm_joint_start_index", 15);

  std::string left_arm_ip;
  std::string right_arm_ip;
  std::string left_push_ip;
  std::string right_push_ip;
  int arm_port = 0;
  int push_port = 0;
  int push_cycle = 0;

  const bool params_ok =
    node_->get_parameter("left_arm_ip", left_arm_ip) &&
    node_->get_parameter("right_arm_ip", right_arm_ip) &&
    node_->get_parameter("arm_port", arm_port) &&
    node_->get_parameter("left_push_ip", left_push_ip) &&
    node_->get_parameter("right_push_ip", right_push_ip) &&
    node_->get_parameter("push_port", push_port) &&
    node_->get_parameter("push_cycle", push_cycle);

  if (!params_ok ||
      left_arm_ip.empty() ||
      right_arm_ip.empty() ||
      left_push_ip.empty() ||
      right_push_ip.empty() ||
      arm_port <= 0 ||
      push_port <= 0 ||
      push_cycle <= 0)
  {
    RCLCPP_ERROR(node_->get_logger(),
      "Missing/invalid parameters. Please load config/dual_arm_action_server.yaml");
    return;
  }

  if (!driver_.initialize(
      node_->get_logger(),
      left_arm_ip,
      right_arm_ip,
      arm_port,
      left_push_ip,
      right_push_ip,
      push_port,
      push_cycle))
  {
    RCLCPP_ERROR(node_->get_logger(), "Failed to initialize RealRobotDriver");
  }

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

  clear_error_service_ = node_->create_service<interfaces::srv::ClearArmError>(
    "clear_arm_error",
    std::bind(
      &DualArmActionServer::clear_arm_error_callback,
      this,
      std::placeholders::_1,
      std::placeholders::_2));

  ik_service_ = node_->create_service<interfaces::srv::InverseKinematics>(
    "inverse_kinematics",
    std::bind(
      &DualArmActionServer::ik_service_callback,
      this,
      std::placeholders::_1,
      std::placeholders::_2));

  arm_error_pub_ = node_->create_publisher<interfaces::msg::ArmError>("arm_errors", 10);

  error_timer_ = node_->create_wall_timer(
    std::chrono::milliseconds(500),
    std::bind(&DualArmActionServer::error_timer_callback, this));

  joint_state_msg_.name = node_->get_parameter("joint_names").as_string_array();
  if (joint_state_msg_.name.empty()) {
    RCLCPP_ERROR(node_->get_logger(), "joint_names is empty, joint_states will not be published");
  } else {
    left_arm_joint_start_index_ = static_cast<size_t>(
      node_->get_parameter("left_arm_joint_start_index").as_int());
    right_arm_joint_start_index_ = static_cast<size_t>(
      node_->get_parameter("right_arm_joint_start_index").as_int());

    joint_state_pub_ = node_->create_publisher<sensor_msgs::msg::JointState>("joint_states", 10);
    joint_state_msg_.position.resize(joint_state_msg_.name.size(), 0.0);
    joint_state_msg_.header.frame_id = "base_link";

    joint_state_timer_ = node_->create_wall_timer(
      std::chrono::milliseconds(50),
      std::bind(&DualArmActionServer::joint_state_timer_callback, this));
  }

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

rclcpp_action::GoalResponse DualArmActionServer::handle_goal(
  const rclcpp_action::GoalUUID & uuid,
  std::shared_ptr<const DualArm::Goal> goal)
{
  (void)uuid;

  if (!driver_.is_connected()) {
    RCLCPP_ERROR(node_->get_logger(), "DualArm rejected: real robot not connected");
    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());
    return rclcpp_action::GoalResponse::REJECT;
  }

  bool need_left = false;
  bool need_right = false;
  get_target_arms(*goal, need_left, need_right);

  {
    std::lock_guard<std::mutex> lock(exec_state_mutex_);
    if ((need_left && executing_left_.load()) || (need_right && executing_right_.load())) {
      RCLCPP_WARN(node_->get_logger(), "DualArm rejected: target arm is executing");
      return rclcpp_action::GoalResponse::REJECT;
    }
    if (need_left) {
      executing_left_.store(true);
    }
    if (need_right) {
      executing_right_.store(true);
    }
  }

  return rclcpp_action::GoalResponse::ACCEPT_AND_EXECUTE;
}

rclcpp_action::CancelResponse DualArmActionServer::handle_cancel(
  const std::shared_ptr<GoalHandleDualArm> goal_handle)
{
  (void)goal_handle;
  bool need_left = false;
  bool need_right = false;
  get_target_arms(*goal_handle->get_goal(), need_left, need_right);
  if (need_left) {
    driver_.stop_arm(LEFT_ARM);
  }
  if (need_right) {
    driver_.stop_arm(RIGHT_ARM);
  }
  RCLCPP_INFO(node_->get_logger(), "DualArm cancel requested");
  return rclcpp_action::CancelResponse::ACCEPT;
}

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

void DualArmActionServer::execute(const std::shared_ptr<GoalHandleDualArm> goal_handle)
{
  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 action failed";

  const auto goal = goal_handle->get_goal();
  const auto params_list = normalize_goal_params(*goal);
  const size_t total = params_list.size();
  bool need_left = false;
  bool need_right = false;
  get_target_arms(*goal, need_left, need_right);

  std::vector<std::string> errors(total);
  std::vector<std::future<bool>> futures;
  futures.reserve(total);

  for (size_t i = 0; i < total; ++i) {
    const auto params = params_list[i];
    futures.emplace_back(std::async(std::launch::async, [this, params, goal, &errors, i]() {
      return driver_.execute_motion(
        params,
        params.arm_id,
        goal->velocity_scaling,
        params.blend_radius,
        &errors[i]);
    }));
  }

  while (rclcpp::ok()) {
    if (goal_handle->is_canceling()) {
      if (need_left) {
        driver_.stop_arm(LEFT_ARM);
      }
      if (need_right) {
        driver_.stop_arm(RIGHT_ARM);
      }
      result->message = "DualArm canceled";
      result->success = false;
      result->final_progress = 0.0;
      goal_handle->canceled(result);
      {
        std::lock_guard<std::mutex> lock(exec_state_mutex_);
        if (need_left) {
          executing_left_.store(false);
        }
        if (need_right) {
          executing_right_.store(false);
        }
      }
      return;
    }

    size_t done_count = 0;
    for (auto & future : futures) {
      if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) {
        ++done_count;
      }
    }

    std::vector<double> joints_left;
    std::vector<double> joints_right;
    driver_.get_current_joints(joints_left, joints_right);
    feedback->joints_left = joints_left;
    feedback->joints_right = joints_right;
    feedback->progress = total > 0 ? static_cast<double>(done_count) / total : 0.0;
    feedback->status = DualArm::Goal::STATUS_EXECUTING;
    goal_handle->publish_feedback(feedback);

    if (done_count == total) {
      break;
    }

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

  bool all_success = true;
  for (size_t i = 0; i < total; ++i) {
    bool ok = futures[i].get();
    if (!ok) {
      all_success = false;
      if (!errors[i].empty()) {
        result->message = errors[i];
      }
    }
  }

  result->success = all_success;
  result->final_progress = all_success ? 1.0 : 0.0;
  if (all_success) {
    result->message = "DualArm succeeded";
    goal_handle->succeed(result);
  } else {
    goal_handle->abort(result);
  }

  {
    std::lock_guard<std::mutex> lock(exec_state_mutex_);
    if (need_left) {
      executing_left_.store(false);
    }
    if (need_right) {
      executing_right_.store(false);
    }
  }
}

bool DualArmActionServer::validate_goal(const DualArm::Goal & goal, std::string & error_msg) const
{
  if (goal.velocity_scaling < 1 || goal.velocity_scaling > 100) {
    error_msg = "velocity_scaling out of range [1, 100]";
    return false;
  }
  if (goal.arm_motion_params.empty() || goal.arm_motion_params.size() > 2) {
    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.motion_type != interfaces::msg::ArmMotionParams::MOVEJ &&
        params.motion_type != interfaces::msg::ArmMotionParams::MOVEL)
    {
      error_msg = "unsupported motion_type";
      return false;
    }
    if (params.blend_radius < 0 || params.blend_radius > 100) {
      error_msg = "blend_radius out of range [0, 100]";
      return false;
    }
    if (params.motion_type == interfaces::msg::ArmMotionParams::MOVEJ &&
        params.target_joints.size() != USED_ARM_DOF)
    {
      error_msg = "MOVEJ expects target_joints size " + std::to_string(USED_ARM_DOF);
      return false;
    }
    if (params.arm_id == LEFT_ARM) {
      has_left = true;
    } else if (params.arm_id == RIGHT_ARM) {
      has_right = true;
    } else {
      error_msg = "invalid arm_id (expected 0=left, 1=right)";
      return false;
    }
  }

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

  return true;
}

std::vector<interfaces::msg::ArmMotionParams> DualArmActionServer::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;
  for (const auto & item : goal.arm_motion_params) {
    if (item.arm_id == LEFT_ARM) {
      left = item;
    } else if (item.arm_id == RIGHT_ARM) {
      right = item;
    }
  }
  if (left.arm_id == LEFT_ARM) {
    params.push_back(left);
  }
  if (right.arm_id == RIGHT_ARM) {
    params.push_back(right);
  }
  return params;
}

void DualArmActionServer::get_target_arms(
  const DualArm::Goal & goal,
  bool & need_left,
  bool & need_right) const
{
  need_left = false;
  need_right = false;
  for (const auto & params : goal.arm_motion_params) {
    if (params.arm_id == LEFT_ARM) {
      need_left = true;
    } else if (params.arm_id == RIGHT_ARM) {
      need_right = true;
    }
  }
}

void DualArmActionServer::clear_arm_error_callback(
  const std::shared_ptr<interfaces::srv::ClearArmError::Request> request,
  std::shared_ptr<interfaces::srv::ClearArmError::Response> response)
{
  if (!driver_.is_connected()) {
    response->success = false;
    response->message = "Real robot not connected";
    RCLCPP_ERROR(node_->get_logger(), "ClearArmError failed: real robot not connected");
    return;
  }

  if (request->arm_id != LEFT_ARM && request->arm_id != RIGHT_ARM) {
    response->success = false;
    response->message = "Invalid arm_id (expected 0=left, 1=right)";
    RCLCPP_ERROR(node_->get_logger(), "ClearArmError failed: invalid arm_id %d", request->arm_id);
    return;
  }

  RCLCPP_INFO(
    node_->get_logger(),
    "Clearing error for arm %d, joint %d",
    request->arm_id,
    request->joint_num);

  const int result = driver_.clear_joint_error(
    static_cast<uint8_t>(request->arm_id),
    request->joint_num);
  if (result == 0) {
    response->success = true;
    response->message = "Clear error success";
    RCLCPP_INFO(node_->get_logger(), "Clear error success");
  } else {
    response->success = false;
    response->message = "Clear error failed with code: " + std::to_string(result);
    RCLCPP_ERROR(node_->get_logger(), "Clear error failed with code: %d", result);
  }
}

void DualArmActionServer::ik_service_callback(
  const std::shared_ptr<interfaces::srv::InverseKinematics::Request> request,
  std::shared_ptr<interfaces::srv::InverseKinematics::Response> response)
{
  if (!driver_.is_connected()) {
    response->result = UNKNOWN_ERR;
    RCLCPP_ERROR(node_->get_logger(), "InverseKinematics failed: real robot not connected");
    return;
  }

  if (request->arm_id != LEFT_ARM && request->arm_id != RIGHT_ARM) {
    response->result = UNKNOWN_ERR;
    RCLCPP_ERROR(node_->get_logger(), "InverseKinematics failed: invalid arm_id %d", request->arm_id);
    return;
  }

  rm_pose_t pose{};
  pose.position.x = static_cast<float>(request->pose.position.x);
  pose.position.y = static_cast<float>(request->pose.position.y);
  pose.position.z = static_cast<float>(request->pose.position.z);
  pose.quaternion.x = static_cast<float>(request->pose.orientation.x);
  pose.quaternion.y = static_cast<float>(request->pose.orientation.y);
  pose.quaternion.z = static_cast<float>(request->pose.orientation.z);
  pose.quaternion.w = static_cast<float>(request->pose.orientation.w);

  std::vector<double> joint_angles;
  const int result = driver_.inverse_kinematics(
    static_cast<uint8_t>(request->arm_id),
    pose,
    joint_angles);

  response->result = result;
  if (result == 0) {
    response->joint_angles.clear();
    response->joint_angles.reserve(joint_angles.size());
    for (const auto angle : joint_angles) {
      response->joint_angles.push_back(static_cast<float>(angle));
    }
  }
}

void DualArmActionServer::error_timer_callback()
{
  if (!driver_.is_connected()) {
    return;
  }

  auto publish_error = [this](uint8_t arm_id) {
    std::vector<uint16_t> err_flag;
    std::vector<uint16_t> brake_state;
    const int result = driver_.get_joint_error_and_brake(arm_id, err_flag, brake_state);
    if (result != 0) {
      return;
    }

    interfaces::msg::ArmError msg;
    msg.arm_id = static_cast<int8_t>(arm_id);
    msg.err_flag.reserve(err_flag.size());
    msg.brake_state.reserve(brake_state.size());
    for (const auto value : err_flag) {
      msg.err_flag.push_back(value);
    }
    for (const auto value : brake_state) {
      msg.brake_state.push_back(value);
    }
    arm_error_pub_->publish(msg);
  };

  publish_error(LEFT_ARM);
  publish_error(RIGHT_ARM);
}

void DualArmActionServer::joint_state_timer_callback()
{
  if (!joint_state_pub_ || !driver_.is_connected()) {
    return;
  }

  constexpr double kDegToRad = 3.14159265358979323846 / 180.0;

  std::vector<double> joints_left;
  std::vector<double> joints_right;
  if (!driver_.get_current_joints(joints_left, joints_right)) {
    return;
  }

  if (joint_state_msg_.position.size() >= right_arm_joint_start_index_ + USED_ARM_DOF) {
    for (size_t i = 0; i < USED_ARM_DOF && i < joints_left.size(); ++i) {
      joint_state_msg_.position[left_arm_joint_start_index_ + i] = joints_left[i] * kDegToRad;
    }
    for (size_t i = 0; i < USED_ARM_DOF && i < joints_right.size(); ++i) {
      joint_state_msg_.position[right_arm_joint_start_index_ + i] = joints_right[i] * kDegToRad;
    }
  }

  joint_state_msg_.header.stamp = node_->get_clock()->now();
  joint_state_msg_.velocity.clear();
  joint_state_msg_.effort.clear();
  joint_state_pub_->publish(joint_state_msg_);
}

}  // namespace dual_arm_action_server

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

