#!/usr/bin/env python3
"""Lightweight tag-localized route follower.

This node intentionally does not start Nav2, costmaps, planners, controllers, or
lidar. It exposes the same action names Apollo's navigation blocks already use
and drives directly from the tag-localized map -> base_footprint TF.

When a GeoJSON route graph is provided (via route_graph_file param or
APOLLO_GRAPH_FILE env var), navigate_to_pose and follow_waypoints goals are
planned through the graph using Dijkstra before executing, so the robot follows
collision-free paths around corners instead of driving straight-line shortcuts.
"""

import heapq
import json
import math
import os
import time
from typing import Dict, Iterable, List, Optional, Tuple

import rclpy
from geometry_msgs.msg import Pose, PoseStamped, Twist
from nav2_msgs.action import FollowWaypoints, NavigateToPose
from nav_msgs.msg import Path
from rclpy.action import ActionServer, CancelResponse, GoalResponse
from rclpy.executors import MultiThreadedExecutor
from rclpy.node import Node
from rclpy.qos import DurabilityPolicy, HistoryPolicy, QoSProfile
from std_msgs.msg import Int32, String
import tf2_ros


def clamp(value: float, low: float, high: float) -> float:
    return max(low, min(high, value))


# ---------------------------------------------------------------------------
# GeoJSON route graph with Dijkstra path planning
# ---------------------------------------------------------------------------

class RouteGraph:
    """Parse a GeoJSON route graph produced by generator_v2.py and plan paths."""

    def __init__(self, path: str):
        self.nodes: Dict[int, Tuple[float, float]] = {}   # id -> (x, y)
        self.adj: Dict[int, List[Tuple[int, float]]] = {}  # id -> [(neighbor, dist)]
        self._load(path)

    def _load(self, path: str) -> None:
        with open(path, "r") as f:
            data = json.load(f)

        edges_raw = []
        for feature in data.get("features", []):
            props = feature.get("properties", {})
            geom = feature.get("geometry", {})
            if geom.get("type") == "Point":
                nid = int(props["id"])
                coords = geom["coordinates"]
                self.nodes[nid] = (float(coords[0]), float(coords[1]))
                self.adj[nid] = []
            elif geom.get("type") == "LineString":
                edges_raw.append((
                    int(props["startid"]),
                    int(props["endid"]),
                    float(props["length_m"]),
                ))

        for startid, endid, dist in edges_raw:
            if startid in self.adj and endid in self.nodes:
                self.adj[startid].append((endid, dist))
            if endid in self.adj and startid in self.nodes:
                self.adj[endid].append((startid, dist))

    def _nearest(self, x: float, y: float) -> Optional[int]:
        best_id, best_d = None, float("inf")
        for nid, (nx, ny) in self.nodes.items():
            d = math.hypot(nx - x, ny - y)
            if d < best_d:
                best_d = d
                best_id = nid
        return best_id

    def plan(
        self, start_x: float, start_y: float, goal_x: float, goal_y: float
    ) -> Optional[List[Tuple[float, float]]]:
        """Return ordered (x, y) graph waypoints from start to goal, or None on failure."""
        if not self.nodes:
            return None
        start_node = self._nearest(start_x, start_y)
        goal_node = self._nearest(goal_x, goal_y)
        if start_node is None or goal_node is None:
            return None
        if start_node == goal_node:
            return []

        dist: Dict[int, float] = {start_node: 0.0}
        prev: Dict[int, int] = {}
        heap = [(0.0, start_node)]
        while heap:
            d, u = heapq.heappop(heap)
            if d > dist.get(u, float("inf")):
                continue
            if u == goal_node:
                break
            for v, w in self.adj.get(u, []):
                nd = d + w
                if nd < dist.get(v, float("inf")):
                    dist[v] = nd
                    prev[v] = u
                    heapq.heappush(heap, (nd, v))

        if goal_node not in prev:
            return None  # disconnected graph

        path: List[Tuple[float, float]] = []
        node = goal_node
        while node in prev:
            path.append(self.nodes[node])
            node = prev[node]
        path.reverse()
        return path


def yaw_from_quaternion(q) -> float:
    siny_cosp = 2.0 * (q.w * q.z + q.x * q.y)
    cosy_cosp = 1.0 - 2.0 * (q.y * q.y + q.z * q.z)
    return math.atan2(siny_cosp, cosy_cosp)


def normalize_angle(angle: float) -> float:
    while angle > math.pi:
        angle -= 2.0 * math.pi
    while angle < -math.pi:
        angle += 2.0 * math.pi
    return angle


class TagRouteFollowerNode(Node):
    def __init__(self) -> None:
        super().__init__("tag_route_follower_node")
        self.map_frame = str(self.declare_parameter(
            "map_frame", os.environ.get("APOLLO_TAG_MAP_FRAME", "map")
        ).value)
        self.base_frame = str(self.declare_parameter(
            "base_frame", os.environ.get("APOLLO_TAG_BASE_FRAME", "base_footprint")
        ).value)
        self.xy_tolerance = float(self.declare_parameter(
            "xy_tolerance", float(os.environ.get("APOLLO_TAG_XY_TOLERANCE", "0.12"))
        ).value)
        self.graph_xy_tolerance = float(self.declare_parameter(
            "graph_xy_tolerance", float(os.environ.get("APOLLO_TAG_GRAPH_XY_TOLERANCE", "0.035"))
        ).value)
        self.yaw_tolerance = float(self.declare_parameter(
            "yaw_tolerance", float(os.environ.get("APOLLO_TAG_YAW_TOLERANCE", "0.20"))
        ).value)
        self.heading_stop_threshold = float(self.declare_parameter(
            "heading_stop_threshold", float(os.environ.get("APOLLO_TAG_HEADING_STOP_THRESHOLD", "0.35"))
        ).value)
        self.max_graph_segment_length = float(self.declare_parameter(
            "max_graph_segment_length", float(os.environ.get("APOLLO_TAG_MAX_GRAPH_SEGMENT_LENGTH", "0.08"))
        ).value)
        self.max_linear = float(self.declare_parameter(
            "max_linear", float(os.environ.get("APOLLO_TAG_MAX_LINEAR", "0.18"))
        ).value)
        self.max_angular = float(self.declare_parameter(
            "max_angular", float(os.environ.get("APOLLO_TAG_MAX_ANGULAR", "0.7"))
        ).value)
        self.k_linear = float(self.declare_parameter(
            "k_linear", float(os.environ.get("APOLLO_TAG_K_LINEAR", "0.7"))
        ).value)
        self.k_angular = float(self.declare_parameter(
            "k_angular", float(os.environ.get("APOLLO_TAG_K_ANGULAR", "1.8"))
        ).value)
        self.control_hz = float(self.declare_parameter(
            "control_hz", float(os.environ.get("APOLLO_TAG_CONTROL_HZ", "20.0"))
        ).value)
        self.tf_timeout = float(self.declare_parameter(
            "tf_timeout", float(os.environ.get("APOLLO_TAG_TF_TIMEOUT", "0.5"))
        ).value)
        self.require_initial_tag_fix = bool(self.declare_parameter(
            "require_initial_tag_fix",
            os.environ.get("APOLLO_TAG_REQUIRE_INITIAL_FIX", "1").strip().lower() not in {"0", "false", "no"}
        ).value)
        self.startup_tag_timeout = float(self.declare_parameter(
            "startup_tag_timeout",
            float(os.environ.get("APOLLO_TAG_STARTUP_FIX_TIMEOUT", "12.0"))
        ).value)
        self.camera_startup_timeout = float(self.declare_parameter(
            "camera_startup_timeout",
            float(os.environ.get("APOLLO_TAG_CAMERA_STARTUP_TIMEOUT", "120.0"))
        ).value)

        graph_file = str(self.declare_parameter(
            "route_graph_file",
            os.environ.get("APOLLO_GRAPH_FILE", ""),
        ).value).strip()
        self.route_graph: Optional[RouteGraph] = None
        if graph_file and os.path.isfile(graph_file):
            try:
                self.route_graph = RouteGraph(graph_file)
                self.get_logger().info(
                    f"route graph loaded: {graph_file} "
                    f"({len(self.route_graph.nodes)} nodes)"
                )
            except Exception as exc:
                self.get_logger().error(f"failed to load route graph {graph_file}: {exc}")
        else:
            if graph_file:
                self.get_logger().warn(f"route_graph_file not found: {graph_file!r}")
            else:
                self.get_logger().warn(
                    "No route graph file configured (route_graph_file param / "
                    "APOLLO_GRAPH_FILE env). Driving direct — robot may cut corners."
                )

        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer, self)
        self.cmd_pub = self.create_publisher(Twist, "/cmd_vel", 10)

        # Route visualization: latch the planned path so the UI gets the current
        # path immediately when it (re)subscribes; the active index is a live
        # pointer into that path (g0, g1, ... matching the navigate log labels).
        latched = QoSProfile(
            depth=1,
            history=HistoryPolicy.KEEP_LAST,
            durability=DurabilityPolicy.TRANSIENT_LOCAL,
        )
        self.planned_path_pub = self.create_publisher(
            Path, "/apollo/tag_route/planned_path", latched
        )
        self.active_node_pub = self.create_publisher(
            Int32, "/apollo/tag_route/active_node", latched
        )

        self.localization_source = "unknown"
        self.latest_vision_time: Optional[float] = None
        self.localization_sub = self.create_subscription(
            String, "/apollo/tag_localization/source", self.on_localization_source, 10
        )
        self.vision_sub = self.create_subscription(
            String, "/apollo/vision/output", self.on_vision_output, 10
        )

        self.nav_server = ActionServer(
            self,
            NavigateToPose,
            "navigate_to_pose",
            execute_callback=self.execute_navigate_to_pose,
            goal_callback=self.accept_goal,
            cancel_callback=self.accept_cancel,
        )
        self.waypoint_server = ActionServer(
            self,
            FollowWaypoints,
            "follow_waypoints",
            execute_callback=self.execute_follow_waypoints,
            goal_callback=self.accept_goal,
            cancel_callback=self.accept_cancel,
        )

        self.get_logger().info(
            "tag_route_follower_node ready: "
            f"frames={self.map_frame}->{self.base_frame} "
            f"xy_tol={self.xy_tolerance:.2f} graph_xy_tol={self.graph_xy_tolerance:.3f} "
            f"yaw_tol={self.yaw_tolerance:.2f} heading_stop={self.heading_stop_threshold:.2f} "
            f"max_graph_segment={self.max_graph_segment_length:.2f} "
            f"max_linear={self.max_linear:.2f} max_angular={self.max_angular:.2f} "
            f"require_initial_tag_fix={self.require_initial_tag_fix}"
        )
        if self.route_graph is None:
            self.get_logger().warn(
                "Tag route follower has no route graph and no lidar obstacle avoidance; "
                "keep tag navigation slow and supervise the robot."
            )
        else:
            self.get_logger().warn(
                "Tag route follower has no lidar obstacle avoidance; "
                "using route graph for path planning around walls."
            )

    def accept_goal(self, _goal_request) -> GoalResponse:
        return GoalResponse.ACCEPT

    def accept_cancel(self, _goal_handle) -> CancelResponse:
        self.stop()
        return CancelResponse.ACCEPT

    def on_localization_source(self, msg: String) -> None:
        self.localization_source = msg.data or "unknown"

    def on_vision_output(self, _msg: String) -> None:
        self.latest_vision_time = time.monotonic()

    def stop(self) -> None:
        try:
            self.cmd_pub.publish(Twist())
        except Exception:
            pass

    def publish_planned_path(self, nodes: List[Tuple[float, float]]) -> None:
        """Publish the ordered route nodes (map frame) for UI visualization."""
        msg = Path()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = self.map_frame
        for nx, ny in nodes:
            ps = PoseStamped()
            ps.header = msg.header
            ps.pose.position.x = float(nx)
            ps.pose.position.y = float(ny)
            ps.pose.orientation.w = 1.0
            msg.poses.append(ps)
        try:
            self.planned_path_pub.publish(msg)
        except Exception:
            pass

    def publish_active_node(self, index: int) -> None:
        """Publish which node along the planned path is the current target."""
        try:
            self.active_node_pub.publish(Int32(data=int(index)))
        except Exception:
            pass

    def clear_route(self) -> None:
        """Clear the planned path / active node so a finished route leaves the map."""
        self.publish_planned_path([])
        self.publish_active_node(-1)

    def wait_for_initial_tag_fix(self) -> bool:
        if not self.require_initial_tag_fix:
            return True
        if self.localization_source.startswith("tag:"):
            return True

        self.get_logger().info(
            f"Waiting for camera frames up to {self.camera_startup_timeout:.1f}s, then "
            f"startup AprilTag localization fix up to {self.startup_tag_timeout:.1f}s..."
        )
        start = time.monotonic()
        tag_wait_start: Optional[float] = None
        last_log = 0
        while rclpy.ok():
            if self.localization_source.startswith("tag:"):
                self.get_logger().info(f"Startup localization fix ready: {self.localization_source}")
                return True

            now = time.monotonic()
            if self.latest_vision_time is not None:
                if tag_wait_start is None:
                    tag_wait_start = now
                    last_log = 0
                    self.get_logger().info(
                        "Camera frames are flowing; waiting for mapped AprilTag localization fix..."
                    )
                deadline_start = tag_wait_start
                deadline = self.startup_tag_timeout
                wait_label = "AprilTag fix"
            else:
                deadline_start = start
                deadline = self.camera_startup_timeout
                wait_label = "camera/vision frames"

            if now - deadline_start >= deadline:
                self.get_logger().error(
                    f"No startup {wait_label} after {deadline:.1f}s; "
                    f"current source={self.localization_source}. Not driving."
                )
                return False

            elapsed = int(now - deadline_start)
            if elapsed > 0 and elapsed % 3 == 0 and elapsed != last_log:
                self.get_logger().info(
                    f"Still waiting for {wait_label}; current source={self.localization_source}"
                )
                last_log = elapsed
            time.sleep(0.05)

        return False

    def current_pose(self) -> Optional[Tuple[float, float, float]]:
        try:
            tf = self.tf_buffer.lookup_transform(
                self.map_frame,
                self.base_frame,
                rclpy.time.Time(),
                timeout=rclpy.duration.Duration(seconds=self.tf_timeout),
            )
        except Exception as exc:
            self.get_logger().warn(f"tag route waiting for localized TF: {exc}")
            return None
        t = tf.transform.translation
        yaw = yaw_from_quaternion(tf.transform.rotation)
        return float(t.x), float(t.y), yaw

    def densify_graph_waypoints(
        self,
        start: Tuple[float, float],
        waypoints: List[Tuple[float, float]],
    ) -> List[Tuple[float, float]]:
        max_len = max(0.01, self.max_graph_segment_length)
        dense: List[Tuple[float, float]] = []
        previous = start
        for waypoint in waypoints:
            dx = waypoint[0] - previous[0]
            dy = waypoint[1] - previous[1]
            distance = math.hypot(dx, dy)
            steps = max(1, int(math.ceil(distance / max_len)))
            for step in range(1, steps + 1):
                t = step / steps
                dense.append((previous[0] + dx * t, previous[1] + dy * t))
            previous = waypoint
        return dense

    def drive_to_pose(
        self,
        goal_handle,
        pose: Pose,
        label: str,
        skip_yaw_align: bool = False,
        xy_tolerance: Optional[float] = None,
        heading_stop_threshold: Optional[float] = None,
    ) -> bool:
        goal_x = float(pose.position.x)
        goal_y = float(pose.position.y)
        goal_yaw = yaw_from_quaternion(pose.orientation)
        xy_tol = self.xy_tolerance if xy_tolerance is None else max(0.0, xy_tolerance)
        stop_threshold = 1.25 if heading_stop_threshold is None else max(0.0, heading_stop_threshold)
        period = 1.0 / max(1.0, self.control_hz)

        self.get_logger().info(
            f"{label}: driving to ({goal_x:.3f}, {goal_y:.3f}, yaw={goal_yaw:.3f}, "
            f"xy_tol={xy_tol:.3f}, heading_stop={stop_threshold:.3f})"
        )

        while rclpy.ok():
            if goal_handle.is_cancel_requested:
                self.stop()
                goal_handle.canceled()
                self.get_logger().warn(f"{label}: canceled")
                return False

            current = self.current_pose()
            if current is None:
                self.stop()
                time.sleep(period)
                continue

            if self.latest_vision_time is not None:
                vision_age = time.monotonic() - self.latest_vision_time
                if vision_age > 3.0:
                    self.get_logger().warn(
                        f"{label}: no tag frame for {vision_age:.1f}s — navigating on odometry only",
                        throttle_duration_sec=5.0,
                    )

            x, y, yaw = current
            dx = goal_x - x
            dy = goal_y - y
            distance = math.hypot(dx, dy)
            heading = math.atan2(dy, dx)
            heading_error = normalize_angle(heading - yaw)
            yaw_error = normalize_angle(goal_yaw - yaw)

            if distance <= xy_tol:
                if skip_yaw_align or abs(yaw_error) <= self.yaw_tolerance:
                    self.stop()
                    self.get_logger().info(f"{label}: goal reached")
                    return True
                twist = Twist()
                twist.angular.z = clamp(self.k_angular * yaw_error, -self.max_angular, self.max_angular)
                self.cmd_pub.publish(twist)
                time.sleep(period)
                continue

            twist = Twist()
            if abs(heading_error) > stop_threshold:
                twist.linear.x = 0.0
            else:
                twist.linear.x = clamp(self.k_linear * distance, 0.04, self.max_linear)
                twist.linear.x *= max(0.2, math.cos(heading_error))
            twist.angular.z = clamp(self.k_angular * heading_error, -self.max_angular, self.max_angular)
            self.cmd_pub.publish(twist)
            time.sleep(period)

        self.stop()
        return False

    def plan_and_drive(self, goal_handle, pose: Pose, label: str) -> bool:
        """Drive to pose via the route graph when available, otherwise drive direct."""
        if self.route_graph is None:
            return self.drive_to_pose(goal_handle, pose, label)

        current = self.current_pose()
        if current is None:
            # TF chain (map -> odom -> base) can lag a few hundred ms behind the
            # localization source at startup. Retry briefly so a goal that arrives
            # right after launch still plans through the route graph instead of
            # demoting to direct drive (which never publishes the planned path).
            deadline = time.monotonic() + 3.0
            while current is None and time.monotonic() < deadline and rclpy.ok():
                if goal_handle.is_cancel_requested:
                    self.stop()
                    goal_handle.canceled()
                    return False
                time.sleep(0.1)
                current = self.current_pose()
        if current is None:
            self.get_logger().warn(f"{label}: no localization yet — driving direct")
            return self.drive_to_pose(goal_handle, pose, label)

        cx, cy, _ = current
        gx, gy = float(pose.position.x), float(pose.position.y)
        waypoints = self.route_graph.plan(cx, cy, gx, gy)

        if waypoints is None:
            self.get_logger().warn(f"{label}: graph has no path to goal — driving direct")
            return self.drive_to_pose(goal_handle, pose, label)

        self.get_logger().info(
            f"{label}: graph path: {len(waypoints)} intermediate nodes → ({gx:.3f}, {gy:.3f})"
        )
        self.get_logger().info(
            f"{label}: tracking {len(waypoints)} graph nodes "
            f"(densified to <= {self.max_graph_segment_length:.2f}m per step)"
        )
        # Publish the planner's actual graph nodes (real GeoJSON node coords) plus
        # the final goal so the UI draws the true route nodes — not the densified
        # sub-targets. We still densify each segment internally for smooth motion,
        # but the active index reported tracks the GeoJSON node being approached.
        self.publish_planned_path(list(waypoints) + [(gx, gy)])
        prev = (cx, cy)
        for wi, (wx, wy) in enumerate(waypoints):
            if goal_handle.is_cancel_requested:
                self.stop()
                goal_handle.canceled()
                return False
            self.publish_active_node(wi)
            for sx, sy in self.densify_graph_waypoints(prev, [(wx, wy)]):
                if goal_handle.is_cancel_requested:
                    self.stop()
                    goal_handle.canceled()
                    return False
                sub_pose = Pose()
                sub_pose.position.x = sx
                sub_pose.position.y = sy
                sub_pose.orientation.w = 1.0
                if not self.drive_to_pose(
                    goal_handle,
                    sub_pose,
                    f"{label}[g{wi}]",
                    skip_yaw_align=True,
                    xy_tolerance=self.graph_xy_tolerance,
                    heading_stop_threshold=self.heading_stop_threshold,
                ):
                    return False
            prev = (wx, wy)

        self.publish_active_node(len(waypoints))
        return self.drive_to_pose(goal_handle, pose, label)

    def execute_navigate_to_pose(self, goal_handle):
        result = NavigateToPose.Result()
        if not self.wait_for_initial_tag_fix():
            goal_handle.abort()
            self.stop()
            return result
        ok = self.plan_and_drive(goal_handle, goal_handle.request.pose.pose, "navigate_to")
        if ok:
            goal_handle.succeed()
        elif not goal_handle.is_cancel_requested:
            goal_handle.abort()
        self.stop()
        self.clear_route()
        return result

    def execute_follow_waypoints(self, goal_handle):
        result = FollowWaypoints.Result()
        result.missed_waypoints = []
        if not self.wait_for_initial_tag_fix():
            result.missed_waypoints = list(range(len(goal_handle.request.poses)))
            goal_handle.abort()
            self.stop()
            return result
        poses: Iterable[PoseStamped] = goal_handle.request.poses
        for index, stamped in enumerate(poses):
            if goal_handle.is_cancel_requested:
                self.stop()
                goal_handle.canceled()
                return result
            ok = self.plan_and_drive(goal_handle, stamped.pose, f"follow_waypoints[{index}]")
            if not ok:
                result.missed_waypoints.append(index)
                break
        if result.missed_waypoints:
            goal_handle.abort()
        else:
            goal_handle.succeed()
        self.stop()
        self.clear_route()
        return result


def main() -> None:
    rclpy.init()
    node = TagRouteFollowerNode()
    executor = MultiThreadedExecutor(num_threads=2)
    executor.add_node(node)
    try:
        executor.spin()
    finally:
        node.stop()
        executor.shutdown()
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
