#!/usr/bin/env python3
import json
import math
import time
import clip
import torch
from copy import deepcopy
import os
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import PoseStamped
from tf2_ros import Buffer, TransformListener
from rclpy.duration import Duration
from rclpy.qos import QoSProfile, ReliabilityPolicy, DurabilityPolicy, HistoryPolicy
from nav_msgs.msg import OccupancyGrid
from nav2_simple_commander.robot_navigator import BasicNavigator, TaskResult
from ament_index_python.packages import get_package_share_directory


class SemanticNavNode(Node):
    def __init__(self, json_path: str):
        super().__init__('semantic_nav_node')
        self.declare_parameter('approach_offset_m', 0.7)
        self.declare_parameter('approach_offset_refrigerator_m', 0.7)
        self.declare_parameter('approach_offset_water_dispenser_m', 0.7)
        self.declare_parameter('approach_offset_sofa_m', 0.5)
        self.declare_parameter('approach_offset_white_toilet_m', 0.7)
        self.declare_parameter('approach_offset_office_chair_with_wheels_m', 0.45)
        self.declare_parameter('approach_candidate_count', 8)
        self.declare_parameter('approach_clearance_m', 0.30)
        self.declare_parameter('approach_clearance_office_chair_with_wheels_m', 0.18)
        self.declare_parameter('approach_candidate_count_office_chair_with_wheels', 16)
        self.declare_parameter('approach_offset_office_chair_with_wheels_search_radii', [0.45, 0.60, 0.75])
        self.map_msg = None

        # Load object library
        with open(json_path) as f:
            self.object_lib = json.load(f)

        # Load CLIP
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model, _ = clip.load("ViT-B/32", device=self.device)

        # Build text embeddings for all object names
        self.object_names = list(self.object_lib.keys())
        text_tokens = clip.tokenize(self.object_names).to(self.device)
        with torch.no_grad():
            text_features = self.model.encode_text(text_tokens)
            self.text_features = text_features / text_features.norm(dim=-1, keepdim=True)

        # Initialize navigator
        self.navigator = BasicNavigator()
        self.navigator.waitUntilNav2Active()
        self.tf_buffer = Buffer(cache_time=Duration(seconds=10))
        self.tf_listener = TransformListener(self.tf_buffer, self, spin_thread=True)
        map_qos = QoSProfile(
            depth=1,
            reliability=ReliabilityPolicy.RELIABLE,
            durability=DurabilityPolicy.TRANSIENT_LOCAL,
            history=HistoryPolicy.KEEP_LAST,
        )
        self.map_sub = self.create_subscription(
            OccupancyGrid,
            "/map",
            self.map_callback,
            map_qos,
        )
        self.get_logger().info("Navigator ready")

    def query_object(self, prompt: str):
        tokens = clip.tokenize([prompt]).to(self.device)
        with torch.no_grad():
            query_feat = self.model.encode_text(tokens)
            query_feat /= query_feat.norm(dim=-1, keepdim=True)

        sims = (query_feat @ self.text_features.T).squeeze(0)
        best_idx = sims.argmax().item()
        best_name = self.object_names[best_idx]

        # Align with vlm_detection: position may be None if not yet detected
        obj_info = self.object_lib.get(best_name, {})
        pos = obj_info.get("position")

        if pos is None:
            self.get_logger().warn(
                f"'{best_name}' matched but has no position in JSON. "
                f"Run vlm_detection in AMCL mode first!"
            )
            return best_name, None

        return best_name, pos

    def navigate_to_object(self, prompt: str):
        name, pos = self.query_object(prompt)
        if pos is None:
            self.get_logger().warn(f"Cannot navigate: no valid position for '{name}'. Please run vlm_detection first!")
            return

        robot_pose = self.get_robot_pose_in_map()
        if robot_pose is None:
            self.get_logger().warn("Cannot navigate: failed to get robot pose in map frame.")
            return

        offset_m = self.get_approach_offset(name)
        goal_x, goal_y, yaw = self.compute_approach_goal(name, robot_pose, pos, offset_m)
        robot_x, robot_y = robot_pose
        distance_to_goal = math.hypot(goal_x - robot_x, goal_y - robot_y)

        goal_pose = PoseStamped()
        goal_pose.header.frame_id = "map"
        goal_pose.header.stamp = self.navigator.get_clock().now().to_msg()
        goal_pose.pose.position.x = goal_x
        goal_pose.pose.position.y = goal_y
        goal_pose.pose.position.z = 0.0
        goal_pose.pose.orientation.z = math.sin(yaw / 2.0)
        goal_pose.pose.orientation.w = math.cos(yaw / 2.0)

        self.navigator.goToPose(deepcopy(goal_pose))
        self.get_logger().info(
            f"Navigating to {name}: object={pos}, goal={{'x': {goal_x:.4f}, 'y': {goal_y:.4f}}}, "
            f"approach_offset={offset_m:.2f}m"
        )
        self.get_logger().info(
            f"Robot pose={{'x': {robot_x:.4f}, 'y': {robot_y:.4f}}}, "
            f"distance_to_goal={distance_to_goal:.4f}m"
        )

        # Wait until navigation completes
        while not self.navigator.isTaskComplete():
            feedback = self.navigator.getFeedback()
            if feedback:
                self.get_logger().info(
                    "Navigation task in progress"
                )

        result = self.navigator.getResult()
        final_robot_pose = self.get_robot_pose_in_map()
        if final_robot_pose is not None:
            final_x, final_y = final_robot_pose
            final_distance_to_goal = math.hypot(goal_x - final_x, goal_y - final_y)
            self.get_logger().info(
                f"Final robot map pose={{'x': {final_x:.4f}, 'y': {final_y:.4f}}}, "
                f"distance_to_goal={final_distance_to_goal:.4f}m"
            )
        if result == TaskResult.SUCCEEDED:
            self.get_logger().info("Navigation succeeded")
        elif result == TaskResult.CANCELED:
            self.get_logger().info("Navigation canceled")
        elif result == TaskResult.FAILED:
            self.get_logger().info("Navigation failed")

    def map_callback(self, msg: OccupancyGrid):
        self.map_msg = msg

    def get_robot_pose_in_map(self):
        target_frame = "map"
        source_frame = "base_footprint"
        last_error = None
        for _ in range(15):
            try:
                if not self.tf_buffer.can_transform(
                    target_frame,
                    source_frame,
                    rclpy.time.Time(),
                    timeout=Duration(seconds=0.2)
                ):
                    continue

                transform = self.tf_buffer.lookup_transform(
                    target_frame,
                    source_frame,
                    rclpy.time.Time(),
                    timeout=Duration(seconds=0.2)
                )
                return (
                    transform.transform.translation.x,
                    transform.transform.translation.y,
                )
            except Exception as e:
                last_error = e

        if last_error is None:
            self.get_logger().warn(
                "Robot pose lookup failed: transform map -> base_footprint was not available within timeout."
            )
        else:
            self.get_logger().warn(f"Robot pose lookup failed: {last_error}")
        return None

    def compute_approach_goal(self, object_name, robot_pose, obj_pos, offset_m):
        robot_x, robot_y = robot_pose
        obj_x = obj_pos["x"]
        obj_y = obj_pos["y"]
        search_radii = self.get_approach_search_radii(object_name, offset_m)
        candidate_count = self.get_approach_candidate_count(object_name)
        approach_candidates = []
        for radius in search_radii:
            approach_candidates.extend(
                self.generate_approach_candidates(obj_x, obj_y, radius, candidate_count)
            )
        feasible_candidates = [
            candidate for candidate in approach_candidates
            if self.is_candidate_feasible(object_name, candidate[0], candidate[1])
        ]

        path_candidates = []
        start_pose = self.build_pose_stamped(robot_x, robot_y, 0.0)
        for candidate in feasible_candidates:
            candidate_x, candidate_y, candidate_yaw = candidate
            candidate_goal_pose = self.build_pose_stamped(candidate_x, candidate_y, candidate_yaw)
            path = self.navigator.getPath(start_pose, candidate_goal_pose, use_start=True)
            if path and len(path.poses) >= 2:
                path_length = self.compute_path_length(path)
                if path_length >= 0.2:
                    path_candidates.append((candidate_x, candidate_y, candidate_yaw, path_length))

        if path_candidates:
            goal_x, goal_y, yaw, path_length = min(
                path_candidates,
                key=lambda candidate: candidate[3]
            )
            self.get_logger().info(
                f"Approach candidates for {object_name}: total={len(approach_candidates)} "
                f"feasible={len(feasible_candidates)} path_feasible={len(path_candidates)} "
                f"selected=({goal_x:.4f}, {goal_y:.4f}) path_length={path_length:.4f}m"
            )
            return goal_x, goal_y, yaw

        if feasible_candidates:
            goal_x, goal_y, yaw = min(
                feasible_candidates,
                key=lambda candidate: math.hypot(candidate[0] - robot_x, candidate[1] - robot_y)
            )
            self.get_logger().warn(
                f"Approach candidates for {object_name}: total={len(approach_candidates)} "
                f"feasible={len(feasible_candidates)} path_feasible=0, "
                f"falling back to nearest map-feasible candidate=({goal_x:.4f}, {goal_y:.4f})"
            )
            return goal_x, goal_y, yaw

        dx = obj_x - robot_x
        dy = obj_y - robot_y
        distance = math.hypot(dx, dy)
        if distance < 1e-6:
            yaw = 0.0
            return obj_x, obj_y, yaw

        ux = dx / distance
        uy = dy / distance
        goal_x = obj_x - ux * min(offset_m, max(distance - 0.05, 0.0))
        goal_y = obj_y - uy * min(offset_m, max(distance - 0.05, 0.0))
        yaw = math.atan2(dy, dx)
        self.get_logger().warn(
            f"Approach candidates for {object_name}: no feasible candidate found, "
            f"falling back to line-of-sight offset goal=({goal_x:.4f}, {goal_y:.4f})"
        )
        return goal_x, goal_y, yaw

    def generate_approach_candidates(self, obj_x, obj_y, offset_m, candidate_count):
        candidates = []
        for i in range(candidate_count):
            angle = (2.0 * math.pi * i) / candidate_count
            goal_x = obj_x + offset_m * math.cos(angle)
            goal_y = obj_y + offset_m * math.sin(angle)
            yaw = math.atan2(obj_y - goal_y, obj_x - goal_x)
            candidates.append((goal_x, goal_y, yaw))
        return candidates

    def is_candidate_feasible(self, object_name: str, world_x, world_y):
        map_msg = self.wait_for_map()
        if map_msg is None:
            return True

        info = map_msg.info
        resolution = info.resolution
        origin_x = info.origin.position.x
        origin_y = info.origin.position.y
        width = info.width
        height = info.height
        grid_x = int((world_x - origin_x) / resolution)
        grid_y = int((world_y - origin_y) / resolution)

        if grid_x < 0 or grid_x >= width or grid_y < 0 or grid_y >= height:
            return False

        clearance_m = self.get_approach_clearance(object_name)
        clearance_cells = max(0, int(math.ceil(clearance_m / resolution)))

        for dy in range(-clearance_cells, clearance_cells + 1):
            for dx in range(-clearance_cells, clearance_cells + 1):
                if dx * dx + dy * dy > clearance_cells * clearance_cells:
                    continue
                xx = grid_x + dx
                yy = grid_y + dy
                if xx < 0 or xx >= width or yy < 0 or yy >= height:
                    return False
                cost = map_msg.data[yy * width + xx]
                if cost < 0 or cost >= 50:
                    return False
        return True

    def build_pose_stamped(self, x, y, yaw):
        pose = PoseStamped()
        pose.header.frame_id = "map"
        pose.header.stamp = self.navigator.get_clock().now().to_msg()
        pose.pose.position.x = x
        pose.pose.position.y = y
        pose.pose.position.z = 0.0
        pose.pose.orientation.z = math.sin(yaw / 2.0)
        pose.pose.orientation.w = math.cos(yaw / 2.0)
        return pose

    def compute_path_length(self, path):
        if not path.poses or len(path.poses) < 2:
            return 0.0
        length = 0.0
        for prev_pose, next_pose in zip(path.poses[:-1], path.poses[1:]):
            dx = next_pose.pose.position.x - prev_pose.pose.position.x
            dy = next_pose.pose.position.y - prev_pose.pose.position.y
            length += math.hypot(dx, dy)
        return length

    def wait_for_map(self, timeout_sec: float = 2.0):
        if self.map_msg is not None:
            return self.map_msg
        end_time = time.time() + timeout_sec
        while self.map_msg is None and time.time() < end_time:
            time.sleep(0.05)
        return self.map_msg

    def get_approach_candidate_count(self, object_name: str) -> int:
        default_count = int(
            self.get_parameter('approach_candidate_count').get_parameter_value().integer_value
        )
        param_name = (
            "approach_candidate_count_"
            + object_name.lower().replace(" ", "_")
        )
        if self.has_parameter(param_name):
            default_count = int(self.get_parameter(param_name).get_parameter_value().integer_value)
        return min(max(default_count, 4), 16)

    def get_approach_search_radii(self, object_name: str, default_offset: float):
        param_name = (
            "approach_offset_"
            + object_name.lower().replace(" ", "_")
            + "_search_radii"
        )
        if self.has_parameter(param_name):
            value = self.get_parameter(param_name).value
            if value:
                return [float(radius) for radius in value]
        return [default_offset]

    def get_approach_clearance(self, object_name: str) -> float:
        default_clearance = self.get_parameter(
            'approach_clearance_m'
        ).get_parameter_value().double_value
        param_name = (
            "approach_clearance_"
            + object_name.lower().replace(" ", "_")
            + "_m"
        )
        if self.has_parameter(param_name):
            return self.get_parameter(param_name).get_parameter_value().double_value
        return default_clearance

    def get_approach_offset(self, object_name: str) -> float:
        default_offset = self.get_parameter(
            'approach_offset_m'
        ).get_parameter_value().double_value
        param_name = (
            "approach_offset_"
            + object_name.lower().replace(" ", "_")
            + "_m"
        )
        if self.has_parameter(param_name):
            return self.get_parameter(param_name).get_parameter_value().double_value
        return default_offset

    def cleanup(self):
        if hasattr(self, 'tf_listener'):
            try:
                if hasattr(self.tf_listener, 'executor'):
                    self.tf_listener.executor.shutdown(timeout_sec=0.5)
                if hasattr(self.tf_listener, 'dedicated_listener_thread'):
                    self.tf_listener.dedicated_listener_thread.join(timeout=0.5)
                self.tf_listener.unregister()
            except Exception:
                pass




def main():
    rclpy.init()
    package_share_dir = get_package_share_directory('vlm_nav_pkg')
    json_file_path = os.path.join(package_share_dir, 'config', 'example_object_detection_vlm.json')
    node = SemanticNavNode(json_file_path)

    try:
        user_input = input("Where do you want to go: ")
        node.navigate_to_object(user_input)
    finally:
        node.cleanup()
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()


if __name__ == "__main__":
    main()
