#include "dual_arm_action_server/real_robot_driver.hpp"
#include <tf2/LinearMath/Matrix3x3.hpp>
#include <tf2/LinearMath/Quaternion.hpp>
#include <algorithm>
#include <cmath>
#include <cstring>

namespace dual_arm_action_server
{

RealRobotDriver * RealRobotDriver::instance_ = nullptr;

RealRobotDriver::RealRobotDriver()
: logger_(rclcpp::get_logger("dual_arm_action_server")),
  left_handle_(nullptr),
  right_handle_(nullptr),
  connected_(false),
  joints_left_(USED_ARM_DOF, 0.0),
  joints_right_(USED_ARM_DOF, 0.0)
{
}

RealRobotDriver::~RealRobotDriver()
{
  shutdown();
}

bool RealRobotDriver::initialize(
  const rclcpp::Logger & logger,
  const std::string & left_ip,
  const std::string & right_ip,
  int arm_port,
  const std::string & left_push_ip,
  const std::string & right_push_ip,
  int push_port,
  int push_cycle)
{
  logger_ = logger;
  left_ip_ = left_ip;
  right_ip_ = right_ip;

  rm_service_ = std::make_unique<RM_Service>();
  if (!rm_service_) {
    RCLCPP_ERROR(logger_, "Failed to allocate RM_Service");
    return false;
  }

  if (rm_service_->rm_init(RM_TRIPLE_MODE_E) != 0) {
    RCLCPP_ERROR(logger_, "rm_init failed");
    rm_service_.reset();
    return false;
  }

  instance_ = this;
  rm_service_->rm_set_log_call_back(api_log_callback, 3);
  rm_service_->rm_realtime_arm_state_call_back(realtime_state_callback);

  if (!init_handle(&left_handle_, left_ip, arm_port, left_push_ip, push_port, push_cycle)) {
    shutdown();
    return false;
  }

  if (!init_handle(&right_handle_, right_ip, arm_port, right_push_ip, push_port, push_cycle)) {
    shutdown();
    return false;
  }

  connected_.store(true);
  return true;
}

void RealRobotDriver::shutdown()
{
  connected_.store(false);
  if (rm_service_) {
    rm_service_->rm_destroy();
    rm_service_.reset();
  }
  left_handle_ = nullptr;
  right_handle_ = nullptr;
  instance_ = nullptr;
}

bool RealRobotDriver::is_connected() const
{
  return connected_.load();
}

bool RealRobotDriver::execute_motion(
  const interfaces::msg::ArmMotionParams & params,
  uint8_t arm_id,
  int velocity_scaling,
  int blend_radius,
  std::string * error_msg)
{
  const int blend_radius_percent = clamp_blend_radius_percent(blend_radius);
  const int trajectory_connect = (blend_radius_percent > 0) ? 1 : 0;
  const int block_mode = (blend_radius_percent > 0) ? 0 : 1;

  rm_robot_handle * handle = handle_for(arm_id);
  if (!handle) {
    if (error_msg) {
      *error_msg = "Invalid arm_id or missing robot handle";
    }
    return false;
  }

  const int speed_percent = clamp_speed_percent(velocity_scaling);

  if (params.motion_type == interfaces::msg::ArmMotionParams::MOVEJ) {
    if (params.target_joints.size() != USED_ARM_DOF) {
      if (error_msg) {
        *error_msg = "MOVEJ expects target_joints size " + std::to_string(USED_ARM_DOF);
      }
      return false;
    }

    float joints[ARM_DOF] = {0.0f};
    for (size_t i = 0; i < USED_ARM_DOF; ++i) {
      joints[i] = static_cast<float>(params.target_joints[i]);
    }

    const int result = rm_service_->rm_movej(
      handle,
      joints,
      speed_percent,
      blend_radius_percent,
      trajectory_connect,
      block_mode);
    if (result != 0) {
      if (error_msg) {
        *error_msg = "rm_movej failed with code " + std::to_string(result);
      }
      return false;
    }
    return true;
  }

  if (params.motion_type == interfaces::msg::ArmMotionParams::MOVEL) {
    rm_pose_t pose{};
    pose.position.x = static_cast<float>(params.target_pose.position.x);
    pose.position.y = static_cast<float>(params.target_pose.position.y);
    pose.position.z = static_cast<float>(params.target_pose.position.z);

    pose.quaternion.w = static_cast<float>(params.target_pose.orientation.w);
    pose.quaternion.x = static_cast<float>(params.target_pose.orientation.x);
    pose.quaternion.y = static_cast<float>(params.target_pose.orientation.y);
    pose.quaternion.z = static_cast<float>(params.target_pose.orientation.z);

    tf2::Quaternion q(
      pose.quaternion.x,
      pose.quaternion.y,
      pose.quaternion.z,
      pose.quaternion.w);
    tf2::Matrix3x3 m(q);
    double roll = 0.0;
    double pitch = 0.0;
    double yaw = 0.0;
    m.getRPY(roll, pitch, yaw);
    pose.euler.rx = static_cast<float>(roll);
    pose.euler.ry = static_cast<float>(pitch);
    pose.euler.rz = static_cast<float>(yaw);

    const int result = rm_service_->rm_movel(
      handle,
      pose,
      speed_percent,
      blend_radius_percent,
      trajectory_connect,
      block_mode);
    if (result != 0) {
      if (error_msg) {
        *error_msg = "rm_movel failed with code " + std::to_string(result);
      }
      return false;
    }
    return true;
  }

  if (error_msg) {
    *error_msg = "Unsupported motion_type " + std::to_string(params.motion_type);
  }
  return false;
}

bool RealRobotDriver::get_current_joints(
  std::vector<double> & joints_left,
  std::vector<double> & joints_right)
{
  std::lock_guard<std::mutex> lock(joints_mutex_);
  joints_left.resize(USED_ARM_DOF);
  joints_right.resize(USED_ARM_DOF);
  for (size_t i = 0; i < USED_ARM_DOF; ++i) {
    joints_left[i] = joints_left_[i];
    joints_right[i] = joints_right_[i];
  }
  return true;
}

int RealRobotDriver::clear_joint_error(uint8_t arm_id, int joint_num)
{
  rm_robot_handle * handle = handle_for(arm_id);
  if (!handle || !rm_service_) {
    return UNKNOWN_ERR;
  }

  const int result = rm_service_->rm_set_joint_clear_err(handle, joint_num);
  if (result == 0) {
    rm_service_->rm_set_joint_en_state(handle, joint_num, 1);
  }
  return result;
}

int RealRobotDriver::get_joint_error_and_brake(
  uint8_t arm_id,
  std::vector<uint16_t> & err_flag,
  std::vector<uint16_t> & brake_state)
{
  rm_robot_handle * handle = handle_for(arm_id);
  if (!handle || !rm_service_) {
    return UNKNOWN_ERR;
  }

  err_flag.assign(ARM_DOF, 0);
  brake_state.assign(ARM_DOF, 0);

  return rm_service_->rm_get_joint_err_flag(handle, err_flag.data(), brake_state.data());
}

int RealRobotDriver::inverse_kinematics(
  uint8_t arm_id,
  const rm_pose_t & pose,
  std::vector<double> & joint_angles)
{
  rm_robot_handle * handle = handle_for(arm_id);
  if (!handle || !rm_service_) {
    return UNKNOWN_ERR;
  }

  rm_inverse_kinematics_params_t params{};
  params.q_pose = pose;
  params.flag = 0;

  {
    std::lock_guard<std::mutex> lock(joints_mutex_);
    const auto & source = (arm_id == LEFT_ARM) ? joints_left_ : joints_right_;
    for (size_t i = 0; i < USED_ARM_DOF && i < source.size(); ++i) {
      params.q_in[i] = static_cast<float>(source[i]);
    }
  }

  float q_out[ARM_DOF] = {0.0f};
  const int result = rm_service_->rm_algo_inverse_kinematics(handle, params, q_out);

  joint_angles.clear();
  if (result == 0) {
    joint_angles.reserve(USED_ARM_DOF);
    for (size_t i = 0; i < USED_ARM_DOF; ++i) {
      joint_angles.push_back(static_cast<double>(q_out[i]));
    }
  }

  return result;
}

void RealRobotDriver::stop_arm(uint8_t arm_id)
{
  rm_robot_handle * handle = handle_for(arm_id);
  if (handle && rm_service_) {
    rm_service_->rm_set_arm_stop(handle);
  }
}

void RealRobotDriver::stop_all()
{
  stop_arm(LEFT_ARM);
  stop_arm(RIGHT_ARM);
}

void RealRobotDriver::realtime_state_callback(rm_realtime_arm_joint_state_t data)
{
  if (instance_) {
    instance_->update_joints(data);
  }
}

void RealRobotDriver::api_log_callback(const char * message, va_list args)
{
  if (!instance_ || !message) {
    return;
  }
  char buffer[1024];
  vsnprintf(buffer, sizeof(buffer), message, args);
  RCLCPP_INFO(instance_->logger_, "%s", buffer);
}

void RealRobotDriver::update_joints(const rm_realtime_arm_joint_state_t & data)
{
  std::lock_guard<std::mutex> lock(joints_mutex_);
  if (std::strcmp(data.arm_ip, left_ip_.c_str()) == 0) {
    for (size_t i = 0; i < USED_ARM_DOF; ++i) {
      joints_left_[i] = data.joint_status.joint_position[i];
    }
    return;
  }
  if (std::strcmp(data.arm_ip, right_ip_.c_str()) == 0) {
    for (size_t i = 0; i < USED_ARM_DOF; ++i) {
      joints_right_[i] = data.joint_status.joint_position[i];
    }
  }
}

bool RealRobotDriver::init_handle(
  rm_robot_handle ** handle,
  const std::string & ip,
  int port,
  const std::string & push_ip,
  int push_port,
  int push_cycle)
{
  if (!rm_service_) {
    return false;
  }

  *handle = rm_service_->rm_create_robot_arm(ip.c_str(), port);
  if (*handle == nullptr || (*handle)->id <= 0) {
    RCLCPP_ERROR(logger_, "Failed to create handle for arm %s", ip.c_str());
    return false;
  }

  rm_realtime_push_config_t config{};
  config.cycle = push_cycle;
  config.enable = true;
  config.port = push_port;
  config.force_coordinate = 0;
  std::snprintf(config.ip, sizeof(config.ip), "%s", push_ip.c_str());

  const int push_result = rm_service_->rm_set_realtime_push(*handle, config);
  if (push_result != 0) {
    RCLCPP_ERROR(
      logger_,
      "Failed to set realtime push for arm %s (code %d)",
      ip.c_str(),
      push_result);
    return false;
  }

  rm_robot_info_t robot_info{};
  const int info_result = rm_service_->rm_get_robot_info(*handle, &robot_info);
  if (info_result != 0 || robot_info.arm_dof != USED_ARM_DOF) {
    RCLCPP_ERROR(
      logger_,
      "Invalid robot info for arm %s (code %d, dof %d)",
      ip.c_str(),
      info_result,
      robot_info.arm_dof);
    return false;
  }

  return true;
}

rm_robot_handle * RealRobotDriver::handle_for(uint8_t arm_id) const
{
  if (arm_id == LEFT_ARM) {
    return left_handle_;
  }
  if (arm_id == RIGHT_ARM) {
    return right_handle_;
  }
  return nullptr;
}

int RealRobotDriver::clamp_speed_percent(int percent)
{
  return std::clamp(percent, 1, 100);
}

int RealRobotDriver::clamp_blend_radius_percent(int percent)
{
  if (percent <= 0) {
    return 0;
  }
  return std::clamp(percent, 0, 100);
}

}  // namespace dual_arm_action_server

