#include <iostream>
#include <vector>
#include <filesystem>

#include <tinyxml2.h>
#include <geometric_shapes/mesh_operations.h>
#include <geometric_shapes/shape_operations.h>
#include <sdf/sdf.hh>
#include <sdf/World.hh>

#include <rclcpp/rclcpp.hpp>
#include <tf2_ros/transform_listener.h>
#include <tf2_ros/static_transform_broadcaster.h>
#include <tf2_ros/buffer.h>

#include <moveit_msgs/msg/collision_object.hpp>
#include <moveit/planning_scene_interface/planning_scene_interface.h>

#include "rbs_skill_interfaces/srv/planning_scene_object_state.hpp"

#define SERVICE_NAME "update_planning_scene"
#define INIT_SCENE_PARAM_NAME "init_scene"
#define CUSTOM_MODEL_PATH_PARAM_NAME "models_paths"

static geometry_msgs::msg::Pose pose_from_pose3d(const ignition::math::Pose3d &pose)
{
  geometry_msgs::msg::Pose result;
  auto position = pose.Pos();
  result.position.set__x(position.X());
  result.position.set__y(position.Y());
  result.position.set__z(position.Z());
  auto orientation = pose.Rot();
  result.orientation.set__x(orientation.X());
  result.orientation.set__y(orientation.Y());
  result.orientation.set__z(orientation.Z());
  result.orientation.set__w(orientation.W());
  return result;
}

static std::string get_correct_mesh_path(
  const std::string& uri, const std::vector<std::string> &resources)
{
  std::string result = "";
  std::regex reg(R"((?:model|package)(?:\:\/)(.*))");
  std::smatch m;
  if (std::regex_match(uri, m, reg))
  {
    std::string rel_path = m[1];
    std::for_each(resources.begin(), resources.end(), 
      [&result, &rel_path](const std::string& res){
        if (result.empty())
          result = std::filesystem::exists(res + rel_path)? std::string(res + rel_path): result;   
    });
  }
  return "file://" + result; 
}

static moveit_msgs::msg::CollisionObject 
get_object(const sdf::Model *model, const std::vector<std::string> &resources)
{
  moveit_msgs::msg::CollisionObject obj;
  obj.header.frame_id = "world";
  obj.id = model->Name();
  obj.pose = pose_from_pose3d(model->RawPose());
  size_t link_count = model->LinkCount();
  for (size_t i = 0; i < link_count; ++i)
  {
    auto link = model->LinkByIndex(i);
    auto collision = link->CollisionByIndex(0);
    auto link_pose = pose_from_pose3d(link->RawPose());
    auto geometry = collision->Geom();
    switch(geometry->Type())
    {
      case sdf::GeometryType::MESH:
      {
        auto path = get_correct_mesh_path(geometry->MeshShape()->Uri(), resources);
        shapes::Mesh *m = shapes::createMeshFromResource(
          !path.empty()? path: geometry->MeshShape()->Uri());
        auto scale = geometry->MeshShape()->Scale().X();
        m->scale(scale);

        Eigen::Vector3d centroid = {0, 0, 0};
        for (size_t i = 0; i < 3 * m->vertex_count; i += 3)
        {
          const auto x = m->vertices[i];
          const auto y = m->vertices[i+1];
          const auto z = m->vertices[i+2];

          centroid.x() += x*(1-scale);
          centroid.y() += y*(1-scale);
          centroid.z() += z*(1-scale);
        }

        centroid = centroid / m->vertex_count;
        for (size_t i = 0; i < 3 * m->vertex_count; i += 3) {
          m->vertices[i] -= centroid.x();
          m->vertices[i + 1] -= centroid.y();
          m->vertices[i + 2] -= centroid.z();
        }

        shape_msgs::msg::Mesh mesh;
        shapes::ShapeMsg m_msg;
        shapes::constructMsgFromShape(m, m_msg);
        mesh = boost::get<shape_msgs::msg::Mesh>(m_msg);
        obj.meshes.push_back(mesh);
        obj.mesh_poses.push_back(link_pose);
        break;
      } case sdf::GeometryType::BOX: {
        auto sdf_box = geometry->BoxShape();
        shape_msgs::msg::SolidPrimitive box;
        box.type = shape_msgs::msg::SolidPrimitive::BOX;
        auto sdf_box_size = sdf_box->Size();
        box.dimensions.push_back(sdf_box_size.X());
        box.dimensions.push_back(sdf_box_size.Y());
        box.dimensions.push_back(sdf_box_size.Z());
        obj.primitives.push_back(box);
        obj.primitive_poses.push_back(link_pose);
        break;
      }
      case sdf::GeometryType::PLANE:
      {
        auto sdf_plane = geometry->PlaneShape();
        shape_msgs::msg::Plane plane;
        auto normal = sdf_plane->Normal();
        plane.coef[0] = normal.X();
        plane.coef[1] = normal.Y();
        plane.coef[2] = normal.Z();
        obj.planes.push_back(plane);
        obj.plane_poses.push_back(link_pose);
        break;
      }
      case sdf::GeometryType::EMPTY:
      case sdf::GeometryType::CYLINDER:
      case sdf::GeometryType::SPHERE:
      case sdf::GeometryType::HEIGHTMAP:
      case sdf::GeometryType::CAPSULE:
      case sdf::GeometryType::ELLIPSOID:
      case sdf::GeometryType::POLYLINE:
        break; 
    }
  }

  return obj;
}

static std::vector<moveit_msgs::msg::CollisionObject> 
get_objects(const sdf::World *world, const std::string &model_resources)
{
  std::vector<std::string> resources;
  std::regex reg("\\:+");
  std::sregex_token_iterator begin(
    model_resources.begin(), model_resources.end(), reg, -1), end;
  std::copy(++begin, end, std::back_inserter(resources));
  std::vector<moveit_msgs::msg::CollisionObject> result;
  auto models_count = world->ModelCount();

  for (size_t i = 0; i < models_count; ++i)
  {
    try{
      auto model = world->ModelByIndex(i);
      result.push_back(get_object(model, resources));
    } catch (std::exception &ex){
      std::cerr << ex.what() << std::endl;
    }
  }

  return result;
}

class UpdatePlanningSceneServer: public rclcpp::Node
{
public:
    UpdatePlanningSceneServer(rclcpp::NodeOptions options) 
    : Node("update_planning_scene_node", options.allow_undeclared_parameters(true).automatically_declare_parameters_from_overrides(true))
    {
      std::string scene = get_parameter(INIT_SCENE_PARAM_NAME).as_string(); 
      std::string model_resources = get_parameter(CUSTOM_MODEL_PATH_PARAM_NAME).as_string();

      if (!scene.empty())
      {
        RCLCPP_INFO(get_logger(), "Init scene from file: %s", scene.c_str());
        init_scene(scene, model_resources);
      }
      service_ = create_service<rbs_skill_interfaces::srv::PlanningSceneObjectState>(
        SERVICE_NAME, std::bind(&UpdatePlanningSceneServer::callback, this, 
          std::placeholders::_1, std::placeholders::_2));
    }

    bool init_scene(const std::string &init_scene, const std::string &model_resources)
    {
      sdf::Root root;
      sdf::ParserConfig config;
      config.AddURIPath("package://", model_resources);
      config.AddURIPath("model://", model_resources);
      sdf::Errors errors = root.Load(init_scene, config);
      if (!errors.empty())
      {
        for (auto const &e: errors)
          RCLCPP_ERROR(get_logger(), "%s", e.Message().c_str());
        return false;
      }
      auto world = root.WorldByIndex(0);
      auto objects = get_objects(world, model_resources);
      planning_scene_.applyCollisionObjects(objects);
      return true;
    }

    void callback(
      std::shared_ptr<rbs_skill_interfaces::srv::PlanningSceneObjectState::Request> request,
      std::shared_ptr<rbs_skill_interfaces::srv::PlanningSceneObjectState::Response> response)
    {
      auto state = static_cast<UpdatePlanningSceneRequestState>(request->req_kind);
      moveit_msgs::msg::CollisionObject obj;
      obj.id = request->frame_name;
      switch(state)
      {
        case ADD:
        case UPDATE:
          obj.meshes.resize(1);
          obj.mesh_poses.resize(1);
          obj.operation = state == ADD? 
            moveit_msgs::msg::CollisionObject::ADD: 
            moveit_msgs::msg::CollisionObject::MOVE;
          obj.meshes.at(0) = std::move(request->mesh);
          obj.mesh_poses.at(0) = std::move(request->pose.pose);
          break;
        case REMOVE:
          obj.operation = moveit_msgs::msg::CollisionObject::REMOVE;
          break;
      }
      
      response->call_status = planning_scene_.applyCollisionObject(obj); 
    }

private:
    enum UpdatePlanningSceneRequestState
    {
      ADD=0,
      REMOVE=1,
      UPDATE=2
    };

    rclcpp::Service<rbs_skill_interfaces::srv::PlanningSceneObjectState>::SharedPtr service_;  
    moveit::planning_interface::PlanningSceneInterface planning_scene_;
};

#include "rclcpp_components/register_node_macro.hpp"

RCLCPP_COMPONENTS_REGISTER_NODE(UpdatePlanningSceneServer);