#include <rclcpp/rclcpp.hpp>
#include <moveit/planning_scene/planning_scene.h>
#include <moveit/planning_scene_interface/planning_scene_interface.h>
#include <moveit/task_constructor/task.h>
#include <moveit/task_constructor/solvers.h>
#include <moveit/task_constructor/stages.h>
#if __has_include(<tf2_geometry_msgs/tf2_geometry_msgs.hpp>)
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>
#else
#include <tf2_geometry_msgs/tf2_geometry_msgs.h>
#endif
#if __has_include(<tf2_eigen/tf2_eigen.hpp>)
#include <tf2_eigen/tf2_eigen.hpp>
#else
#include <tf2_eigen/tf2_eigen.h>
#endif

#include <chrono>
#include <thread>
#include <functional>
#include <memory>


#include "rclcpp_action/rclcpp_action.hpp"
#include "rclcpp_components/register_node_macro.hpp"
#include "action_pick_place_interfaces/action/pick_place.hpp"


static const rclcpp::Logger LOGGER = rclcpp::get_logger("scara_arm_mus_pick_place");
namespace mtc = moveit::task_constructor;

class MTCTaskNode : public rclcpp::Node
{
public:

  using PickPlace = action_pick_place_interfaces::action::PickPlace;
  using GoalHandlePickPlace = rclcpp_action::ServerGoalHandle<PickPlace>;

  explicit MTCTaskNode(const rclcpp::NodeOptions & options = rclcpp::NodeOptions())
  : Node("pick_place_node", options)
  // : node_{ std::make_shared<rclcpp::Node>("pick_place_node", options) }
{
  using namespace std::placeholders;

  this->action_server_ = rclcpp_action::create_server<PickPlace>(
    this,
    "pick_place",
    std::bind(&MTCTaskNode::handle_goal, this, _1, _2),
    std::bind(&MTCTaskNode::handle_cancel, this, _1),
    std::bind(&MTCTaskNode::handle_accepted, this, _1));
}   

  rclcpp::node_interfaces::NodeBaseInterface::SharedPtr getNodeBaseInterface();
  void doTask(const std::shared_ptr<GoalHandlePickPlace> goal_handle, 
              const std::shared_ptr<const PickPlace::Goal> goal, 
              const std::shared_ptr<PickPlace::Result> result);
  void setupPlanningScene(const std::shared_ptr<GoalHandlePickPlace> goal_handle, 
                          const std::shared_ptr<const PickPlace::Goal> goal, 
                          const std::shared_ptr<PickPlace::Result> result);

private:
  // Compose an MTC task from a series of stages.
  mtc::Task createTask(int place_pos_x, int place_pos_y);
  mtc::Task task_;
  rclcpp::Node::SharedPtr node_;
  rclcpp_action::Server<PickPlace>::SharedPtr action_server_;

  rclcpp_action::GoalResponse handle_goal(
    const rclcpp_action::GoalUUID & uuid,
    std::shared_ptr<const PickPlace::Goal> goal)
  {
    RCLCPP_INFO(LOGGER, "Received goal request with order %f", goal->pick_x);
    (void)uuid;
    return rclcpp_action::GoalResponse::ACCEPT_AND_EXECUTE;
  }

  rclcpp_action::CancelResponse handle_cancel(
    const std::shared_ptr<GoalHandlePickPlace> goal_handle)
  {
    RCLCPP_INFO(LOGGER, "Received request to cancel goal");
    (void)goal_handle;
    return rclcpp_action::CancelResponse::ACCEPT;
  }

  void handle_accepted(const std::shared_ptr<GoalHandlePickPlace> goal_handle)
  {
    using namespace std::placeholders;
    // this needs to return quickly to avoid blocking the executor, so spin up a new thread
    std::thread{std::bind(&MTCTaskNode::execute, this, _1), goal_handle}.detach();
  }

  void execute(const std::shared_ptr<GoalHandlePickPlace> goal_handle)
  {
    RCLCPP_INFO(LOGGER, "Executing goal");

    // // rclcpp::Rate loop_rate(1);
    const auto goal = goal_handle->get_goal();
    auto result = std::make_shared<PickPlace::Result>();
    
    this->setupPlanningScene(goal_handle, goal, result);
    this->doTask(goal_handle, goal, result);
    
    // Check if goal is done
    if (rclcpp::ok()) {
      result->status = true;
      goal_handle->succeed(result);
      RCLCPP_INFO(LOGGER, "Goal succeeded");
    }
  }

  
};

rclcpp::node_interfaces::NodeBaseInterface::SharedPtr MTCTaskNode::getNodeBaseInterface()
{
  return node_->get_node_base_interface();
}

void MTCTaskNode::setupPlanningScene(const std::shared_ptr<GoalHandlePickPlace> goal_handle, 
                                     const std::shared_ptr<const PickPlace::Goal> goal, 
                                     const std::shared_ptr<PickPlace::Result> result)
{
  if (goal_handle->is_canceling()) {
    result->status = false;
    goal_handle->canceled(result);
    RCLCPP_INFO(LOGGER, "Goal canceled");
    return;
  }
  RCLCPP_INFO(LOGGER, "Running...");
  
  moveit_msgs::msg::CollisionObject object;
  object.id = "object";
  object.header.frame_id = "world";
  object.primitives.resize(1);
  object.primitives[0].type = shape_msgs::msg::SolidPrimitive::BOX;
  object.primitives[0].dimensions = { 0.02, 0.02, 0.1 };

  geometry_msgs::msg::Pose pose;
  pose.position.x = goal->pick_x;
  pose.position.y = goal->pick_y;
  pose.position.z = 0.05;
  pose.orientation.w = 1.0;
  object.pose = pose;

  moveit::planning_interface::PlanningSceneInterface psi;
  psi.applyCollisionObject(object);
}

void MTCTaskNode::doTask(const std::shared_ptr<GoalHandlePickPlace> goal_handle, 
                         const std::shared_ptr<const PickPlace::Goal> goal, 
                         const std::shared_ptr<PickPlace::Result> result)
{
  if (goal_handle->is_canceling()) {
    result->status = false;
    goal_handle->canceled(result);
    RCLCPP_INFO(LOGGER, "Goal canceled");
    return;
  }
  RCLCPP_INFO(LOGGER, "Running...");

  int place_pos_x = goal->place_x;
  int place_pos_y = goal->place_y;
  
  task_ = createTask(place_pos_x, place_pos_y);

  try
  {
    task_.init();
  }
  catch (mtc::InitStageException& e)
  {
    RCLCPP_ERROR_STREAM(LOGGER, e);
    return;
  }

  if (!task_.plan(5)) //generates a plan, stopping after 5 successful plans are found
  {
    RCLCPP_ERROR_STREAM(LOGGER, "Task planning failed");
    return;
  }
  task_.introspection().publishSolution(*task_.solutions().front());

  // auto result = task_.execute(*task_.solutions().front());
  // if (result.val != moveit_msgs::msg::MoveItErrorCodes::SUCCESS)
  // {
  //   RCLCPP_ERROR_STREAM(LOGGER, "Task execution failed");
  //   return;
  // }

  return;
}

mtc::Task MTCTaskNode::createTask(int place_pos_x, int place_pos_y)
{
  mtc::Task task;
  task.stages()->setName("demo task"); //task name (can be anything)
  task.loadRobotModel(node_);

  const auto& arm_group_name = "arm";
  const auto& hand_group_name = "gripper";
  const auto& hand_frame = "Link_4";

  // Set task properties
  task.setProperty("group", arm_group_name);
  task.setProperty("eef", hand_group_name);
  task.setProperty("ik_frame", hand_frame);

// Disable warnings for this line, as it's a variable that's set but not used in this example
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-but-set-variable"
  mtc::Stage* current_state_ptr = nullptr;  // Forward current_state on to grasp pose generator
#pragma GCC diagnostic pop

  auto stage_state_current = std::make_unique<mtc::stages::CurrentState>("current");
  current_state_ptr = stage_state_current.get();
  task.add(std::move(stage_state_current));

  //3 different planners to choose from 
  auto sampling_planner = std::make_shared<mtc::solvers::PipelinePlanner>(node_);
  auto interpolation_planner = std::make_shared<mtc::solvers::JointInterpolationPlanner>();

  auto cartesian_planner = std::make_shared<mtc::solvers::CartesianPath>();
  cartesian_planner->setMaxVelocityScalingFactor(1.0);
  cartesian_planner->setMaxAccelerationScalingFactor(1.0);
  cartesian_planner->setStepSize(0.01);

  auto stage_open_hand = std::make_unique<mtc::stages::MoveTo>("open_hand", sampling_planner);
  stage_open_hand->setGroup(hand_group_name);
  stage_open_hand->setGoal("open"); //defined in SRDF while setting up moveit_setup_assistant
  task.add(std::move(stage_open_hand));

  auto stage_move_to_pick = std::make_unique<mtc::stages::Connect>("move to pick",
    mtc::stages::Connect::GroupPlannerVector{ { arm_group_name, interpolation_planner } });
  stage_move_to_pick->setTimeout(5.0); 
  stage_move_to_pick->properties().configureInitFrom(mtc::Stage::PARENT);
  task.add(std::move(stage_move_to_pick));

  mtc::Stage* attach_object_stage = nullptr;  // Forward attach_object_stage to place pose generator

  // serial container containing the stages relevant to the picking action
  {
    auto grasp = std::make_unique<mtc::SerialContainer>("pick object");
    task.properties().exposeTo(grasp->properties(), { "eef", "group", "ik_frame" });
    grasp->properties().configureInitFrom(mtc::Stage::PARENT,{ "eef", "group", "ik_frame" });

    {
    auto stage = std::make_unique<mtc::stages::MoveRelative>("approach object", interpolation_planner);
    stage->properties().set("marker_ns", "approach_object");
    stage->properties().set("link", hand_frame);
    stage->properties().configureInitFrom(mtc::Stage::PARENT, { "group" });
    stage->setMinMaxDistance(0.01, 0.15);

    // Set hand forward direction
    geometry_msgs::msg::Vector3Stamped vec;
    vec.header.frame_id = "world";
    vec.vector.y = -1.0;
    stage->setDirection(vec);
    grasp->insert(std::move(stage));
    }

    {
    // Sample grasp pose
    auto stage = std::make_unique<mtc::stages::GenerateGraspPose>("generate grasp pose");
    stage->properties().configureInitFrom(mtc::Stage::PARENT);
    stage->properties().set("marker_ns", "grasp_pose");
    stage->setPreGraspPose("open");
    stage->setObject("object");
    stage->setAngleDelta(M_PI / 12);
    stage->setMonitoredStage(current_state_ptr);  // Hook into current state
    Eigen::Isometry3d grasp_frame_transform;
    Eigen::Quaterniond q = Eigen::AngleAxisd(-M_PI / 2, Eigen::Vector3d::UnitX()) *
                           Eigen::AngleAxisd(0 / 2, Eigen::Vector3d::UnitY()) *
                           Eigen::AngleAxisd(M_PI / 2, Eigen::Vector3d::UnitZ());
    grasp_frame_transform.linear() = q.matrix();
    grasp_frame_transform.translation().y() = -0.06;
    
    // Compute IK
    auto wrapper = std::make_unique<mtc::stages::ComputeIK>("grasp pose IK", std::move(stage));
    wrapper->setMaxIKSolutions(8);
    wrapper->setMinSolutionDistance(1.0);
    wrapper->setIKFrame(grasp_frame_transform, hand_frame);
    wrapper->properties().configureInitFrom(mtc::Stage::PARENT, { "eef", "group" });
    wrapper->properties().configureInitFrom(mtc::Stage::INTERFACE, { "target_pose" });
    grasp->insert(std::move(wrapper));
    }

    {
    auto stage = std::make_unique<mtc::stages::ModifyPlanningScene>("allow collision (hand,object)");
    stage->allowCollisions("object", task.getRobotModel()
                                        ->getJointModelGroup(hand_group_name)
                                        ->getLinkModelNamesWithCollisionGeometry(), true);
    grasp->insert(std::move(stage));
    }

    {
    auto stage = std::make_unique<mtc::stages::MoveTo>("close hand", interpolation_planner);
    stage->setGroup(hand_group_name);
    stage->setGoal("close");
    grasp->insert(std::move(stage));
    }

    {
    auto stage = std::make_unique<mtc::stages::ModifyPlanningScene>("attach object");
    stage->attachObject("object", hand_frame);
    attach_object_stage = stage.get();
    grasp->insert(std::move(stage));
    }

    {
    auto stage = std::make_unique<mtc::stages::MoveRelative>("lift object", interpolation_planner);
    stage->properties().configureInitFrom(mtc::Stage::PARENT, { "group" });
    stage->setMinMaxDistance(0.01, 0.03);
    stage->setIKFrame(hand_frame);
    stage->properties().set("marker_ns", "lift_object");

    geometry_msgs::msg::Vector3Stamped vec;
    vec.header.frame_id = "world";
    vec.vector.z = 1.0;
    stage->setDirection(vec);
    grasp->insert(std::move(stage));
    }
    task.add(std::move(grasp));
  }

  {
    auto stage_move_to_place = std::make_unique<mtc::stages::Connect>("move to place",
                               mtc::stages::Connect::GroupPlannerVector{ {arm_group_name, interpolation_planner} });
    stage_move_to_place->setTimeout(5.0);
    stage_move_to_place->properties().configureInitFrom(mtc::Stage::PARENT);
    task.add(std::move(stage_move_to_place));
  }

  // serial container containing the stages relevant to the placing action
  {
    auto place = std::make_unique<mtc::SerialContainer>("place object");
    task.properties().exposeTo(place->properties(), { "eef", "group", "ik_frame" });
    place->properties().configureInitFrom(mtc::Stage::PARENT,{ "eef", "group", "ik_frame" });

    {
      // Sample place pose
      auto stage = std::make_unique<mtc::stages::GeneratePlacePose>("generate place pose");
      stage->properties().configureInitFrom(mtc::Stage::PARENT);
      stage->properties().set("marker_ns", "place_pose");
      stage->setObject("object");

      geometry_msgs::msg::PoseStamped target_pose_msg;
      target_pose_msg.header.frame_id = "world";
      target_pose_msg.pose.position.x = place_pos_x;
      target_pose_msg.pose.position.y = place_pos_y;
      target_pose_msg.pose.position.z = 0.05;
      target_pose_msg.pose.orientation.w = 1.0;
      stage->setPose(target_pose_msg);
      stage->setMonitoredStage(attach_object_stage);  // Hook into attach_object_stage

      // Compute IK
      auto wrapper = std::make_unique<mtc::stages::ComputeIK>("place pose IK", std::move(stage));
      wrapper->setMaxIKSolutions(2);
      wrapper->setMinSolutionDistance(1.0);
      wrapper->setIKFrame("object");
      wrapper->properties().configureInitFrom(mtc::Stage::PARENT, { "eef", "group" });
      wrapper->properties().configureInitFrom(mtc::Stage::INTERFACE, { "target_pose" });
      place->insert(std::move(wrapper));
    }

    {
      auto stage = std::make_unique<mtc::stages::MoveTo>("open hand", interpolation_planner);
      stage->setGroup(hand_group_name);
      stage->setGoal("open");
      place->insert(std::move(stage));
    }

    {
      auto stage = std::make_unique<mtc::stages::ModifyPlanningScene>("forbid collision (hand,object)");
      stage->allowCollisions("object",task.getRobotModel()
                                         ->getJointModelGroup(hand_group_name)
                                         ->getLinkModelNamesWithCollisionGeometry(),
                                         false);
      place->insert(std::move(stage));
    }

    {
      auto stage = std::make_unique<mtc::stages::ModifyPlanningScene>("detach object");
      stage->detachObject("object", hand_frame);
      place->insert(std::move(stage));
    }

    {
      auto stage = std::make_unique<mtc::stages::MoveRelative>("retreat", cartesian_planner);
      stage->properties().configureInitFrom(mtc::Stage::PARENT, { "group" });
      stage->setMinMaxDistance(0.01, 0.03);
      stage->setIKFrame(hand_frame);
      stage->properties().set("marker_ns", "retreat");

      // Set retreat direction
      geometry_msgs::msg::Vector3Stamped vec;
      vec.header.frame_id = "world";
      vec.vector.z = 0.1;
      stage->setDirection(vec);
      place->insert(std::move(stage));
    }
    task.add(std::move(place));
  }

  {
    auto stage = std::make_unique<mtc::stages::MoveTo>("return home", interpolation_planner);
    stage->properties().configureInitFrom(mtc::Stage::PARENT, { "group" });
    stage->setGoal("initial");
    task.add(std::move(stage));
  }

  return task;
}

int main(int argc, char * argv[])
{
  rclcpp::init(argc, argv);
  rclcpp::NodeOptions options;
  options.automatically_declare_parameters_from_overrides(true);
  rclcpp::spin(std::make_shared<MTCTaskNode>());
  rclcpp::shutdown();
  return 0;
}