#include "racing_control/candidate_waypoint_selector.hpp"

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <limits>
#include <nlohmann/json.hpp>

namespace racing_control
{
namespace
{

double normalizeAngle(double angle)
{
  while (angle > M_PI) {
    angle -= 2.0 * M_PI;
  }
  while (angle < -M_PI) {
    angle += 2.0 * M_PI;
  }
  return angle;
}

double yawFromQuaternionInternal(
  const double x, const double y, const double z, const double w)
{
  return std::atan2(
    2.0 * (w * z + x * y),
    1.0 - 2.0 * (y * y + z * z));
}

std::vector<std::filesystem::path> sortedJsonFiles(const std::filesystem::path & json_dir)
{
  std::vector<std::filesystem::path> files;
  for (const auto & entry : std::filesystem::directory_iterator(json_dir)) {
    if (entry.is_regular_file() && entry.path().extension() == ".json") {
      files.push_back(entry.path());
    }
  }
  std::sort(
    files.begin(), files.end(), [](const auto & a, const auto & b) {
      return a.filename().string() < b.filename().string();
    });
  return files;
}

geometry_msgs::msg::PoseStamped poseFromJsonPoint(
  const nlohmann::json & point,
  const std::string & frame_id)
{
  const auto & odom = point.at("odom");
  const auto & pose = odom.at("pose").at("pose");
  const auto & position = pose.at("position");
  const auto & orientation = pose.at("orientation");

  const double x = position.at("x").get<double>();
  const double y = position.at("y").get<double>();
  const double yaw = point.contains("yaw_degrees") ?
    point.at("yaw_degrees").get<double>() * M_PI / 180.0 :
    yawFromQuaternionInternal(
    orientation.at("x").get<double>(), orientation.at("y").get<double>(),
    orientation.at("z").get<double>(), orientation.at("w").get<double>());

  geometry_msgs::msg::PoseStamped stamped;
  stamped.header.frame_id = frame_id;
  stamped.pose.position.x = x;
  stamped.pose.position.y = y;
  stamped.pose.position.z = 0.0;
  stamped.pose.orientation.z = std::sin(yaw * 0.5);
  stamped.pose.orientation.w = std::cos(yaw * 0.5);
  return stamped;
}

bool fileContainsGroup(
  const std::string & source_file,
  const std::string & preferred_source_file)
{
  if (preferred_source_file.empty()) {
    return false;
  }
  return source_file == preferred_source_file;
}

}  // namespace

CandidateWaypointGroups loadCandidateWaypointGroups(
  const std::string & json_dir,
  const std::string & frame_id)
{
  CandidateWaypointGroups groups;
  const auto root = std::filesystem::path(json_dir);
  if (json_dir.empty() || !std::filesystem::exists(root)) {
    return groups;
  }

  for (const auto & path : sortedJsonFiles(root)) {
    std::ifstream input(path);
    if (!input) {
      continue;
    }
    nlohmann::json data;
    try {
      input >> data;
    } catch (...) {
      continue;
    }

    const auto points = data.is_array() ? data : data.value("points", nlohmann::json::array());
    if (!points.is_array()) {
      continue;
    }

    for (const auto & point : points) {
      if (!point.is_object() || !point.contains("odom") || !point["odom"].is_object()) {
        continue;
      }
      CandidateWaypoint candidate;
      candidate.source_file = path.filename().string();
      candidate.point_name = point.value("name", path.stem().string());
      candidate.group_name = candidate.point_name;
      candidate.pose = poseFromJsonPoint(point, frame_id);
      groups[candidate.group_name].push_back(std::move(candidate));
    }
  }

  return groups;
}

void CandidateWaypointSelector::setGroups(CandidateWaypointGroups groups)
{
  groups_ = std::move(groups);
}

void CandidateWaypointSelector::setCostmap(const nav_msgs::msg::OccupancyGrid & costmap)
{
  costmap_ = costmap;
  has_costmap_ = true;
}

std::optional<CandidateWaypoint> CandidateWaypointSelector::select(
  const std::string & group_name,
  const int occupied_threshold,
  const bool treat_unknown_as_occupied,
  const std::string & preferred_source_file) const
{
  const auto group_it = groups_.find(group_name);
  if (group_it == groups_.end() || group_it->second.empty() || !has_costmap_) {
    return std::nullopt;
  }

  std::vector<const CandidateWaypoint *> ordered;
  ordered.reserve(group_it->second.size());
  for (const auto & candidate : group_it->second) {
    ordered.push_back(&candidate);
  }
  std::stable_sort(
    ordered.begin(), ordered.end(), [&](const auto * lhs, const auto * rhs) {
      const bool lhs_preferred = fileContainsGroup(lhs->source_file, preferred_source_file);
      const bool rhs_preferred = fileContainsGroup(rhs->source_file, preferred_source_file);
      if (lhs_preferred != rhs_preferred) {
        return lhs_preferred;
      }
      return lhs->source_file < rhs->source_file;
    });

  for (const auto * candidate : ordered) {
    if (!isCandidateOccupied(candidate->pose, occupied_threshold, treat_unknown_as_occupied)) {
      return *candidate;
    }
  }
  return std::nullopt;
}

std::optional<RecoveryCommand> nearestFreeRecoveryCommand(
  const nav_msgs::msg::OccupancyGrid & costmap,
  const geometry_msgs::msg::PoseStamped & current,
  const int occupied_threshold,
  const bool treat_unknown_as_occupied,
  const double search_radius_m)
{
  CandidateWaypointSelector selector;
  selector.setCostmap(costmap);
  return selector.nearestFreeRecoveryCommand(
    current, occupied_threshold, treat_unknown_as_occupied, search_radius_m);
}

std::optional<RecoveryCommand> rearClearRecoveryCommand(
  const nav_msgs::msg::OccupancyGrid & costmap,
  const geometry_msgs::msg::PoseStamped & current,
  const int occupied_threshold,
  const bool treat_unknown_as_occupied,
  const double rear_clear_distance_m)
{
  CandidateWaypointSelector selector;
  selector.setCostmap(costmap);
  const double yaw = 2.0 * std::atan2(current.pose.orientation.z, current.pose.orientation.w);
  geometry_msgs::msg::PoseStamped rear = current;
  rear.pose.position.x -= std::cos(yaw) * rear_clear_distance_m;
  rear.pose.position.y -= std::sin(yaw) * rear_clear_distance_m;
  if (selector.isPoseOccupied(rear, occupied_threshold, treat_unknown_as_occupied)) {
    return std::nullopt;
  }
  return RecoveryCommand{-0.25, 0.0, rear_clear_distance_m / 0.25, "rear-clear-backward"};
}

std::optional<RecoveryCommand> CandidateWaypointSelector::nearestFreeRecoveryCommand(
  const geometry_msgs::msg::PoseStamped & current,
  const int occupied_threshold,
  const bool treat_unknown_as_occupied,
  const double search_radius_m) const
{
  if (!has_costmap_) {
    return std::nullopt;
  }

  const double yaw = 2.0 * std::atan2(current.pose.orientation.z, current.pose.orientation.w);
  const std::vector<double> distances{0.25, 0.5, 0.75, std::max(0.25, search_radius_m)};
  const std::vector<double> angles{0.0, M_PI / 4.0, -M_PI / 4.0, M_PI / 2.0, -M_PI / 2.0, M_PI};

  double best_score = std::numeric_limits<double>::infinity();
  std::optional<RecoveryCommand> best_command;
  for (const double distance : distances) {
    for (const double angle : angles) {
      const double sample_angle = normalizeAngle(yaw + angle);
      geometry_msgs::msg::PoseStamped sample = current;
      sample.pose.position.x += std::cos(sample_angle) * distance;
      sample.pose.position.y += std::sin(sample_angle) * distance;
      if (isCandidateOccupied(sample, occupied_threshold, treat_unknown_as_occupied)) {
        continue;
      }
      const double score = distance + std::abs(angle) * 0.1;
      if (score < best_score) {
        best_score = score;
        best_command = commandForRecoveryAngle(angle);
      }
    }
  }
  return best_command;
}

bool CandidateWaypointSelector::isPoseOccupied(
  const geometry_msgs::msg::PoseStamped & pose,
  const int occupied_threshold,
  const bool treat_unknown_as_occupied) const
{
  return isCandidateOccupied(pose, occupied_threshold, treat_unknown_as_occupied);
}

bool CandidateWaypointSelector::isCandidateOccupied(
  const geometry_msgs::msg::PoseStamped & pose,
  const int occupied_threshold,
  const bool treat_unknown_as_occupied) const
{
  const auto & costmap = costmap_;
  const auto & origin = costmap.info.origin;
  const double resolution = costmap.info.resolution;
  const auto width = static_cast<int>(costmap.info.width);
  const auto height = static_cast<int>(costmap.info.height);
  if (resolution <= 0.0 || width <= 0 || height <= 0) {
    return true;
  }

  const double origin_yaw = yawFromQuaternionInternal(
    origin.orientation.x, origin.orientation.y, origin.orientation.z, origin.orientation.w);
  const double dx = pose.pose.position.x - origin.position.x;
  const double dy = pose.pose.position.y - origin.position.y;
  const double cos_yaw = std::cos(-origin_yaw);
  const double sin_yaw = std::sin(-origin_yaw);
  const double local_x = cos_yaw * dx - sin_yaw * dy;
  const double local_y = sin_yaw * dx + cos_yaw * dy;
  const auto mx = static_cast<int>(std::floor(local_x / resolution));
  const auto my = static_cast<int>(std::floor(local_y / resolution));
  if (mx < 0 || my < 0 || mx >= width || my >= height) {
    return true;
  }

  const auto index = my * width + mx;
  if (index < 0 || index >= static_cast<int>(costmap.data.size())) {
    return true;
  }

  const auto cost = static_cast<int>(costmap.data[static_cast<std::size_t>(index)]);
  if (cost < 0) {
    return treat_unknown_as_occupied;
  }
  return cost >= occupied_threshold;
}

std::optional<RecoveryCommand> CandidateWaypointSelector::commandForRecoveryAngle(
  const double angle_rad) const
{
  const double angle = normalizeAngle(angle_rad);
  if (std::abs(angle) < M_PI / 8.0) {
    return RecoveryCommand{0.25, 0.0, 0.6, "forward"};
  }
  if (std::abs(std::abs(angle) - M_PI) < M_PI / 8.0) {
    return RecoveryCommand{-0.25, 0.0, 0.6, "backward"};
  }
  if (angle > 0.0 && angle < M_PI / 2.0) {
    return RecoveryCommand{0.25, 0.7, 0.6, "forward-left"};
  }
  if (angle < 0.0 && angle > -M_PI / 2.0) {
    return RecoveryCommand{0.25, -0.7, 0.6, "forward-right"};
  }
  if (angle >= M_PI / 2.0) {
    return RecoveryCommand{0.0, 0.7, 0.4, "left"};
  }
  return RecoveryCommand{0.0, -0.7, 0.4, "right"};
}

}  // namespace racing_control
