#include "saltybot_segmentation_costmap/segmentation_costmap_layer.hpp"

#include <algorithm>
#include <string>

#include "nav2_costmap_2d/costmap_2d.hpp"
#include "pluginlib/class_list_macros.hpp"

PLUGINLIB_EXPORT_CLASS(
  saltybot_segmentation_costmap::SegmentationCostmapLayer,
  nav2_costmap_2d::Layer)

namespace saltybot_segmentation_costmap
{

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

SegmentationCostmapLayer::SegmentationCostmapLayer()
: last_min_x_(-1e9), last_min_y_(-1e9),
  last_max_x_(1e9),  last_max_y_(1e9),
  need_update_(false)
{}

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

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

  declareParameter("enabled",            rclcpp::ParameterValue(true));
  declareParameter("topic",              rclcpp::ParameterValue(std::string("/segmentation/costmap")));
  declareParameter("combination_method", rclcpp::ParameterValue(std::string("max")));

  enabled_            = node->get_parameter(name_ + "." + "enabled").as_bool();
  topic_              = node->get_parameter(name_ + "." + "topic").as_string();
  combination_method_ = node->get_parameter(name_ + "." + "combination_method").as_string();

  RCLCPP_INFO(
    node->get_logger(),
    "SegmentationCostmapLayer: topic=%s method=%s",
    topic_.c_str(), combination_method_.c_str());

  // Transient local QoS to match the sidewalk_seg_node publisher
  rclcpp::QoS qos(rclcpp::KeepLast(1));
  qos.transient_local();
  qos.reliable();

  sub_ = node->create_subscription<nav_msgs::msg::OccupancyGrid>(
    topic_, qos,
    std::bind(
      &SegmentationCostmapLayer::segCostmapCallback, this,
      std::placeholders::_1));

  current_ = true;
}

// ── Subscription callback ─────────────────────────────────────────────────────

void SegmentationCostmapLayer::segCostmapCallback(
  const nav_msgs::msg::OccupancyGrid::SharedPtr msg)
{
  std::lock_guard<std::mutex> lock(grid_mutex_);
  latest_grid_ = msg;
  need_update_ = true;
}

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

void SegmentationCostmapLayer::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;}

  std::lock_guard<std::mutex> lock(grid_mutex_);
  if (!latest_grid_) {return;}

  const auto & info = latest_grid_->info;
  double ox = info.origin.position.x;
  double oy = info.origin.position.y;
  double gw = info.width  * info.resolution;
  double gh = info.height * info.resolution;

  last_min_x_ = ox;
  last_min_y_ = oy;
  last_max_x_ = ox + gw;
  last_max_y_ = oy + gh;

  *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 SegmentationCostmapLayer::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(grid_mutex_);
  if (!latest_grid_ || !need_update_) {return;}

  const auto & info  = latest_grid_->info;
  const auto & data  = latest_grid_->data;

  float seg_res = info.resolution;
  float seg_ox  = static_cast<float>(info.origin.position.x);
  float seg_oy  = static_cast<float>(info.origin.position.y);
  int   seg_w   = static_cast<int>(info.width);
  int   seg_h   = static_cast<int>(info.height);

  float master_res = static_cast<float>(master.getResolution());

  // Iterate over the update window
  for (int j = min_j; j < max_j; ++j) {
    for (int i = min_i; i < max_i; ++i) {
      // World coords of this master cell centre
      double wx, wy;
      master.mapToWorld(
        static_cast<unsigned int>(i), static_cast<unsigned int>(j), wx, wy);

      // Corresponding cell in segmentation grid
      int seg_col = static_cast<int>((wx - seg_ox) / seg_res);
      int seg_row = static_cast<int>((wy - seg_oy) / seg_res);

      if (seg_col < 0 || seg_col >= seg_w ||
          seg_row < 0 || seg_row >= seg_h) {
        continue;
      }

      int seg_idx = seg_row * seg_w + seg_col;
      if (seg_idx < 0 || seg_idx >= static_cast<int>(data.size())) {
        continue;
      }

      int8_t occ = data[static_cast<size_t>(seg_idx)];
      unsigned char new_cost = 0;
      if (!occupancyToCost(occ, new_cost)) {
        continue;  // unknown cell — leave master unchanged
      }

      unsigned char existing = master.getCost(
        static_cast<unsigned int>(i), static_cast<unsigned int>(j));

      unsigned char final_cost = existing;
      if (combination_method_ == "override") {
        final_cost = new_cost;
      } else if (combination_method_ == "min") {
        final_cost = std::min(existing, new_cost);
      } else {
        // "max" (default) — most conservative
        final_cost = std::max(existing, new_cost);
      }

      master.setCost(
        static_cast<unsigned int>(i), static_cast<unsigned int>(j),
        final_cost);
    }
  }

  need_update_ = false;
}

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

void SegmentationCostmapLayer::reset()
{
  std::lock_guard<std::mutex> lock(grid_mutex_);
  latest_grid_.reset();
  need_update_ = false;
  current_ = true;
}

// ── occupancyToCost ───────────────────────────────────────────────────────────

bool SegmentationCostmapLayer::occupancyToCost(int8_t occ, unsigned char & cost_out)
{
  if (occ < 0) {
    return false;  // unknown — leave cell unchanged
  }
  if (occ == 0) {
    cost_out = FREE_SPACE;
    return true;
  }
  if (occ >= 100) {
    cost_out = LETHAL_OBSTACLE;
    return true;
  }
  // Scale 1–99 → 1 to (INSCRIBED_INFLATED_OBSTACLE-1) = 1 to 252
  int scaled = 1 + static_cast<int>(
    (occ - 1) * (INSCRIBED_INFLATED_OBSTACLE - 2) / 98.0f);
  cost_out = static_cast<unsigned char>(
    std::min(scaled, static_cast<int>(INSCRIBED_INFLATED_OBSTACLE - 1)));
  return true;
}

}  // namespace saltybot_segmentation_costmap
