#!/usr/bin/env python3
"""Q_i candidate waypoint selector for Nav2 NavigateToPose.

用途
----
这个节点用于比赛/实车场景下的多路线候选导航。它会从同一个 JSON
目录中读取多个 ``*.json`` 文件。每个 JSON 文件表示一条完整路线，文件
中的 ``points`` 数组表示该路线按顺序经过的导航点。

多 JSON 编组规则
--------------
假设目录中存在多个文件：

  test1.json: p_0, p_1, p_2, ...
  test2.json: p_0, p_1, p_2, ...
  test3.json: p_0, p_1, p_2, ...

节点会把每个文件的第 i 个点组合成 Q_i：

  Q_0 = [test1[0], test2[0], test3[0], ...]
  Q_1 = [test1[1], test2[1], test3[1], ...]
  Q_2 = [test1[2], test2[2], test3[2], ...]

文件按自然顺序排序，因此 ``test2.json`` 会排在 ``test10.json`` 前面。
如果某条路线点数较少，缺失的索引会被跳过，不影响其他路线候选。

JSON 格式
---------
推荐格式为：

  {
    "points": [
      {
        "name": "goal_001",
        "yaw_degrees": 4.39,
        "odom": {
          "pose": {
            "pose": {
              "position": {"x": 1.48, "y": 0.18, "z": 0.0},
              "orientation": {"x": 0.0, "y": 0.0, "z": 0.03, "w": 0.99}
            }
          }
        }
      }
    ]
  }

同时也支持顶层直接是 list 的 JSON。坐标 ``x/y`` 必须与 ``goal_frame``
一致，例如 ``odom`` 或 ``odom_combined``。

候选点选择规则
--------------
1. 第一次优先选择第 0 条路线，也就是自然排序后的第一个 JSON 文件。
2. 到达某个 Q_i 后，进入 Q_{i+1} 时优先保持同一条路线。
3. 如果该候选点在全局 costmap 中不可用，则按路线顺序循环尝试下一个候选。
4. 如果某个 Q_i 中所有候选都不可用，则跳过该 Q，进入 Q_{i+1}。
5. 如果 Nav2 action 被拒绝、失败、取消或超时，则当前候选在本轮 Q 中
   被标记为失败，下一轮 tick 会尝试下一个候选。

不可用判断
---------
发送目标前会查询 ``/global_costmap/costmap``：

* 点在 costmap 外：不可用。
* costmap 数据异常：不可用。
* cost >= ``occupied_threshold``：不可用。
* cost = -1：由 ``treat_unknown_as_occupied`` 决定。

当前默认适配无静态地图场景：``treat_unknown_as_occupied=False``，因此未知
区域不会直接导致候选点被拒绝；真正障碍仍由 cost 阈值过滤。

关键参数
--------
* ``json_dir``: JSON 路线目录。
* ``goal_frame``: 下发 NavigateToPose 目标使用的坐标系，默认 ``odom``。
* ``costmap_topic``: 用于判断候选点占用状态的 costmap，默认
  ``/global_costmap/costmap``。
* ``navigate_action``: Nav2 action 名称，默认 ``/navigate_to_pose``。
* ``occupied_threshold``: costmap 占用阈值，默认 50。
* ``treat_unknown_as_occupied``: 是否把 cost=-1 视为不可通行，默认 False。
* ``goal_timeout_sec``: 单个候选目标超时时间，默认 180 秒。
* ``selection_period_sec``: 主循环周期，默认 0.5 秒。
"""

from __future__ import annotations

import json
import math
import re
import sys
import traceback
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Set

ROS_IMPORT_ERROR: Optional[BaseException] = None

try:
    import rclpy
    from geometry_msgs.msg import PoseStamped
    from nav2_msgs.action import NavigateToPose
    from nav_msgs.msg import OccupancyGrid
    from rclpy.action import ActionClient
    from rclpy.duration import Duration
    from rclpy.node import Node
    from rclpy.qos import QoSProfile
except BaseException as exc:  # pragma: no cover - used for remote runtime diagnosis
    ROS_IMPORT_ERROR = exc
    rclpy = None  # type: ignore[assignment]
    PoseStamped = object  # type: ignore[assignment]
    NavigateToPose = object  # type: ignore[assignment]
    OccupancyGrid = object  # type: ignore[assignment]
    ActionClient = object  # type: ignore[assignment]
    Duration = object  # type: ignore[assignment]
    Node = object  # type: ignore[assignment]
    QoSProfile = object  # type: ignore[assignment]


@dataclass(frozen=True)
class CandidateWaypoint:
    q_index: int
    route_index: int
    file_name: str
    point_name: str
    x: float
    y: float
    yaw: float


def yaw_from_quaternion(q: Dict[str, float]) -> float:
    x = float(q.get("x", 0.0))
    y = float(q.get("y", 0.0))
    z = float(q.get("z", 0.0))
    w = float(q.get("w", 1.0))
    return math.atan2(
        2.0 * (w * z + x * y),
        1.0 - 2.0 * (y * y + z * z),
    )


def pose_from_json_point(point: Dict, q_index: int, route_index: int, file_name: str) -> CandidateWaypoint:
    odom = point.get("odom", {})
    pose = odom.get("pose", {}).get("pose", {})
    position = pose.get("position", {})
    orientation = pose.get("orientation", {})

    x = float(position["x"])
    y = float(position["y"])
    yaw_degrees = point.get("yaw_degrees", None)
    if yaw_degrees is None:
        yaw = yaw_from_quaternion(orientation)
    else:
        yaw = math.radians(float(yaw_degrees))

    return CandidateWaypoint(
        q_index=q_index,
        route_index=route_index,
        file_name=file_name,
        point_name=str(point.get("name", f"{Path(file_name).stem}_{q_index}")),
        x=x,
        y=y,
        yaw=yaw,
    )


def natural_json_sort_key(path: Path) -> List[object]:
    return [
        int(part) if part.isdigit() else part
        for part in re.split(r"(\d+)", path.name.lower())
    ]


def load_points_from_json(path: Path) -> List[Dict]:
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception as exc:
        raise RuntimeError(f"failed to parse JSON file {path}: {exc}") from exc

    if isinstance(data, list):
        points = data
    elif isinstance(data, dict):
        points = data.get("points", [])
    else:
        raise RuntimeError(
            f"{path} must contain either a top-level list or a dict field named 'points'"
        )

    if not isinstance(points, list):
        raise RuntimeError(f"{path} does not contain a list field named 'points'")
    return points


def load_q_groups(json_dir: Path) -> tuple[List[List[CandidateWaypoint]], List[str], List[int]]:
    if not json_dir.exists():
        raise RuntimeError(f"JSON directory does not exist: {json_dir}")
    if not json_dir.is_dir():
        raise RuntimeError(f"JSON path is not a directory: {json_dir}")

    files = sorted(json_dir.glob("*.json"), key=natural_json_sort_key)
    if not files:
        raise RuntimeError(f"no JSON files found in {json_dir}")

    routes: List[List[Dict]] = []
    for path in files:
        routes.append(load_points_from_json(path))

    max_points = max(len(route) for route in routes)
    groups: List[List[CandidateWaypoint]] = []
    for q_index in range(max_points):
        group: List[CandidateWaypoint] = []
        for route_index, points in enumerate(routes):
            if q_index >= len(points):
                continue
            group.append(
                pose_from_json_point(
                    points[q_index],
                    q_index=q_index,
                    route_index=route_index,
                    file_name=files[route_index].name,
                )
            )
        groups.append(group)
    return groups, [path.name for path in files], [len(route) for route in routes]


class QWaypointSelector(Node):
    def __init__(self) -> None:
        super().__init__("q_waypoint_selector")

        self._declare_parameter_if_needed("use_sim_time", False)
        self._declare_parameter_if_needed("json_dir", "")
        self._declare_parameter_if_needed("goal_frame", "odom")
        self._declare_parameter_if_needed("costmap_topic", "/global_costmap/costmap")
        self._declare_parameter_if_needed("navigate_action", "/navigate_to_pose")
        self._declare_parameter_if_needed("occupied_threshold", 50)
        self._declare_parameter_if_needed("treat_unknown_as_occupied", False)
        self._declare_parameter_if_needed("goal_timeout_sec", 180.0)
        self._declare_parameter_if_needed("selection_period_sec", 0.5)

        json_dir = Path(str(self.get_parameter("json_dir").value)).expanduser()
        if not json_dir.is_absolute():
            json_dir = Path.cwd() / json_dir
        self.goal_frame = str(self.get_parameter("goal_frame").value)
        self.costmap_topic = str(self.get_parameter("costmap_topic").value)
        self.navigate_action = str(self.get_parameter("navigate_action").value)
        self.occupied_threshold = int(self.get_parameter("occupied_threshold").value)
        self.treat_unknown_as_occupied = bool(
            self.get_parameter("treat_unknown_as_occupied").value
        )
        self.goal_timeout_sec = max(
            1.0, float(self.get_parameter("goal_timeout_sec").value)
        )
        self.selection_period_sec = max(
            0.1, float(self.get_parameter("selection_period_sec").value)
        )

        self.q_groups, self.route_file_names, self.route_point_counts = load_q_groups(json_dir)
        self.current_q_index = 0
        self.preferred_route_index = 0
        self.rejected_routes_for_current_q: Set[int] = set()
        self.active_candidate: Optional[CandidateWaypoint] = None
        self.active_goal_handle = None
        self.active_goal_start_time = None
        self.active_goal_token: Optional[int] = None
        self.next_goal_token = 0
        self.done = False
        self.latest_costmap: Optional[OccupancyGrid] = None
        self.last_costmap_wait_log_ns = 0

        qos = QoSProfile(depth=1)
        self.costmap_sub = self.create_subscription(
            OccupancyGrid, self.costmap_topic, self._costmap_callback, qos
        )
        self.action_client = ActionClient(self, NavigateToPose, self.navigate_action)
        self.timer = self.create_timer(self.selection_period_sec, self._tick)

        self.get_logger().info(
            "loaded Q waypoint groups "
            f"json_dir={json_dir} groups={len(self.q_groups)} "
            f"routes={len(self.route_file_names)} "
            f"goal_frame={self.goal_frame} costmap_topic={self.costmap_topic}"
        )
        route_summary = ", ".join(
            f"{name}:{count}"
            for name, count in zip(self.route_file_names, self.route_point_counts)
        )
        self.get_logger().info(f"loaded JSON waypoint routes {route_summary}")

    def _declare_parameter_if_needed(self, name: str, default_value) -> None:
        if not self.has_parameter(name):
            self.declare_parameter(name, default_value)

    def _costmap_callback(self, msg: OccupancyGrid) -> None:
        self.latest_costmap = msg

    def _log_waiting_for_costmap(self) -> None:
        now_ns = self.get_clock().now().nanoseconds
        if now_ns - self.last_costmap_wait_log_ns < 2_000_000_000:
            return
        self.last_costmap_wait_log_ns = now_ns

        publishers = self.get_publishers_info_by_topic(self.costmap_topic)
        if not publishers:
            self.get_logger().warn(
                f"waiting for costmap {self.costmap_topic}; publisher_count=0; "
                "no goal will be sent yet"
            )
            return

        details = []
        for info in publishers[:3]:
            details.append(
                f"{info.node_namespace}/{info.node_name} "
                f"type={info.topic_type} qos={info.qos_profile}"
            )
        more = "" if len(publishers) <= 3 else f" ... +{len(publishers) - 3} more"
        self.get_logger().warn(
            f"waiting for costmap {self.costmap_topic}; "
            f"publisher_count={len(publishers)} publishers={details}{more}; "
            "no goal will be sent yet"
        )

    def _tick(self) -> None:
        if self.done:
            return

        if self.current_q_index >= len(self.q_groups):
            self.done = True
            self.get_logger().info("all Q waypoint groups completed")
            return

        if not self.action_client.server_is_ready():
            self.get_logger().info(
                f"waiting for NavigateToPose action server {self.navigate_action}"
            )
            self.action_client.wait_for_server(timeout_sec=0.1)
            return

        if self.latest_costmap is None:
            self._log_waiting_for_costmap()
            return

        if self.active_goal_handle is not None:
            if self._active_goal_timed_out():
                self.get_logger().warn(
                    "active goal timed out; cancelling and trying next candidate "
                    f"q={self.current_q_index} candidate={self._candidate_label(self.active_candidate)}"
                )
                cancel_future = self.active_goal_handle.cancel_goal_async()
                cancel_future.add_done_callback(lambda _: None)
                self._reject_active_candidate()
            return

        candidate = self._select_candidate()
        if candidate is None:
            self.get_logger().warn(
                f"all candidates in Q_{self.current_q_index} are occupied/unusable; skipping group"
            )
            self.current_q_index += 1
            self.rejected_routes_for_current_q.clear()
            return

        self._send_goal(candidate)

    def _select_candidate(self) -> Optional[CandidateWaypoint]:
        group = self.q_groups[self.current_q_index]
        if not group:
            return None

        route_to_candidate = {candidate.route_index: candidate for candidate in group}
        ordered_route_indices = self._cyclic_route_order(
            sorted(route_to_candidate.keys()), self.preferred_route_index
        )

        for route_index in ordered_route_indices:
            if route_index in self.rejected_routes_for_current_q:
                continue
            candidate = route_to_candidate[route_index]
            occupied, reason = self._candidate_occupied(candidate)
            if occupied:
                self.rejected_routes_for_current_q.add(route_index)
                self.get_logger().info(
                    f"skip occupied candidate Q_{candidate.q_index} "
                    f"{self._candidate_label(candidate)} reason={reason}"
                )
                continue
            return candidate
        return None

    @staticmethod
    def _cyclic_route_order(route_indices: Sequence[int], preferred: int) -> List[int]:
        if not route_indices:
            return []
        if preferred not in route_indices:
            preferred = route_indices[0]
        start = route_indices.index(preferred)
        return list(route_indices[start:]) + list(route_indices[:start])

    def _candidate_occupied(self, candidate: CandidateWaypoint) -> tuple[bool, str]:
        costmap = self.latest_costmap
        if costmap is None:
            return True, "no_costmap"

        origin = costmap.info.origin
        resolution = costmap.info.resolution
        width = int(costmap.info.width)
        height = int(costmap.info.height)
        if resolution <= 0.0 or width <= 0 or height <= 0:
            return True, "invalid_costmap"

        origin_yaw = yaw_from_quaternion(
            {
                "x": origin.orientation.x,
                "y": origin.orientation.y,
                "z": origin.orientation.z,
                "w": origin.orientation.w,
            }
        )
        dx = candidate.x - origin.position.x
        dy = candidate.y - origin.position.y
        cos_yaw = math.cos(-origin_yaw)
        sin_yaw = math.sin(-origin_yaw)
        local_x = cos_yaw * dx - sin_yaw * dy
        local_y = sin_yaw * dx + cos_yaw * dy
        mx = int(math.floor(local_x / resolution))
        my = int(math.floor(local_y / resolution))

        if mx < 0 or my < 0 or mx >= width or my >= height:
            return True, f"outside_costmap cell=({mx},{my})"

        index = my * width + mx
        if index < 0 or index >= len(costmap.data):
            return True, "costmap_data_too_short"

        cost = int(costmap.data[index])
        if cost < 0:
            return self.treat_unknown_as_occupied, f"unknown cost={cost}"
        if cost >= self.occupied_threshold:
            return True, f"occupied cost={cost}"
        return False, f"free cost={cost}"

    def _send_goal(self, candidate: CandidateWaypoint) -> None:
        goal_msg = NavigateToPose.Goal()
        goal_msg.pose = self._to_pose_stamped(candidate)
        self.next_goal_token += 1
        token = self.next_goal_token
        self.active_candidate = candidate
        self.active_goal_start_time = self.get_clock().now()
        self.active_goal_token = token

        self.get_logger().info(
            f"send Q_{candidate.q_index} candidate {self._candidate_label(candidate)} "
            f"pose=({candidate.x:.3f},{candidate.y:.3f},{candidate.yaw:.3f})"
        )
        send_future = self.action_client.send_goal_async(goal_msg)
        send_future.add_done_callback(
            lambda future, goal_token=token: self._goal_response_callback(future, goal_token)
        )

    def _to_pose_stamped(self, candidate: CandidateWaypoint) -> PoseStamped:
        pose = PoseStamped()
        pose.header.frame_id = self.goal_frame
        pose.header.stamp = self.get_clock().now().to_msg()
        pose.pose.position.x = candidate.x
        pose.pose.position.y = candidate.y
        pose.pose.position.z = 0.0
        pose.pose.orientation.z = math.sin(candidate.yaw * 0.5)
        pose.pose.orientation.w = math.cos(candidate.yaw * 0.5)
        return pose

    def _goal_response_callback(self, future, goal_token: int) -> None:
        if goal_token != self.active_goal_token:
            return
        try:
            goal_handle = future.result()
        except Exception as exc:
            self.get_logger().error(
                "NavigateToPose send_goal_async failed; try next candidate: "
                f"{type(exc).__name__}: {exc}"
            )
            self._reject_active_candidate()
            return
        if not goal_handle.accepted:
            self.get_logger().warn(
                f"goal rejected by Nav2 {self._candidate_label(self.active_candidate)}"
            )
            self._reject_active_candidate()
            return

        self.active_goal_handle = goal_handle
        result_future = goal_handle.get_result_async()
        result_future.add_done_callback(
            lambda future, token=goal_token: self._goal_result_callback(future, token)
        )

    def _goal_result_callback(self, future, goal_token: int) -> None:
        if goal_token != self.active_goal_token:
            return
        try:
            result = future.result()
        except Exception as exc:
            self.get_logger().error(
                "NavigateToPose result failed; try next candidate: "
                f"{type(exc).__name__}: {exc}"
            )
            self._reject_active_candidate()
            return
        candidate = self.active_candidate
        status = int(result.status)

        if status == 4:
            if candidate is not None:
                self.preferred_route_index = candidate.route_index
                self.get_logger().info(
                    f"reached Q_{candidate.q_index} {self._candidate_label(candidate)}; "
                    f"advance to Q_{candidate.q_index + 1}"
                )
            self.current_q_index += 1
            self.rejected_routes_for_current_q.clear()
            self.active_candidate = None
            self.active_goal_handle = None
            self.active_goal_start_time = None
            self.active_goal_token = None
            return

        self.get_logger().warn(
            f"goal failed status={status}; try next candidate "
            f"{self._candidate_label(candidate)}"
        )
        self._reject_active_candidate()

    def _active_goal_timed_out(self) -> bool:
        if self.active_goal_start_time is None:
            return False
        elapsed = self.get_clock().now() - self.active_goal_start_time
        return elapsed > Duration(seconds=self.goal_timeout_sec)

    def _reject_active_candidate(self) -> None:
        if self.active_candidate is not None:
            self.rejected_routes_for_current_q.add(self.active_candidate.route_index)
        self.active_candidate = None
        self.active_goal_handle = None
        self.active_goal_start_time = None
        self.active_goal_token = None

    @staticmethod
    def _candidate_label(candidate: Optional[CandidateWaypoint]) -> str:
        if candidate is None:
            return "none"
        return (
            f"route={candidate.route_index} file={candidate.file_name} "
            f"name={candidate.point_name}"
        )


def main() -> int:
    if ROS_IMPORT_ERROR is not None:
        print(
            "FATAL q_waypoint_selector: failed to import ROS/Nav2 Python modules.",
            file=sys.stderr,
        )
        traceback.print_exception(
            type(ROS_IMPORT_ERROR),
            ROS_IMPORT_ERROR,
            ROS_IMPORT_ERROR.__traceback__,
            file=sys.stderr,
        )
        print(
            "Hint: source /opt/ros/humble/setup.bash and install/source Nav2 "
            "packages, especially nav2_msgs.",
            file=sys.stderr,
        )
        return 1

    rclpy.init()
    node = None
    try:
        node = QWaypointSelector()
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    except BaseException:
        print("FATAL q_waypoint_selector: unhandled exception", file=sys.stderr)
        traceback.print_exc(file=sys.stderr)
        return 1
    finally:
        if node is not None:
            node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
