#include "obstacle_nav2/trajectory_guard.hpp"

#include <algorithm>
#include <cmath>
#include <limits>

namespace obstacle_nav2
{
namespace
{

bool worldToMap(
  const nav_msgs::msg::OccupancyGrid & grid,
  double wx,
  double wy,
  int & mx,
  int & my)
{
  const double resolution = grid.info.resolution;
  if (resolution <= 0.0) {
    return false;
  }

  const double ox = grid.info.origin.position.x;
  const double oy = grid.info.origin.position.y;
  if (wx < ox || wy < oy) {
    return false;
  }

  mx = static_cast<int>((wx - ox) / resolution);
  my = static_cast<int>((wy - oy) / resolution);
  return mx >= 0 && my >= 0 &&
    mx < static_cast<int>(grid.info.width) &&
    my < static_cast<int>(grid.info.height);
}

bool occupiedAt(
  const nav_msgs::msg::OccupancyGrid & grid,
  double wx,
  double wy,
  const GuardSettings & settings)
{
  int mx = 0;
  int my = 0;
  if (!worldToMap(grid, wx, wy, mx, my)) {
    return settings.treat_unknown_as_occupied;
  }

  const auto value = grid.data[my * grid.info.width + mx];
  if (value < 0) {
    return settings.treat_unknown_as_occupied;
  }
  return value >= settings.occupied_threshold;
}

bool footprintOccupied(
  const nav_msgs::msg::OccupancyGrid & costmap,
  double cx,
  double cy,
  double yaw,
  const GuardSettings & settings)
{
  const double half_l = settings.footprint_half_length + settings.footprint_padding;
  const double half_w = settings.footprint_half_width + settings.footprint_padding;
  const double step = std::max(0.02, settings.footprint_sample_step);
  const double c = std::cos(yaw);
  const double s = std::sin(yaw);

  auto sample = [&](double x, double y) {
      const double wx = cx + c * x - s * y;
      const double wy = cy + s * x + c * y;
      return occupiedAt(costmap, wx, wy, settings);
    };

  if (sample(0.0, 0.0)) {
    return true;
  }
  for (double x = -half_l; x <= half_l + 1.0e-9; x += step) {
    if (sample(x, -half_w) || sample(x, 0.0) || sample(x, half_w)) {
      return true;
    }
  }
  for (double y = -half_w; y <= half_w + 1.0e-9; y += step) {
    if (sample(-half_l, y) || sample(0.0, y) || sample(half_l, y)) {
      return true;
    }
  }
  return false;
}

}  // namespace

double yawFromPose(const geometry_msgs::msg::PoseStamped & pose)
{
  const auto & q = pose.pose.orientation;
  return std::atan2(
    2.0 * (q.w * q.z + q.x * q.y),
    1.0 - 2.0 * (q.y * q.y + q.z * q.z));
}

double distance2d(
  const geometry_msgs::msg::PoseStamped & a,
  const geometry_msgs::msg::PoseStamped & b)
{
  const double dx = a.pose.position.x - b.pose.position.x;
  const double dy = a.pose.position.y - b.pose.position.y;
  return std::hypot(dx, dy);
}

std::size_t nearestPathIndex(
  const nav_msgs::msg::Path & path,
  const geometry_msgs::msg::PoseStamped & pose)
{
  std::size_t best_index = 0;
  double best_distance = std::numeric_limits<double>::infinity();

  for (std::size_t i = 0; i < path.poses.size(); ++i) {
    const double d = distance2d(path.poses[i], pose);
    if (d < best_distance) {
      best_distance = d;
      best_index = i;
    }
  }
  return best_index;
}

std::size_t advanceByDistance(
  const nav_msgs::msg::Path & path,
  std::size_t start_index,
  double distance_m)
{
  if (path.poses.empty()) {
    return 0;
  }

  std::size_t index = std::min(start_index, path.poses.size() - 1);
  double traveled = 0.0;
  while (index + 1 < path.poses.size() && traveled < distance_m) {
    traveled += distance2d(path.poses[index], path.poses[index + 1]);
    ++index;
  }
  return index;
}

CollisionCheckResult checkPathAhead(
  const nav_msgs::msg::Path & path,
  std::size_t start_index,
  const nav_msgs::msg::OccupancyGrid & costmap,
  const GuardSettings & settings)
{
  if (path.poses.empty()) {
    return CollisionCheckResult{true, 0, "empty_path"};
  }

  const auto clamped_start = std::min(start_index, path.poses.size() - 1);
  const auto end_index =
    advanceByDistance(path, clamped_start, settings.lookahead_distance);
  const double step = std::max(0.02, settings.footprint_sample_step);

  for (std::size_t i = clamped_start; i <= end_index && i < path.poses.size(); ++i) {
    if (i + 1 >= path.poses.size() || i == end_index) {
      const auto & pose = path.poses[i];
      if (footprintOccupied(
          costmap, pose.pose.position.x, pose.pose.position.y, yawFromPose(pose), settings))
      {
        return CollisionCheckResult{true, i, "occupied"};
      }
      continue;
    }

    const auto & a = path.poses[i];
    const auto & b = path.poses[i + 1];
    const double dx = b.pose.position.x - a.pose.position.x;
    const double dy = b.pose.position.y - a.pose.position.y;
    const double segment_length = std::hypot(dx, dy);
    const double yaw = segment_length > 1.0e-6 ? std::atan2(dy, dx) : yawFromPose(a);
    const int samples = std::max(1, static_cast<int>(std::ceil(segment_length / step)));

    for (int sample_index = 0; sample_index <= samples; ++sample_index) {
      const double t = static_cast<double>(sample_index) / static_cast<double>(samples);
      const double x = a.pose.position.x + t * dx;
      const double y = a.pose.position.y + t * dy;
      if (footprintOccupied(costmap, x, y, yaw, settings)) {
        return CollisionCheckResult{true, i, "occupied"};
      }
    }
  }

  return CollisionCheckResult{false, end_index, "clear"};
}

std::optional<std::size_t> findClearRejoinIndex(
  const nav_msgs::msg::Path & path,
  std::size_t start_index,
  const nav_msgs::msg::OccupancyGrid & costmap,
  const GuardSettings & settings)
{
  if (path.poses.empty()) {
    return std::nullopt;
  }

  const auto first = advanceByDistance(path, start_index, settings.rejoin_min_distance);
  const auto last = advanceByDistance(path, start_index, settings.rejoin_max_distance);
  GuardSettings single_pose_settings = settings;
  single_pose_settings.lookahead_distance = 0.0;

  for (std::size_t i = first; i <= last && i < path.poses.size(); ++i) {
    const auto result = checkPathAhead(path, i, costmap, single_pose_settings);
    if (!result.blocked) {
      return i;
    }
  }
  return std::nullopt;
}

std::optional<std::size_t> findClearRejoinIndexWithPlannerCostmap(
  const nav_msgs::msg::Path & path,
  std::size_t start_index,
  const nav_msgs::msg::OccupancyGrid & local_costmap,
  const nav_msgs::msg::OccupancyGrid * planner_costmap,
  const GuardSettings & settings)
{
  if (planner_costmap != nullptr) {
    return findClearRejoinIndex(path, start_index, *planner_costmap, settings);
  }
  return findClearRejoinIndex(path, start_index, local_costmap, settings);
}

bool shouldRetryBlockedRepair(
  const std::optional<std::size_t> & last_repair_nearest_index,
  std::size_t nearest_index,
  int min_repair_progress_indices,
  double seconds_since_last_repair,
  double blocked_retry_wait_sec,
  bool costmap_updated_since_repair)
{
  if (!last_repair_nearest_index.has_value()) {
    return true;
  }

  const auto min_progress =
    static_cast<std::size_t>(std::max(0, min_repair_progress_indices));
  if (nearest_index > *last_repair_nearest_index + min_progress) {
    return true;
  }

  return seconds_since_last_repair >= blocked_retry_wait_sec &&
    costmap_updated_since_repair;
}

nav_msgs::msg::Path slicePath(
  const nav_msgs::msg::Path & path,
  std::size_t start_index,
  std::size_t end_index_inclusive)
{
  nav_msgs::msg::Path out;
  out.header = path.header;
  if (path.poses.empty()) {
    return out;
  }

  const auto begin = std::min(start_index, path.poses.size() - 1);
  const auto end = std::min(end_index_inclusive, path.poses.size() - 1);
  for (std::size_t i = begin; i <= end; ++i) {
    out.poses.push_back(path.poses[i]);
  }
  return out;
}

nav_msgs::msg::Path stitchPaths(
  const nav_msgs::msg::Path & bypass,
  const nav_msgs::msg::Path & original,
  std::size_t rejoin_index)
{
  nav_msgs::msg::Path out;
  out.header = bypass.header.frame_id.empty() ? original.header : bypass.header;
  out.poses = bypass.poses;

  if (original.poses.empty()) {
    return out;
  }

  const auto begin = std::min(rejoin_index, original.poses.size());
  for (std::size_t i = begin; i < original.poses.size(); ++i) {
    out.poses.push_back(original.poses[i]);
  }
  return out;
}

}  // namespace obstacle_nav2
