#include "dual_arm_remote_teleop/real_robot_driver.hpp"
#include "interfaces/msg/meta_key.hpp"
#include <geometry_msgs/msg/pose.hpp>
#include <rclcpp/rclcpp.hpp>
#include <tf2/LinearMath/Matrix3x3.hpp>
#include <tf2/LinearMath/Quaternion.hpp>
#include <algorithm>
#include <chrono>
#include <deque>
#include <string>

namespace dual_arm_remote_teleop
{

namespace
{
constexpr uint8_t kLeftArmId = 0;
constexpr uint8_t kRightArmId = 1;
}  // namespace

class RemoteTeleopNode : public rclcpp::Node
{
public:
  RemoteTeleopNode()
  : rclcpp::Node("dual_arm_remote_teleop")
  {
    declare_parameter<std::string>("left_arm_ip");
    declare_parameter<std::string>("right_arm_ip");
    declare_parameter<int>("arm_port");
    declare_parameter<std::string>("left_push_ip");
    declare_parameter<std::string>("right_push_ip");
    declare_parameter<int>("push_port");
    declare_parameter<int>("push_cycle");
    declare_parameter<std::string>("meta_topic", "/meta_key");
    declare_parameter<double>("command_rate_hz", 250.0);
    declare_parameter<int>("interpolation_delay_ms", 30);
    declare_parameter<double>("smooth_alpha", 0.3);
    declare_parameter<int>("stale_timeout_ms", 500);
    declare_parameter<bool>("require_trigger", false);
    declare_parameter<int>("trigger_threshold", 1);
    declare_parameter<bool>("canfd_follow", true);
    declare_parameter<int>("canfd_trajectory_mode", 2);
    declare_parameter<int>("canfd_smooth_ratio", 200);

    load_parameters();

    if (!driver_.initialize(
        get_logger(),
        left_arm_ip_,
        right_arm_ip_,
        arm_port_,
        left_push_ip_,
        right_push_ip_,
        push_port_,
        push_cycle_))
    {
      RCLCPP_ERROR(get_logger(), "Failed to initialize RM driver");
    }

    meta_sub_ = create_subscription<interfaces::msg::MetaKey>(
      meta_topic_,
      10,
      std::bind(&RemoteTeleopNode::on_meta_key, this, std::placeholders::_1));

    const auto period = std::chrono::duration<double>(1.0 / command_rate_hz_);
    timer_ = create_wall_timer(
      std::chrono::duration_cast<std::chrono::nanoseconds>(period),
      std::bind(&RemoteTeleopNode::on_timer, this));
  }

private:
  struct TeleopState {
    struct Sample {
      geometry_msgs::msg::Pose pose;
      rclcpp::Time stamp;
    };
    std::deque<Sample> samples;
    geometry_msgs::msg::Pose pose;
    rclcpp::Time stamp;
    bool active{false};
    bool has_pose{false};
    geometry_msgs::msg::Pose smoothed_pose;
    bool has_smoothed{false};
  };

  void load_parameters()
  {
    if (!get_parameter("left_arm_ip", left_arm_ip_) ||
      !get_parameter("right_arm_ip", right_arm_ip_) ||
      !get_parameter("arm_port", arm_port_) ||
      !get_parameter("left_push_ip", left_push_ip_) ||
      !get_parameter("right_push_ip", right_push_ip_) ||
      !get_parameter("push_port", push_port_) ||
      !get_parameter("push_cycle", push_cycle_))
    {
      RCLCPP_ERROR(get_logger(),
        "Missing/invalid parameters. Please load config/dual_arm_remote_teleop.yaml");
    }

    get_parameter("meta_topic", meta_topic_);
    get_parameter("command_rate_hz", command_rate_hz_);
    get_parameter("interpolation_delay_ms", interpolation_delay_ms_);
    get_parameter("smooth_alpha", smooth_alpha_);
    get_parameter("stale_timeout_ms", stale_timeout_ms_);
    get_parameter("require_trigger", require_trigger_);
    get_parameter("trigger_threshold", trigger_threshold_);
    get_parameter("canfd_follow", canfd_follow_);
    get_parameter("canfd_trajectory_mode", canfd_trajectory_mode_);
    get_parameter("canfd_smooth_ratio", canfd_smooth_ratio_);

    if (command_rate_hz_ <= 0.0) {
      command_rate_hz_ = 10.0;
    }
    if (smooth_alpha_ <= 0.0) {
      smooth_alpha_ = 0.1;
    } else if (smooth_alpha_ > 1.0) {
      smooth_alpha_ = 1.0;
    }
  }

  void on_meta_key(const interfaces::msg::MetaKey::SharedPtr msg)
  {
    const bool active = !require_trigger_ || msg->trigger >= trigger_threshold_;
    if (msg->type == "left") {
      update_state(left_state_, msg->pose, active);
      return;
    }
    if (msg->type == "right") {
      update_state(right_state_, msg->pose, active);
      return;
    }
  }

  void update_state(TeleopState & state, const geometry_msgs::msg::Pose & pose, bool active)
  {
    state.active = active;
    state.pose = pose;
    state.stamp = now();
    state.has_pose = true;
    state.samples.push_back({pose, state.stamp});
    while (state.samples.size() > 4) {
      state.samples.pop_front();
    }
  }

  void on_timer()
  {
    if (!driver_.is_connected()) {
      return;
    }

    const auto now_time = now();
    const auto timeout = rclcpp::Duration::from_nanoseconds(
      static_cast<int64_t>(stale_timeout_ms_) * 1000000LL);

    const auto target_time = now_time - rclcpp::Duration::from_nanoseconds(
      static_cast<int64_t>(interpolation_delay_ms_) * 1000000LL);
    send_if_ready(kLeftArmId, left_state_, now_time, timeout, target_time);
    send_if_ready(kRightArmId, right_state_, now_time, timeout, target_time);
  }

  void send_if_ready(
    uint8_t arm_id,
    TeleopState & state,
    const rclcpp::Time & now_time,
    const rclcpp::Duration & timeout,
    const rclcpp::Time & target_time)
  {
    if (!state.has_pose || !state.active) {
      return;
    }
    if ((now_time - state.stamp) > timeout) {
      return;
    }

    geometry_msgs::msg::Pose target_pose;
    if (!interpolate_pose(state, target_time, &target_pose)) {
      return;
    }

    geometry_msgs::msg::Pose smoothed = apply_smoothing(state, target_pose);
    state.smoothed_pose = smoothed;
    state.has_smoothed = true;
    rm_pose_t target = to_rm_pose(smoothed);
    std::string error_msg;
    if (!driver_.send_pose_canfd(
        arm_id,
        target,
        canfd_follow_,
        canfd_trajectory_mode_,
        canfd_smooth_ratio_,
        &error_msg))
    {
      RCLCPP_WARN_THROTTLE(get_logger(), *get_clock(), 2000, "Teleop send failed: %s",
        error_msg.c_str());
    }
  }

  bool interpolate_pose(
    const TeleopState & state,
    const rclcpp::Time & target_time,
    geometry_msgs::msg::Pose * out) const
  {
    if (state.samples.empty()) {
      return false;
    }

    const auto & first = state.samples.front();
    const auto & last = state.samples.back();
    if (target_time <= first.stamp || state.samples.size() == 1) {
      *out = first.pose;
      return true;
    }
    if (target_time >= last.stamp) {
      *out = last.pose;
      return true;
    }

    for (size_t i = 1; i < state.samples.size(); ++i) {
      const auto & prev = state.samples[i - 1];
      const auto & next = state.samples[i];
      if (target_time <= next.stamp) {
        const double dt = (next.stamp - prev.stamp).seconds();
        double t = 0.0;
        if (dt > 1e-6) {
          t = (target_time - prev.stamp).seconds() / dt;
        }
        t = std::clamp(t, 0.0, 1.0);

        out->position.x = prev.pose.position.x + (next.pose.position.x - prev.pose.position.x) * t;
        out->position.y = prev.pose.position.y + (next.pose.position.y - prev.pose.position.y) * t;
        out->position.z = prev.pose.position.z + (next.pose.position.z - prev.pose.position.z) * t;

        tf2::Quaternion q_prev(
          prev.pose.orientation.x,
          prev.pose.orientation.y,
          prev.pose.orientation.z,
          prev.pose.orientation.w);
        tf2::Quaternion q_next(
          next.pose.orientation.x,
          next.pose.orientation.y,
          next.pose.orientation.z,
          next.pose.orientation.w);
        tf2::Quaternion q_interp = q_prev.slerp(q_next, t);
        q_interp.normalize();
        out->orientation.x = q_interp.x();
        out->orientation.y = q_interp.y();
        out->orientation.z = q_interp.z();
        out->orientation.w = q_interp.w();
        return true;
      }
    }

    *out = last.pose;
    return true;
  }

  geometry_msgs::msg::Pose apply_smoothing(
    const TeleopState & state,
    const geometry_msgs::msg::Pose & input) const
  {
    if (!state.has_smoothed) {
      return input;
    }

    geometry_msgs::msg::Pose output;
    const double a = smooth_alpha_;
    output.position.x = state.smoothed_pose.position.x +
      (input.position.x - state.smoothed_pose.position.x) * a;
    output.position.y = state.smoothed_pose.position.y +
      (input.position.y - state.smoothed_pose.position.y) * a;
    output.position.z = state.smoothed_pose.position.z +
      (input.position.z - state.smoothed_pose.position.z) * a;

    tf2::Quaternion q_prev(
      state.smoothed_pose.orientation.x,
      state.smoothed_pose.orientation.y,
      state.smoothed_pose.orientation.z,
      state.smoothed_pose.orientation.w);
    tf2::Quaternion q_input(
      input.orientation.x,
      input.orientation.y,
      input.orientation.z,
      input.orientation.w);
    tf2::Quaternion q_out = q_prev.slerp(q_input, a);
    q_out.normalize();
    output.orientation.x = q_out.x();
    output.orientation.y = q_out.y();
    output.orientation.z = q_out.z();
    output.orientation.w = q_out.w();
    return output;
  }

  rm_pose_t to_rm_pose(const geometry_msgs::msg::Pose & pose) const
  {
    rm_pose_t target{};
    target.position.x = static_cast<float>(pose.position.x);
    target.position.y = static_cast<float>(pose.position.y);
    target.position.z = static_cast<float>(pose.position.z);

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

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

  RealRobotDriver driver_;
  rclcpp::Subscription<interfaces::msg::MetaKey>::SharedPtr meta_sub_;
  rclcpp::TimerBase::SharedPtr timer_;

  TeleopState left_state_;
  TeleopState right_state_;

  std::string left_arm_ip_;
  std::string right_arm_ip_;
  int arm_port_{0};
  std::string left_push_ip_;
  std::string right_push_ip_;
  int push_port_{0};
  int push_cycle_{0};
  std::string meta_topic_;
  double command_rate_hz_{10.0};
  int interpolation_delay_ms_{30};
  double smooth_alpha_{0.3};
  int stale_timeout_ms_{500};
  bool require_trigger_{false};
  int trigger_threshold_{1};
  bool canfd_follow_{true};
  int canfd_trajectory_mode_{2};
  int canfd_smooth_ratio_{200};
};

}  // namespace dual_arm_remote_teleop

int main(int argc, char ** argv)
{
  rclcpp::init(argc, argv);
  auto node = std::make_shared<dual_arm_remote_teleop::RemoteTeleopNode>();
  rclcpp::spin(node);
  rclcpp::shutdown();
  return 0;
}

