#include "obstacle_nav2/obstacle_array_layer.hpp"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <optional>
#include <string>
#include <vector>

#include "geometry_msgs/msg/point_stamped.hpp"
#include "pluginlib/class_list_macros.hpp"
#include "tf2_geometry_msgs/tf2_geometry_msgs.hpp"

namespace obstacle_nav2
{

ObstacleArrayLayer::ObstacleArrayLayer() = default;
ObstacleArrayLayer::~ObstacleArrayLayer()
{
  callback_stop_.store(true, std::memory_order_release);
  if (callback_executor_) {
    callback_executor_->cancel();
  }
  if (callback_thread_.joinable()) {
    callback_thread_.join();
  }
}

void ObstacleArrayLayer::onInitialize()
{
  // CRITICAL: Match the internal costmap size to the master costmap.
  // CostmapLayer::matchSize() is normally called during updateMap(),
  // but obstacles arrive during CONFIGURE phase when subscriptions
  // are active. We need the buffer sized now so rasterizeCircle works.
  CostmapLayer::onInitialize();
  if (layered_costmap_) {
    auto * master = layered_costmap_->getCostmap();
    resizeMap(
      master->getSizeInCellsX(), master->getSizeInCellsY(),
      master->getResolution(),
      master->getOriginX(), master->getOriginY());
  }

  // Lock the base-class LifecycleNode weak_ptr
  auto node = node_.lock();
  if (!node) {
    RCLCPP_ERROR(
      rclcpp::get_logger("obstacle_array_layer"),
      "Failed to lock node in onInitialize");
    return;
  }

  // Parameters (prefixed with layer name per Nav2 convention)
  node->declare_parameter(name_ + ".enabled", true);
  node->declare_parameter(name_ + ".topic", std::string("/obstacles"));
  node->declare_parameter(name_ + ".obstacle_timeout", 0.5);
  node->declare_parameter(name_ + ".transform_tolerance", 0.02);
  node->declare_parameter(name_ + ".default_obstacle_radius", 0.05);
  node->declare_parameter(name_ + ".minimum_obstacle_radius", 0.02);
  node->declare_parameter(name_ + ".maximum_obstacle_radius", 0.50);
  node->declare_parameter(name_ + ".extra_inflation", 0.02);
  node->declare_parameter(name_ + ".retain_previous_on_empty_snapshot", true);

  node->get_parameter(name_ + ".enabled", enabled_);
  node->get_parameter(name_ + ".topic", topic_);
  node->get_parameter(name_ + ".obstacle_timeout", obstacle_timeout_);
  node->get_parameter(name_ + ".transform_tolerance", transform_tolerance_);
  node->get_parameter(name_ + ".default_obstacle_radius", default_obstacle_radius_);
  node->get_parameter(name_ + ".minimum_obstacle_radius", minimum_obstacle_radius_);
  node->get_parameter(name_ + ".maximum_obstacle_radius", maximum_obstacle_radius_);
  node->get_parameter(name_ + ".extra_inflation", extra_inflation_);
  node->get_parameter(
    name_ + ".retain_previous_on_empty_snapshot", retain_previous_on_empty_snapshot_);

  global_frame_ = layered_costmap_->getGlobalFrameID();

  // Costmap plugins are configured after the parent node has entered its
  // executor. Use an explicit callback group so this late-created subscription
  // is guaranteed to be serviced on Humble.
  callback_group_ = node->create_callback_group(
    rclcpp::CallbackGroupType::MutuallyExclusive, false);
  rclcpp::SubscriptionOptions subscription_options;
  subscription_options.callback_group = callback_group_;
  obstacle_sub_ = node->create_subscription<obstacle_scanner::msg::ObstacleArray>(
    topic_, rclcpp::QoS(rclcpp::KeepLast(1)).reliable(),
    std::bind(&ObstacleArrayLayer::obstacleCallback, this, std::placeholders::_1),
    subscription_options);
  callback_executor_ = std::make_unique<rclcpp::executors::SingleThreadedExecutor>();
  callback_executor_->add_callback_group(
    callback_group_, node->get_node_base_interface());
  callback_stop_.store(false, std::memory_order_release);
  callback_thread_ = std::thread(
    [this]() {
      callback_executor_->spin();
    });

  current_ = true;
  has_received_obstacles_ = false;
  last_obstacle_time_ = node->now();
}

void ObstacleArrayLayer::activate() {}
void ObstacleArrayLayer::deactivate() {}

void ObstacleArrayLayer::reset()
{
  std::lock_guard<std::mutex> lock(data_mutex_);
  pending_clear_bounds_ = mergeBounds(pending_clear_bounds_, current_bounds_);
  if (pending_clear_bounds_.valid) {
    ++bounds_generation_;
  }
  resetMap(0, 0, getSizeInCellsX(), getSizeInCellsY());
  current_bounds_ = SnapshotBounds();
  has_received_obstacles_ = false;
}

bool ObstacleArrayLayer::isClearable()
{
  return true;
}

void ObstacleArrayLayer::updateBounds(
  double robot_x, double robot_y, double /*robot_yaw*/,
  double * min_x, double * min_y,
  double * max_x, double * max_y)
{
  auto node = node_.lock();
  if (!node) {return;}

  bool enabled;
  node->get_parameter_or<bool>(name_ + ".enabled", enabled, true);
  if (!enabled) {return;}

  std::lock_guard<std::mutex> lock(data_mutex_);
  if (layered_costmap_->isRolling()) {
    updateOrigin(
      robot_x - getSizeInMetersX() / 2.0,
      robot_y - getSizeInMetersY() / 2.0);
  }
  clearExpiredObstacles(node->now());

  const auto bounds = mergeBounds(current_bounds_, pending_clear_bounds_);
  if (bounds.valid) {
    touch(bounds.min_x, bounds.min_y, min_x, min_y, max_x, max_y);
    touch(bounds.max_x, bounds.max_y, min_x, min_y, max_x, max_y);
  }
  bounds_generation_used_ = bounds_generation_;
}

void ObstacleArrayLayer::updateCosts(
  nav2_costmap_2d::Costmap2D & master_grid,
  int min_i, int min_j, int max_i, int max_j)
{
  auto node = node_.lock();
  if (!node) {return;}

  bool enabled;
  node->get_parameter_or<bool>(name_ + ".enabled", enabled, true);
  if (!enabled) {return;}

  std::lock_guard<std::mutex> lock(data_mutex_);
  updateWithMax(master_grid, min_i, min_j, max_i, max_j);
  if (bounds_generation_used_ == bounds_generation_) {
    pending_clear_bounds_ = SnapshotBounds();
  }
}

// ============================================================================
// Static helpers
// ============================================================================

void ObstacleArrayLayer::rasterizeCircle(
  nav2_costmap_2d::Costmap2D & grid,
  double cx, double cy, double radius,
  double resolution, double origin_x, double origin_y)
{
  if (radius < 0.0) {
    return;
  }

  // Handle zero radius: mark only the cell containing the center
  if (radius <= 0.0) {
    unsigned int cx_cell, cy_cell;
    if (grid.worldToMap(cx, cy, cx_cell, cy_cell)) {
      grid.setCost(cx_cell, cy_cell, nav2_costmap_2d::LETHAL_OBSTACLE);
    }
    return;
  }

  double radius_sq = radius * radius;
  unsigned int size_x = grid.getSizeInCellsX();
  unsigned int size_y = grid.getSizeInCellsY();

  // Bounding box in cell indices
  int ix_min = static_cast<int>(std::floor((cx - radius - origin_x) / resolution));
  int iy_min = static_cast<int>(std::floor((cy - radius - origin_y) / resolution));
  int ix_max = static_cast<int>(std::ceil((cx + radius - origin_x) / resolution));
  int iy_max = static_cast<int>(std::ceil((cy + radius - origin_y) / resolution));

  if (ix_min < 0) {ix_min = 0;}
  if (iy_min < 0) {iy_min = 0;}
  if (ix_max >= static_cast<int>(size_x)) {ix_max = static_cast<int>(size_x) - 1;}
  if (iy_max >= static_cast<int>(size_y)) {iy_max = static_cast<int>(size_y) - 1;}

  for (int j = iy_min; j <= iy_max; ++j) {
    for (int i = ix_min; i <= ix_max; ++i) {
      double wx, wy;
      grid.mapToWorld(
        static_cast<unsigned int>(i),
        static_cast<unsigned int>(j), wx, wy);
      double dx = wx - cx;
      double dy = wy - cy;
      if (dx * dx + dy * dy <= radius_sq) {
        grid.setCost(
          static_cast<unsigned int>(i),
          static_cast<unsigned int>(j),
          nav2_costmap_2d::LETHAL_OBSTACLE);
      }
    }
  }
}

double ObstacleArrayLayer::clampRadius(
  double radius, double min_r, double max_r, double default_r)
{
  if (!std::isfinite(radius) || radius <= 0.0) {
    return default_r;
  }
  if (radius < min_r) {
    return min_r;
  }
  if (radius > max_r) {
    return max_r;
  }
  return radius;
}

std::vector<ObstacleArrayLayer::CircleObstacle> ObstacleArrayLayer::transformSnapshot(
  const obstacle_scanner::msg::ObstacleArray & snapshot,
  const std::string & global_frame,
  tf2_ros::Buffer & tf_buffer,
  const rclcpp::Duration & timeout)
{
  std::vector<CircleObstacle> transformed;
  transformed.reserve(snapshot.obstacles.size());
  if (snapshot.obstacles.empty()) {
    return transformed;
  }

  const rclcpp::Time stamp(snapshot.header.stamp);
  const auto transform = timeout.nanoseconds() > 0 ?
    tf_buffer.lookupTransform(global_frame, snapshot.header.frame_id, stamp, timeout) :
    tf_buffer.lookupTransform(global_frame, snapshot.header.frame_id, stamp);

  for (const auto & obstacle : snapshot.obstacles) {
    geometry_msgs::msg::PointStamped input;
    geometry_msgs::msg::PointStamped output;
    input.header = snapshot.header;
    input.point.x = obstacle.center_x;
    input.point.y = obstacle.center_y;
    input.point.z = 0.0;
    tf2::doTransform(input, output, transform);
    transformed.push_back({output.point.x, output.point.y, obstacle.radius});
  }
  return transformed;
}

ObstacleArrayLayer::SnapshotBounds ObstacleArrayLayer::applySnapshot(
  nav2_costmap_2d::Costmap2D & grid,
  const std::vector<CircleObstacle> & obstacles,
  double resolution, double origin_x, double origin_y)
{
  grid.resetMap(0, 0, grid.getSizeInCellsX(), grid.getSizeInCellsY());
  SnapshotBounds bounds;
  for (const auto & obstacle : obstacles) {
    rasterizeCircle(
      grid, obstacle.center_x, obstacle.center_y, obstacle.radius,
      resolution, origin_x, origin_y);
    if (!bounds.valid) {
      bounds.valid = true;
      bounds.min_x = obstacle.center_x - obstacle.radius;
      bounds.min_y = obstacle.center_y - obstacle.radius;
      bounds.max_x = obstacle.center_x + obstacle.radius;
      bounds.max_y = obstacle.center_y + obstacle.radius;
    } else {
      bounds.min_x = std::min(bounds.min_x, obstacle.center_x - obstacle.radius);
      bounds.min_y = std::min(bounds.min_y, obstacle.center_y - obstacle.radius);
      bounds.max_x = std::max(bounds.max_x, obstacle.center_x + obstacle.radius);
      bounds.max_y = std::max(bounds.max_y, obstacle.center_y + obstacle.radius);
    }
  }
  return bounds;
}

std::optional<ObstacleArrayLayer::SnapshotBounds> ObstacleArrayLayer::applySnapshotIfNotEmpty(
  nav2_costmap_2d::Costmap2D & grid,
  const std::vector<CircleObstacle> & obstacles,
  double resolution, double origin_x, double origin_y)
{
  if (obstacles.empty()) {
    return std::nullopt;
  }
  return applySnapshot(grid, obstacles, resolution, origin_x, origin_y);
}

ObstacleArrayLayer::SnapshotBounds ObstacleArrayLayer::mergeBounds(
  const SnapshotBounds & first, const SnapshotBounds & second)
{
  if (!first.valid) {
    return second;
  }
  if (!second.valid) {
    return first;
  }
  SnapshotBounds merged;
  merged.valid = true;
  merged.min_x = std::min(first.min_x, second.min_x);
  merged.min_y = std::min(first.min_y, second.min_y);
  merged.max_x = std::max(first.max_x, second.max_x);
  merged.max_y = std::max(first.max_y, second.max_y);
  return merged;
}

void ObstacleArrayLayer::clearExpiredObstacles(const rclcpp::Time & now)
{
  if (!has_received_obstacles_) {
    return;
  }
  if ((now - last_obstacle_time_).seconds() <= obstacle_timeout_) {
    return;
  }
  pending_clear_bounds_ = mergeBounds(pending_clear_bounds_, current_bounds_);
  if (pending_clear_bounds_.valid) {
    ++bounds_generation_;
  }
  resetMap(0, 0, getSizeInCellsX(), getSizeInCellsY());
  current_bounds_ = SnapshotBounds();
  has_received_obstacles_ = false;
}

// ============================================================================
// Private
// ============================================================================

void ObstacleArrayLayer::obstacleCallback(
  const obstacle_scanner::msg::ObstacleArray::SharedPtr msg)
{
  auto node = node_.lock();
  if (!node) {return;}

  bool enabled;
  node->get_parameter_or<bool>(name_ + ".enabled", enabled, true);
  if (!enabled) {return;}

  std::vector<CircleObstacle> transformed;
  try {
    transformed = transformSnapshot(
      *msg, global_frame_, *tf_,
      rclcpp::Duration::from_seconds(transform_tolerance_));
  } catch (const tf2::TransformException & e) {
    RCLCPP_WARN_THROTTLE(
      node->get_logger(), *node->get_clock(), 5000,
      "ObstacleArrayLayer: TF failed at message stamp: %s", e.what());
    return;
  }

  std::vector<CircleObstacle> valid_obstacles;
  valid_obstacles.reserve(transformed.size());
  const double origin_x = getOriginX();
  const double origin_y = getOriginY();
  const double map_x_max = origin_x + getSizeInMetersX();
  const double map_y_max = origin_y + getSizeInMetersY();

  for (const auto & obs : transformed) {
    if (!std::isfinite(obs.center_x) || !std::isfinite(obs.center_y)) {
      continue;
    }

    double radius = clampRadius(
      obs.radius, minimum_obstacle_radius_, maximum_obstacle_radius_,
      default_obstacle_radius_);
    double effective_r = radius + extra_inflation_;

    if (obs.center_x + effective_r < origin_x ||
      obs.center_x - effective_r > map_x_max ||
      obs.center_y + effective_r < origin_y ||
      obs.center_y - effective_r > map_y_max)
    {
      continue;
    }
    valid_obstacles.push_back({obs.center_x, obs.center_y, effective_r});
  }

  if (valid_obstacles.empty() && retain_previous_on_empty_snapshot_) {
    return;
  }

  std::lock_guard<std::mutex> lock(data_mutex_);
  const SnapshotBounds previous_bounds = current_bounds_;
  current_bounds_ = applySnapshot(
    *this, valid_obstacles, getResolution(), origin_x, origin_y);
  pending_clear_bounds_ = mergeBounds(pending_clear_bounds_, previous_bounds);
  if (previous_bounds.valid || current_bounds_.valid) {
    ++bounds_generation_;
  }
  has_received_obstacles_ = current_bounds_.valid;
  last_obstacle_time_ = node->now();
}

}  // namespace obstacle_nav2

PLUGINLIB_EXPORT_CLASS(
  obstacle_nav2::ObstacleArrayLayer,
  nav2_costmap_2d::Layer)
