AI_RDKS100P_Robot/world2/world/world_node.py

#!/usr/bin/env python3
"""
world_node.py - ROS2 World Model 节点

功能：
    - 订阅配置的 Topic，将数据写入 WorldModel
    - 定时发布世界状态 Snapshot
    - YAML 配置驱动

使用：
    ros2 run world world_node --ros-args --params-file config/config.yaml
"""

import json
import yaml
import os
from typing import Optional

import rclpy
from rclpy.node import Node
from rclpy.callback_groups import MutuallyExclusiveCallbackGroup
from std_msgs.msg import String
from ament_index_python.packages import get_package_share_directory

from world.world_model import WorldModel


class TopicSubscription:
    """Topic 订阅配置封装"""
    
    def __init__(
        self,
        topic: str,
        update_type: str,
        entity_type: Optional[str] = None,
        ttl: Optional[float] = None,
        mapping: Optional[dict] = None,
    ):
        self.topic = topic
        self.update_type = update_type
        self.entity_type = entity_type
        self.ttl = ttl
        self.mapping = mapping or {}
    
    def extract_entity_id(self, data: dict) -> Optional[str]:
        """从数据中提取 entity_id"""
        field = self.mapping.get("entity_id")
        if not field:
            return None
        return self._get_nested(data, field)
    
    def extract_state(self, data: dict) -> dict:
        """提取 state 字段"""
        state_mapping = self.mapping.get("state", {})
        result = {}
        for target_key, source_key in state_mapping.items():
            value = self._get_nested(data, source_key)
            if value is not None:
                result[target_key] = value
        return result
    
    def extract_metadata(self, data: dict) -> Optional[dict]:
        """提取 metadata"""
        meta_mapping = self.mapping.get("metadata", {})
        if not meta_mapping:
            return None
        result = {}
        for target_key, source_key in meta_mapping.items():
            value = self._get_nested(data, source_key)
            if value is not None:
                result[target_key] = value
        return result if result else None
    
    @staticmethod
    def _get_nested(data: dict, path: str):
        """从嵌套字典中按路径获取值"""
        keys = path.split(".")
        current = data
        for key in keys:
            if isinstance(current, dict) and key in current:
                current = current[key]
            else:
                return None
        return current


class WorldNode(Node):
    """ROS2 World Model 节点"""
    
    def __init__(self) -> None:
        super().__init__("world_node")
        
        self.declare_parameter("publish_snapshot_topic", "/world/snapshot")
        self.declare_parameter("snapshot_rate", 1.0)
        self.declare_parameter("snapshot_mode", "minimal")
        
        self.publish_snapshot_topic = self.get_parameter("publish_snapshot_topic").value
        self.snapshot_rate = self.get_parameter("snapshot_rate").value
        self.snapshot_mode = self.get_parameter("snapshot_mode").value
        
        self._world_model = WorldModel()
        self.get_logger().info("WorldModel initialized")
        
        self._subscription_configs: list[TopicSubscription] = []
        self._load_subscriptions()
        self._create_subscriptions()
        
        self._snapshot_publisher = self.create_publisher(
            String,
            self.publish_snapshot_topic,
            qos_profile=10,
        )
        
        timer_period = 1.0 / self.snapshot_rate if self.snapshot_rate > 0 else 1.0
        self._timer = self.create_timer(timer_period, self._publish_snapshot)
        
        self.get_logger().info(
            f"WorldNode started, snapshot rate: {self.snapshot_rate} Hz"
        )
    
    def _load_subscriptions(self) -> None:
        """从 YAML 加载订阅配置"""
        config_path = os.environ.get(
            'WORLD_CONFIG_PATH',
            os.path.join(get_package_share_directory('world'), 'config', 'config.yaml'),
        )
        
        try:
            config_path = os.path.normpath(config_path)
            with open(config_path, 'r') as f:
                config_data = yaml.safe_load(f)
            
            params = config_data.get('/**', {}).get('ros__parameters', {})
            raw_subs = params.get('subscriptions', [])
            
            for item in raw_subs:
                if isinstance(item, str):
                    self._subscription_configs.append(item)
                elif isinstance(item, bytes):
                    self._subscription_configs.append(item.decode('utf-8'))
            
            self.get_logger().info(f"Loaded {len(self._subscription_configs)} subscription configs")
        except Exception as e:
            self.get_logger().warn(f"Failed to load config: {e}")
    
    def _create_subscriptions(self) -> None:
        """根据配置创建订阅"""
        callback_group = MutuallyExclusiveCallbackGroup()
        
        for config_str in self._subscription_configs:
            try:
                config = json.loads(config_str)
            except json.JSONDecodeError as e:
                self.get_logger().error(f"Failed to parse config: {e}")
                continue
            
            topic = config.get("topic")
            update_type = config.get("update_type", "entity")
            
            if not topic:
                continue
            
            sub_config = TopicSubscription(
                topic=topic,
                update_type=update_type,
                entity_type=config.get("entity_type"),
                ttl=config.get("ttl"),
                mapping=config.get("mapping", {}),
            )
            
            self.create_subscription(
                String,
                topic,
                lambda msg, sc=sub_config: self._topic_callback(msg, sc),
                qos_profile=10,
                callback_group=callback_group,
            )
            
            self.get_logger().info(f"Subscribed: {topic}")
        
        self.get_logger().info(f"Created {len(self._subscription_configs)} subscriptions")
    
    def _topic_callback(self, msg: String, sub_config: TopicSubscription) -> None:
        """处理接收到的消息"""
        try:
            data = json.loads(msg.data)
            self._process_message(data, sub_config)
        except json.JSONDecodeError as e:
            self.get_logger().warn(f"JSON parse failed [{sub_config.topic}]: {e}")
        except Exception as e:
            self.get_logger().error(f"Process message failed [{sub_config.topic}]: {e}")
    
    def _process_message(self, data: dict, sub_config: TopicSubscription) -> None:
        """根据 update_type 处理消息"""
        update_type = sub_config.update_type
        
        if update_type == "entity":
            self._update_entity(data, sub_config)
        elif update_type == "environment":
            self._update_environment(data, sub_config)
        elif update_type == "system":
            self._update_system(data, sub_config)
    
    def _update_entity(self, data: dict, sub_config: TopicSubscription) -> None:
        """更新 Entity"""
        entity_id = sub_config.extract_entity_id(data)
        if not entity_id:
            return
        
        entity_type = sub_config.entity_type or "unknown"
        state = sub_config.extract_state(data)
        metadata = sub_config.extract_metadata(data)
        
        self._world_model.update_entity(
            entity_id=entity_id,
            entity_type=entity_type,
            state=state,
            metadata=metadata,
            ttl=sub_config.ttl,
        )
    
    def _update_environment(self, data: dict, sub_config: TopicSubscription) -> None:
        """更新环境状态"""
        state = sub_config.extract_state(data)
        if state:
            self._world_model.update_environment_batch(state)
    
    def _update_system(self, data: dict, sub_config: TopicSubscription) -> None:
        """更新系统状态"""
        state = sub_config.extract_state(data)
        if state:
            self._world_model.update_system_batch(state)
    
    def _publish_snapshot(self) -> None:
        """定时发布 Snapshot"""
        try:
            snapshot = self._world_model.snapshot(mode=self.snapshot_mode)
            msg = String()
            msg.data = json.dumps(snapshot, ensure_ascii=False)
            self._snapshot_publisher.publish(msg)
        except Exception as e:
            self.get_logger().error(f"Publish snapshot failed: {e}")
    
    def destroy_node(self) -> None:
        self.get_logger().info("WorldNode shutting down")
        super().destroy_node()


def main(args=None):
    rclpy.init(args=args)
    try:
        node = WorldNode()
        rclpy.spin(node)
    finally:
        if rclpy.ok():
            rclpy.shutdown()


if __name__ == '__main__':
    main()