#include "saltybot_depth_costmap/depth_costmap_layer.hpp"

#include <algorithm>
#include <cmath>
#include <cstring>
#include <string>
#include <vector>

#include "nav2_costmap_2d/costmap_2d.hpp"
#include "pluginlib/class_list_macros.hpp"
#include "tf2/exceptions.h"

PLUGINLIB_EXPORT_CLASS(
  saltybot_depth_costmap::DepthCostmapLayer,
  nav2_costmap_2d::Layer)

namespace saltybot_depth_costmap
{

using nav2_costmap_2d::FREE_SPACE;
using nav2_costmap_2d::INSCRIBED_INFLATED_OBSTACLE;
using nav2_costmap_2d::LETHAL_OBSTACLE;
using nav2_costmap_2d::NO_INFORMATION;

DepthCostmapLayer::DepthCostmapLayer()
: last_min_x_(-1e9), last_min_y_(-1e9),
  last_max_x_(1e9), last_max_y_(1e9),
  logger_(rclcpp::get_logger("DepthCostmapLayer"))
{}

// ── onInitialize ──────────────────────────────────────────────────────────────

void DepthCostmapLayer::onInitialize()
{
  auto node = node_.lock();
  if (!node) {
    throw std::runtime_error("DepthCostmapLayer: node handle expired");
  }

  logger_ = node->get_logger();
  clock_  = node->get_clock();

  declareParameter("enabled",            rclcpp::ParameterValue(true));
  declareParameter("depth_topic",        rclcpp::ParameterValue(
    std::string("/camera/depth/image_rect_raw")));
  declareParameter("camera_info_topic",  rclcpp::ParameterValue(
    std::string("/camera/depth/camera_info")));
  declareParameter("min_height",         rclcpp::ParameterValue(0.05));
  declareParameter("max_height",         rclcpp::ParameterValue(0.80));
  declareParameter("obstacle_range",     rclcpp::ParameterValue(3.5));
  declareParameter("clearing_range",     rclcpp::ParameterValue(4.0));
  declareParameter("inflation_radius",   rclcpp::ParameterValue(0.10));
  declareParameter("downsample_factor",  rclcpp::ParameterValue(4));

  enabled_           = node->get_parameter(name_ + ".enabled").as_bool();
  depth_topic_       = node->get_parameter(name_ + ".depth_topic").as_string();
  camera_info_topic_ = node->get_parameter(name_ + ".camera_info_topic").as_string();
  min_height_        = node->get_parameter(name_ + ".min_height").as_double();
  max_height_        = node->get_parameter(name_ + ".max_height").as_double();
  obstacle_range_    = node->get_parameter(name_ + ".obstacle_range").as_double();
  clearing_range_    = node->get_parameter(name_ + ".clearing_range").as_double();
  inflation_radius_  = node->get_parameter(name_ + ".inflation_radius").as_double();
  downsample_factor_ = node->get_parameter(name_ + ".downsample_factor").as_int();

  if (downsample_factor_ < 1) downsample_factor_ = 1;

  RCLCPP_INFO(
    logger_,
    "DepthCostmapLayer '%s': depth=%s height=[%.2f, %.2f]m range=%.1fm inflate=%.2fm ds=%d",
    name_.c_str(), depth_topic_.c_str(),
    min_height_, max_height_, obstacle_range_, inflation_radius_, downsample_factor_);

  // TF2
  tf_buffer_   = std::make_shared<tf2_ros::Buffer>(clock_);
  tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);

  // Subscriptions
  depth_sub_ = node->create_subscription<sensor_msgs::msg::Image>(
    depth_topic_, rclcpp::SensorDataQoS(),
    std::bind(&DepthCostmapLayer::depthCallback, this, std::placeholders::_1));

  info_sub_ = node->create_subscription<sensor_msgs::msg::CameraInfo>(
    camera_info_topic_, rclcpp::QoS(1).reliable(),
    std::bind(&DepthCostmapLayer::cameraInfoCallback, this, std::placeholders::_1));

  current_ = true;
}

// ── cameraInfoCallback ────────────────────────────────────────────────────────

void DepthCostmapLayer::cameraInfoCallback(
  const sensor_msgs::msg::CameraInfo::SharedPtr msg)
{
  if (have_camera_info_) return;  // Intrinsics are stable — only need once

  // Pinhole intrinsics from K matrix (row-major 3x3)
  fx_ = msg->k[0];   // K[0,0]
  fy_ = msg->k[4];   // K[1,1]
  cx_ = msg->k[2];   // K[0,2]
  cy_ = msg->k[5];   // K[1,2]
  camera_frame_ = msg->header.frame_id;

  if (fx_ <= 0.0 || fy_ <= 0.0) {
    RCLCPP_WARN(logger_, "DepthCostmapLayer: invalid camera intrinsics (fx=%.1f fy=%.1f)", fx_, fy_);
    return;
  }

  have_camera_info_ = true;
  RCLCPP_INFO(
    logger_,
    "DepthCostmapLayer: camera_info received — frame=%s fx=%.1f fy=%.1f cx=%.1f cy=%.1f",
    camera_frame_.c_str(), fx_, fy_, cx_, cy_);
}

// ── depthCallback / processDepthImage ────────────────────────────────────────

void DepthCostmapLayer::depthCallback(
  const sensor_msgs::msg::Image::SharedPtr msg)
{
  if (!enabled_ || !have_camera_info_) return;
  processDepthImage(msg);
}

void DepthCostmapLayer::processDepthImage(
  const sensor_msgs::msg::Image::SharedPtr & img)
{
  // Resolve encoding
  bool is_16u = (img->encoding == "16UC1" || img->encoding == "mono16");
  bool is_32f = (img->encoding == "32FC1");
  if (!is_16u && !is_32f) {
    RCLCPP_WARN_ONCE(logger_, "DepthCostmapLayer: unsupported encoding '%s'", img->encoding.c_str());
    return;
  }

  // Get TF from camera → costmap global frame (latest available)
  const std::string global_frame = layered_costmap_->getGlobalFrameID();
  geometry_msgs::msg::TransformStamped tf_stamped;
  try {
    tf_stamped = tf_buffer_->lookupTransform(
      global_frame, camera_frame_,
      rclcpp::Time(0),
      rclcpp::Duration::from_seconds(0.1));
  } catch (const tf2::TransformException & ex) {
    RCLCPP_WARN_THROTTLE(logger_, *clock_, 5000,
      "DepthCostmapLayer: TF %s→%s failed: %s",
      camera_frame_.c_str(), global_frame.c_str(), ex.what());
    return;
  }

  // Build rotation matrix from quaternion (camera frame → world frame)
  // p_world = R * p_cam + t
  const auto & t = tf_stamped.transform.translation;
  const auto & q = tf_stamped.transform.rotation;
  double qw = q.w, qx = q.x, qy = q.y, qz = q.z;

  double R[3][3] = {
    {1.0 - 2.0*(qy*qy + qz*qz),       2.0*(qx*qy - qz*qw),  2.0*(qx*qz + qy*qw)},
    {      2.0*(qx*qy + qz*qw), 1.0 - 2.0*(qx*qx + qz*qz),  2.0*(qy*qz - qx*qw)},
    {      2.0*(qx*qz - qy*qw),       2.0*(qy*qz + qx*qw), 1.0 - 2.0*(qx*qx + qy*qy)}
  };

  int rows = static_cast<int>(img->height);
  int cols = static_cast<int>(img->width);
  int step = static_cast<int>(img->step);

  const double obstacle_range_sq = obstacle_range_ * obstacle_range_;

  std::vector<ObstaclePoint> new_points;
  new_points.reserve(512);

  for (int v = 0; v < rows; v += downsample_factor_) {
    const uint8_t * row_ptr = img->data.data() + v * step;

    for (int u = 0; u < cols; u += downsample_factor_) {
      double depth_m = 0.0;

      if (is_16u) {
        uint16_t raw;
        std::memcpy(&raw, row_ptr + u * 2, sizeof(raw));
        if (raw == 0) continue;
        depth_m = raw * 0.001;  // mm → m
      } else {
        float raw;
        std::memcpy(&raw, row_ptr + u * 4, sizeof(raw));
        if (!std::isfinite(raw) || raw <= 0.0f) continue;
        depth_m = static_cast<double>(raw);
      }

      if (depth_m > clearing_range_) continue;

      // Back-project: depth optical frame (Z forward, X right, Y down)
      double xc = (u - cx_) * depth_m / fx_;
      double yc = (v - cy_) * depth_m / fy_;
      double zc = depth_m;

      // Distance check in camera frame (cheaper than world-frame dist)
      double dist2 = xc*xc + yc*yc + zc*zc;
      if (dist2 > obstacle_range_sq) continue;

      // Transform to world frame
      double xw = R[0][0]*xc + R[0][1]*yc + R[0][2]*zc + t.x;
      double yw = R[1][0]*xc + R[1][1]*yc + R[1][2]*zc + t.y;
      double zw = R[2][0]*xc + R[2][1]*yc + R[2][2]*zc + t.z;

      // Height filter (zw = height above ground in world frame)
      if (zw < min_height_ || zw > max_height_) continue;

      new_points.push_back({xw, yw});
    }
  }

  {
    std::lock_guard<std::mutex> lock(points_mutex_);
    obstacle_points_ = std::move(new_points);
    need_update_ = true;
  }
}

// ── updateBounds ──────────────────────────────────────────────────────────────

void DepthCostmapLayer::updateBounds(
  double robot_x, double robot_y, double /*robot_yaw*/,
  double * min_x, double * min_y,
  double * max_x, double * max_y)
{
  if (!enabled_) return;

  // Mark region around robot that depth camera can reach
  double margin = obstacle_range_ + inflation_radius_;
  last_min_x_ = robot_x - margin;
  last_min_y_ = robot_y - margin;
  last_max_x_ = robot_x + margin;
  last_max_y_ = robot_y + margin;

  *min_x = std::min(*min_x, last_min_x_);
  *min_y = std::min(*min_y, last_min_y_);
  *max_x = std::max(*max_x, last_max_x_);
  *max_y = std::max(*max_y, last_max_y_);
}

// ── updateCosts ───────────────────────────────────────────────────────────────

void DepthCostmapLayer::updateCosts(
  nav2_costmap_2d::Costmap2D & master,
  int min_i, int min_j, int max_i, int max_j)
{
  if (!enabled_) return;

  std::lock_guard<std::mutex> lock(points_mutex_);
  if (!need_update_ || obstacle_points_.empty()) return;

  double res = master.getResolution();
  // Number of cells to inflate in each direction
  int inflate_cells = (inflation_radius_ > 0.0)
    ? static_cast<int>(std::ceil(inflation_radius_ / res))
    : 0;

  for (const auto & pt : obstacle_points_) {
    unsigned int mx, my;
    if (!master.worldToMap(pt.wx, pt.wy, mx, my)) continue;

    int ci = static_cast<int>(mx);
    int cj = static_cast<int>(my);

    // Mark the obstacle cell as lethal
    if (ci >= min_i && ci < max_i && cj >= min_j && cj < max_j) {
      master.setCost(mx, my, LETHAL_OBSTACLE);
    }

    // In-layer inflation: mark neighbouring cells within inflation_radius
    // as INSCRIBED_INFLATED_OBSTACLE (if not already lethal)
    for (int di = -inflate_cells; di <= inflate_cells; ++di) {
      for (int dj = -inflate_cells; dj <= inflate_cells; ++dj) {
        if (di == 0 && dj == 0) continue;

        double dist = std::hypot(static_cast<double>(di), static_cast<double>(dj)) * res;
        if (dist > inflation_radius_) continue;

        int ni = ci + di;
        int nj = cj + dj;
        if (ni < min_i || ni >= max_i || nj < min_j || nj >= max_j) continue;

        auto uni = static_cast<unsigned int>(ni);
        auto unj = static_cast<unsigned int>(nj);
        unsigned char existing = master.getCost(uni, unj);
        if (existing < INSCRIBED_INFLATED_OBSTACLE) {
          master.setCost(uni, unj, INSCRIBED_INFLATED_OBSTACLE);
        }
      }
    }
  }

  need_update_ = false;
}

// ── reset ─────────────────────────────────────────────────────────────────────

void DepthCostmapLayer::reset()
{
  std::lock_guard<std::mutex> lock(points_mutex_);
  obstacle_points_.clear();
  need_update_ = false;
  current_ = true;
}

}  // namespace saltybot_depth_costmap
