#include <chrono>
#include <memory>
#include <mutex>
#include <cmath>

#include "rclcpp/rclcpp.hpp"
#include "rclcpp_action/rclcpp_action.hpp"
#include "nav_msgs/msg/path.hpp"
#include "nav_msgs/msg/odometry.hpp"
#include "geometry_msgs/msg/twist.hpp"
#include "geometry_msgs/msg/pose_stamped.hpp"
#include "nav2_msgs/action/compute_path_to_pose.hpp"
#include "origincar_msg/msg/sign.hpp"
#include "tf2/LinearMath/Quaternion.h"
#include "tf2/LinearMath/Matrix3x3.h"

#include "past_control/tools.h"

using namespace std::chrono_literals;

// Task states
enum class TaskState
{
  IDLE    = 0,
  GOING   = 1,
  RETURN  = 2,
  QR_SCAN = 3,
  RESET   = 4
};

class RacingOrchestrator : public rclcpp::Node
{
public:
  using ComputePathToPose = nav2_msgs::action::ComputePathToPose;
  using GoalHandleComputePathToPose = rclcpp_action::ClientGoalHandle<ComputePathToPose>;

  RacingOrchestrator()
  : Node("racing_orchestrator"),
    state_(TaskState::IDLE),
    has_odom_(false), has_goal_(false),
    has_path_(false), got_lane_cmd_(false),
    is_back_(false), activate_avoid_(false),
    waypoint_index_(0), guide_step_(5),
    follow_linear_speed_(0.4),
    sub_target_(true),
    cmd_vel_linear_x_(0.0), cmd_vel_angular_z_(0.0),
    cur_x_(0.0), cur_y_(0.0), cur_yaw_(0.0)
  {
    // Parameters
    // PID path following
    this->declare_parameter("follow_linear_speed", 0.4);
    this->declare_parameter("guide_step", 5);
    this->declare_parameter("arrive_square", 0.25);
    this->declare_parameter("angular_kp", 1.5);
    this->declare_parameter("angular_ki", 0.0);
    this->declare_parameter("angular_kd", 0.1);
    this->declare_parameter("angular_integral_max", 1.0);
    this->declare_parameter("angular_output_max", 2.0);
    this->declare_parameter("angular_max_err", 3.14);

    // Obstacle avoidance
    this->declare_parameter("avoid_linear_speed", 0.2);
    this->declare_parameter("avoid_angular_z", 0.8);

    // QR scan params
    this->declare_parameter("qr_scan_speed", 0.15);
    this->declare_parameter("qr_scan_duration", 3.0);

    // Control rate
    this->declare_parameter("control_rate", 20.0);

    // Load params
    follow_linear_speed_ = this->get_parameter("follow_linear_speed").as_double();
    guide_step_ = this->get_parameter("guide_step").as_int();
    arrive_square_ = this->get_parameter("arrive_square").as_double();
    avoid_linear_speed_ = this->get_parameter("avoid_linear_speed").as_double();
    avoid_angular_z_ = this->get_parameter("avoid_angular_z").as_double();
    qr_scan_speed_ = this->get_parameter("qr_scan_speed").as_double();
    qr_scan_duration_ = this->get_parameter("qr_scan_duration").as_double();

    // Init PID
    double kp = this->get_parameter("angular_kp").as_double();
    double ki = this->get_parameter("angular_ki").as_double();
    double kd = this->get_parameter("angular_kd").as_double();
    double i_max = this->get_parameter("angular_integral_max").as_double();
    double o_max = this->get_parameter("angular_output_max").as_double();
    double max_err = this->get_parameter("angular_max_err").as_double();

    angular_pid_.init(kp, ki, kd, i_max, o_max, max_err);

    // Subscribers
    sign_sub_ = this->create_subscription<origincar_msg::msg::Sign>(
      "sign4return", 10,
      std::bind(&RacingOrchestrator::signCallback, this, std::placeholders::_1));

    sign_foxglove_sub_ = this->create_subscription<origincar_msg::msg::Sign>(
      "sign_foxglove", 10,
      std::bind(&RacingOrchestrator::signFoxgloveCallback, this, std::placeholders::_1));

    path_sub_ = this->create_subscription<nav_msgs::msg::Path>(
      "/plan", 10,
      std::bind(&RacingOrchestrator::pathCallback, this, std::placeholders::_1));

    goal_sub_ = this->create_subscription<geometry_msgs::msg::PoseStamped>(
      "/goal_pose", 10,
      std::bind(&RacingOrchestrator::goalCallback, this, std::placeholders::_1));

    odom_sub_ = this->create_subscription<nav_msgs::msg::Odometry>(
      "/odom", 10,
      std::bind(&RacingOrchestrator::odomCallback, this, std::placeholders::_1));

    lane_cmd_sub_ = this->create_subscription<geometry_msgs::msg::Twist>(
      "/lane_cmd_vel", 10,
      std::bind(&RacingOrchestrator::laneCmdCallback, this, std::placeholders::_1));

    // Publishers
    cmd_vel_pub_ = this->create_publisher<geometry_msgs::msg::Twist>("/cmd_vel", 10);

    // ComputePathToPose action client (triggers Nav2 planner_server)
    planner_client_ = rclcpp_action::create_client<ComputePathToPose>(
      this, "compute_path_to_pose");

    // Control loop
    double rate = this->get_parameter("control_rate").as_double();
    int period_ms = static_cast<int>(1000.0 / rate);
    timer_ = this->create_wall_timer(
      std::chrono::milliseconds(period_ms),
      std::bind(&RacingOrchestrator::controlLoop, this));

    RCLCPP_INFO(this->get_logger(), "RacingOrchestrator started. State: IDLE");
  }

private:
  // ─────────────────────────────────────────────
  //  A. Task State Machine — sign callbacks
  // ─────────────────────────────────────────────

  void signCallback(const origincar_msg::msg::Sign::SharedPtr msg)
  {
    processSign(msg->sign_data);
  }

  void signFoxgloveCallback(const origincar_msg::msg::Sign::SharedPtr msg)
  {
    processSign(msg->sign_data);
  }

  void processSign(int sign_data)
  {
    RCLCPP_INFO(this->get_logger(), "Received sign: %d", sign_data);

    if (sign_data == -1) {
      state_ = TaskState::RESET;
      has_path_ = false;
      global_path_ = nav_msgs::msg::Path();
      waypoint_index_ = 0;
      angular_pid_.reset();
      RCLCPP_INFO(this->get_logger(), "State -> RESET");
    }
    else if (sign_data == 3 || sign_data == 4) {
      state_ = TaskState::QR_SCAN;
      qr_scan_start_ = this->now();
      RCLCPP_INFO(this->get_logger(), "State -> QR_SCAN");
    }
    else if (sign_data == 5) {
      state_ = TaskState::GOING;
      is_back_ = false;
      sub_target_ = true;
      has_path_ = false;
      waypoint_index_ = 0;
      angular_pid_.reset();
      // Trigger Nav2 replanning if we have a goal
      if (has_goal_) {
        requestPlan();
      }
      RCLCPP_INFO(this->get_logger(), "State -> GOING (forward)");
    }
    else if (sign_data == -2) {
      state_ = TaskState::RETURN;
      is_back_ = true;
      has_path_ = false;
      waypoint_index_ = 0;
      angular_pid_.reset();
      if (has_goal_) {
        requestPlan();
      }
      RCLCPP_INFO(this->get_logger(), "State -> RETURN");
    }
  }

  // ─────────────────────────────────────────────
  //  B. Callbacks for path, goal, odom, lane cmd_vel
  // ─────────────────────────────────────────────

  void pathCallback(const nav_msgs::msg::Path::SharedPtr msg)
  {
    std::lock_guard<std::mutex> lock(mutex_);
    if (msg->poses.empty()) {
      has_path_ = false;
      return;
    }
    global_path_ = *msg;
    has_path_ = true;
    waypoint_index_ = 0;
    angular_pid_.reset();

    // Auto-transition IDLE -> GOING when receiving an external path
    // (e.g. publish_sine_path.py publishes /plan without /goal_pose)
    if (state_ == TaskState::IDLE || state_ == TaskState::RESET) {
      state_ = TaskState::GOING;
      is_back_ = false;
      sub_target_ = true;
      RCLCPP_INFO(this->get_logger(),
        "Auto State -> GOING (external path, %zu waypoints)",
        msg->poses.size());
    }
  }

  void goalCallback(const geometry_msgs::msg::PoseStamped::SharedPtr msg)
  {
    {
      std::lock_guard<std::mutex> lock(mutex_);
      goal_pose_ = *msg;
      has_goal_ = true;
      RCLCPP_INFO(this->get_logger(), "New goal: (%.2f, %.2f)",
                  msg->pose.position.x, msg->pose.position.y);

      // Auto-transition IDLE → GOING if not already active
      if (state_ == TaskState::IDLE || state_ == TaskState::RESET) {
        state_ = TaskState::GOING;
        is_back_ = false;
        sub_target_ = true;
        has_path_ = false;
        waypoint_index_ = 0;
        angular_pid_.reset();
        RCLCPP_INFO(this->get_logger(), "Auto State -> GOING (goal received)");
      }
    }
    // Request plan outside mutex to avoid blocking other callbacks
    requestPlan();
  }

  void odomCallback(const nav_msgs::msg::Odometry::SharedPtr msg)
  {
    std::lock_guard<std::mutex> lock(mutex_);
    cur_x_ = msg->pose.pose.position.x;
    cur_y_ = msg->pose.pose.position.y;

    double qx = msg->pose.pose.orientation.x;
    double qy = msg->pose.pose.orientation.y;
    double qz = msg->pose.pose.orientation.z;
    double qw = msg->pose.pose.orientation.w;
    cur_yaw_ = std::atan2(2.0 * (qw * qz + qx * qy),
                          1.0 - 2.0 * (qy * qy + qz * qz));
    has_odom_ = true;
  }

  void laneCmdCallback(const geometry_msgs::msg::Twist::SharedPtr msg)
  {
    std::lock_guard<std::mutex> lock(mutex_);
    lane_cmd_vel_ = *msg;
    got_lane_cmd_ = true;
  }

  // ─────────────────────────────────────────────
  //  C. ComputePathToPose action client
  // ─────────────────────────────────────────────

  void requestPlan()
  {
    if (!planner_client_->wait_for_action_server(std::chrono::seconds(2))) {
      RCLCPP_WARN(this->get_logger(),
                  "Planner action server (compute_path_to_pose) not available");
      return;
    }

    auto goal_msg = ComputePathToPose::Goal();
    goal_msg.goal = goal_pose_;
    goal_msg.planner_id = "GridBased";

    auto send_goal_options = rclcpp_action::Client<ComputePathToPose>::SendGoalOptions();
    send_goal_options.result_callback =
      [this](const GoalHandleComputePathToPose::WrappedResult& result) {
        if (result.code == rclcpp_action::ResultCode::SUCCEEDED) {
          RCLCPP_INFO(this->get_logger(), "Plan received: %zu waypoints",
                      result.result->path.poses.size());
        } else {
          RCLCPP_WARN(this->get_logger(), "Plan request failed");
        }
      };

    RCLCPP_INFO(this->get_logger(), "Requesting plan from Nav2 planner_server...");
    planner_client_->async_send_goal(goal_msg, send_goal_options);
  }

  // ─────────────────────────────────────────────
  //  D. Main Control Loop — cmd_vel arbitration
  // ─────────────────────────────────────────────

  void controlLoop()
  {
    std::lock_guard<std::mutex> lock(mutex_);

    // Priority 1: Task state control
    switch (state_)
    {
      case TaskState::IDLE:
        publishStop();
        return;

      case TaskState::RESET:
        publishStop();
        state_ = TaskState::IDLE;
        return;

      case TaskState::QR_SCAN:
      {
        auto elapsed = this->now() - qr_scan_start_;
        if (elapsed.seconds() < 1.0) {
          geometry_msgs::msg::Twist cmd;
          cmd.linear.x = qr_scan_speed_;
          cmd.angular.z = 0.0;
          cmd_vel_pub_->publish(cmd);
        }
        else if (elapsed.seconds() < qr_scan_duration_) {
          publishStop();
        }
        else {
          state_ = TaskState::GOING;
          RCLCPP_INFO(this->get_logger(), "QR scan complete, resuming");
        }
        return;
      }

      case TaskState::GOING:
      case TaskState::RETURN:
        break;
    }

    // Priority 2: Obstacle avoidance
    if (activate_avoid_) {
      geometry_msgs::msg::Twist cmd;
      cmd.linear.x = avoid_linear_speed_;
      cmd.angular.z = avoid_angular_z_;
      cmd_vel_pub_->publish(cmd);
      return;
    }

    // Priority 3: Lane following (sub_target_ == false)
    if (!sub_target_ && got_lane_cmd_) {
      cmd_vel_pub_->publish(lane_cmd_vel_);
      return;
    }

    // Priority 4: PID path following (path from Nav2 planner_server via /plan)
    if (sub_target_ && has_path_ && has_odom_) {
      if (waypoint_index_ >= static_cast<int>(global_path_.poses.size())) {
        RCLCPP_INFO(this->get_logger(), "Path complete - stopping");
        publishStop();
        has_path_ = false;
        return;
      }

      int target_idx = std::min(waypoint_index_ + guide_step_,
                                static_cast<int>(global_path_.poses.size()) - 1);
      const auto& target = global_path_.poses[target_idx];

      double dx = target.pose.position.x - cur_x_;
      double dy = target.pose.position.y - cur_y_;
      double distance_sq = dx * dx + dy * dy;

      if (distance_sq < arrive_square_) {
        waypoint_index_++;
        angular_pid_.reset();
      }

      double target_yaw = std::atan2(dy, dx);
      double heading_error = target_yaw - cur_yaw_;

      while (heading_error > M_PI)  heading_error -= 2.0 * M_PI;
      while (heading_error < -M_PI) heading_error += 2.0 * M_PI;

      double dt = 1.0 / this->get_parameter("control_rate").as_double();
      double angular_z = angular_pid_.update(heading_error, dt);

      geometry_msgs::msg::Twist cmd;
      cmd.linear.x = follow_linear_speed_;
      cmd.angular.z = angular_z;
      cmd_vel_pub_->publish(cmd);
      return;
    }

    // Priority 5: No path - stop and wait
    publishStop();
  }

  void publishStop()
  {
    geometry_msgs::msg::Twist cmd;
    cmd.linear.x = 0.0;
    cmd.angular.z = 0.0;
    cmd_vel_pub_->publish(cmd);
  }

  // ─────────────────────────────────────────────
  //  Members (order must match initializer list)
  // ─────────────────────────────────────────────

  TaskState state_;
  bool has_odom_, has_goal_;
  bool has_path_, got_lane_cmd_;
  bool is_back_, activate_avoid_;
  int waypoint_index_;
  int guide_step_;
  double follow_linear_speed_;
  bool sub_target_;
  double cmd_vel_linear_x_, cmd_vel_angular_z_;
  double cur_x_, cur_y_, cur_yaw_;
  double arrive_square_;
  past_control::PID angular_pid_;
  double avoid_linear_speed_;
  double avoid_angular_z_;
  double qr_scan_speed_;
  double qr_scan_duration_;
  rclcpp::Time qr_scan_start_;
  nav_msgs::msg::Path global_path_;
  geometry_msgs::msg::Twist lane_cmd_vel_;
  geometry_msgs::msg::PoseStamped goal_pose_;

  rclcpp::Subscription<origincar_msg::msg::Sign>::SharedPtr sign_sub_;
  rclcpp::Subscription<origincar_msg::msg::Sign>::SharedPtr sign_foxglove_sub_;
  rclcpp::Subscription<nav_msgs::msg::Path>::SharedPtr path_sub_;
  rclcpp::Subscription<geometry_msgs::msg::PoseStamped>::SharedPtr goal_sub_;
  rclcpp::Subscription<nav_msgs::msg::Odometry>::SharedPtr odom_sub_;
  rclcpp::Subscription<geometry_msgs::msg::Twist>::SharedPtr lane_cmd_sub_;
  rclcpp::Publisher<geometry_msgs::msg::Twist>::SharedPtr cmd_vel_pub_;
  rclcpp_action::Client<ComputePathToPose>::SharedPtr planner_client_;
  rclcpp::TimerBase::SharedPtr timer_;
  std::mutex mutex_;
};

int main(int argc, char* argv[])
{
  rclcpp::init(argc, argv);
  rclcpp::spin(std::make_shared<RacingOrchestrator>());
  rclcpp::shutdown();
  return 0;
}
