#ifndef RACING_CONTROL__CANDIDATE_WAYPOINT_SELECTOR_HPP_
#define RACING_CONTROL__CANDIDATE_WAYPOINT_SELECTOR_HPP_

#include <cmath>
#include <filesystem>
#include <fstream>
#include <map>
#include <optional>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include <nlohmann/json.hpp>

#include "nav_msgs/msg/occupancy_grid.hpp"
#include "racing_control/racing_control.hpp"

namespace racing_control
{

struct CandidateWaypoint
{
  std::string group;
  std::string source_file;
  std::string point_name;
  geometry_msgs::msg::PoseStamped pose;
};

struct RecoveryCommand
{
  double linear_x{0.0};
  double angular_z{0.0};
};

using CandidateWaypointGroups = std::map<std::string, std::vector<CandidateWaypoint>>;

inline geometry_msgs::msg::PoseStamped candidatePoseFromJson(
  const nlohmann::json & point, const std::string & frame_id)
{
  const auto & pose = point.at("odom").at("pose").at("pose");
  const auto & position = pose.at("position");
  const auto & orientation = pose.at("orientation");
  geometry_msgs::msg::PoseStamped result;
  result.header.frame_id = frame_id;
  result.pose.position.x = position.value("x", 0.0);
  result.pose.position.y = position.value("y", 0.0);
  result.pose.position.z = position.value("z", 0.0);
  result.pose.orientation.x = orientation.value("x", 0.0);
  result.pose.orientation.y = orientation.value("y", 0.0);
  result.pose.orientation.z = orientation.value("z", 0.0);
  result.pose.orientation.w = orientation.value("w", 1.0);
  if (point.contains("yaw_degrees")) {
    const auto yaw = point.at("yaw_degrees").get<double>() * M_PI / 180.0;
    result.pose.orientation.z = std::sin(yaw * 0.5);
    result.pose.orientation.w = std::cos(yaw * 0.5);
  }
  return result;
}

inline CandidateWaypointGroups loadCandidateWaypointGroups(
  const std::string & directory, const std::string & frame_id)
{
  CandidateWaypointGroups groups;
  for (const auto & entry : std::filesystem::directory_iterator(directory)) {
    if (!entry.is_regular_file() || entry.path().extension() != ".json") {
      continue;
    }
    std::ifstream input(entry.path());
    const auto document = nlohmann::json::parse(input);
    for (const auto & point : document.value("points", nlohmann::json::array())) {
      const auto point_name = point.value("name", std::string{});
      if (point_name.empty()) {
        continue;
      }
      groups[point_name].push_back(
        CandidateWaypoint{
          point_name, entry.path().filename().string(), point_name,
          candidatePoseFromJson(point, frame_id)});
    }
  }
  return groups;
}

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

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

  std::optional<CandidateWaypoint> select(
    const std::string & group, const int occupied_threshold, const bool allow_unknown) const
  {
    const auto it = groups_.find(group);
    if (it == groups_.end()) {
      return std::nullopt;
    }
    for (const auto & candidate : it->second) {
      if (isFree(candidate.pose, occupied_threshold, allow_unknown)) {
        return candidate;
      }
    }
    return std::nullopt;
  }

private:
  bool isFree(
    const geometry_msgs::msg::PoseStamped & pose,
    const int occupied_threshold,
    const bool allow_unknown) const
  {
    if (costmap_.info.resolution <= 0.0 || costmap_.info.width == 0 || costmap_.info.height == 0) {
      return true;
    }
    const auto mx = static_cast<int>(std::floor(
      (pose.pose.position.x - costmap_.info.origin.position.x) / costmap_.info.resolution));
    const auto my = static_cast<int>(std::floor(
      (pose.pose.position.y - costmap_.info.origin.position.y) / costmap_.info.resolution));
    if (mx < 0 || my < 0 || mx >= static_cast<int>(costmap_.info.width) ||
      my >= static_cast<int>(costmap_.info.height))
    {
      return false;
    }
    const auto value = costmap_.data[
      static_cast<std::size_t>(my) * costmap_.info.width + static_cast<std::size_t>(mx)];
    return value < 0 ? allow_unknown : value < occupied_threshold;
  }

  CandidateWaypointGroups groups_;
  nav_msgs::msg::OccupancyGrid costmap_;
};

inline std::optional<RecoveryCommand> nearestFreeRecoveryCommand(
  const nav_msgs::msg::OccupancyGrid & costmap,
  const geometry_msgs::msg::PoseStamped & current,
  const int occupied_threshold,
  const bool allow_unknown,
  const double max_distance)
{
  if (costmap.info.resolution <= 0.0 || costmap.info.width == 0 || costmap.info.height == 0) {
    return std::nullopt;
  }

  const auto yaw = std::atan2(
    2.0 * (current.pose.orientation.w * current.pose.orientation.z),
    1.0 - 2.0 * current.pose.orientation.z * current.pose.orientation.z);
  const auto cos_yaw = std::cos(yaw);
  const auto sin_yaw = std::sin(yaw);
  double best_distance = max_distance;
  double best_forward = 0.0;
  bool found = false;
  for (std::size_t my = 0; my < costmap.info.height; ++my) {
    for (std::size_t mx = 0; mx < costmap.info.width; ++mx) {
      const auto value = costmap.data[my * costmap.info.width + mx];
      if (value < 0 ? !allow_unknown : value >= occupied_threshold) {
        continue;
      }
      const auto world_x = costmap.info.origin.position.x +
        (static_cast<double>(mx) + 0.5) * costmap.info.resolution;
      const auto world_y = costmap.info.origin.position.y +
        (static_cast<double>(my) + 0.5) * costmap.info.resolution;
      const auto dx = world_x - current.pose.position.x;
      const auto dy = world_y - current.pose.position.y;
      const auto forward = cos_yaw * dx + sin_yaw * dy;
      const auto lateral = -sin_yaw * dx + cos_yaw * dy;
      const auto distance = std::hypot(forward, lateral);
      if (distance <= 0.0 || distance > max_distance || distance >= best_distance) {
        continue;
      }
      best_distance = distance;
      best_forward = forward;
      found = true;
    }
  }
  if (!found) {
    return std::nullopt;
  }
  return RecoveryCommand{best_forward >= 0.0 ? 0.2 : -0.2, 0.0};
}

}  // namespace racing_control

#endif  // RACING_CONTROL__CANDIDATE_WAYPOINT_SELECTOR_HPP_
