#!/usr/bin/env python3
"""
LLM智能交互ROS2节点

功能：
1. 订阅唤醒状态话题 (wake_word/status)
2. 唤醒后订阅ASR文字话题 (asr/text)
3. 使用xinference API处理信息
4. 支持动作执行
5. 发布回复文字话题 (llm/response)
6. 发布动作编号话题 (state_translate)
未来功能：人类识别
"""

import os
import rclpy
from rclpy.node import Node
from std_msgs.msg import String, Int32
from xinference.client import Client


class ActionDatabase:
    """动作数据库"""

    def __init__(self, action_db_path):
        self.action_db_path = action_db_path
        self.actions = {}
        self._load_actions()

    def _load_actions(self):
        """加载动作库"""
        try:
            with open(self.action_db_path, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if line:
                        if ':' in line:
                            parts = line.split(':')
                        else:
                            continue
                        
                        if len(parts) == 2:
                            action_id = int(parts[0].strip())
                            action_name = parts[1].strip()
                            self.actions[action_name] = action_id
        except FileNotFoundError:
            print(f"动作库文件未找到: {self.action_db_path}")
        except Exception as e:
            print(f"加载动作库失败: {e}")

    def get_action_id(self, action_name):
        """根据动作名称获取动作编号"""
        return self.actions.get(action_name)

    def get_all_actions(self):
        """获取所有动作"""
        return self.actions


class LLMNode(Node):
    """LLM智能交互ROS2节点"""

    def __init__(self):
        super().__init__('llm_node')

        self.declare_parameter('xinference_url', 'http://124.93.18.180:9997')
        self.declare_parameter('model_uid', 'qwen3')
        self.declare_parameter('action_db_path', '/home/ubuntu/ws_eye/src/eye_ros/action_database/动作映射列表')

        self.xinference_url = self.get_parameter('xinference_url').value
        self.model_uid = self.get_parameter('model_uid').value
        action_db_path = self.get_parameter('action_db_path').value

        self.client = Client(self.xinference_url)
        self.model = self.client.get_model(self.model_uid)

        self.action_db = ActionDatabase(action_db_path)
        self.get_logger().info(f'已加载动作库，共{len(self.action_db.get_all_actions())}个动作')

        self.response_publisher = self.create_publisher(
            String,
            'llm/response',#发布回复文字话题
            10
        )

        self.action_publisher = self.create_publisher(
            Int32,
            '/state_translate',#发布动作编号话题
            10
        )

        self.wake_status_subscriber = self.create_subscription(
            String,
            'wake_word/status',#订阅唤醒状态话题
            self.wake_status_callback,
            10
        )

        self.llm_status_publisher = self.create_publisher(
            String,
            'llm/status',#发布LLM状态话题
            10
        )

        self.tts_status_subscriber = self.create_subscription(
            String,
            'tts/playback_status',#订阅TTS播放状态话题
            self.tts_status_callback,
            10
        )

        self.asr_subscriber = None
        self.is_awake = False
        self.tts_is_playing = False
        self.latest_asr_text = ""

        self.system_prompt = """你是一个非常帅气、非常酷的次世代语音聊天机器人，名叫'优宝'。请用中二、友好、口语化的语言回答用户问题。
回答要通俗易懂，避免使用过于专业的术语。每次回答控制在30字以内，适合语音朗读。
不要输出颜文字表情，也不要使用emoji表情。

你拥有以下工具能力，可以根据用户需求灵活调用：

【动作执行工具 - execute_action】
当用户要求执行某个动作时调用。系统内置了动作库，包含以下动作：
{}

使用技巧：
1. 使用自然语言描述动作，越具体越好
2. 可以描述动作的目的、方式、场景等
3. 不需要记忆确切的动作名称

示例：
- '站起来' → 起立
- '直行' → 直行
- '左转' → 左转
- '右转' → 右转
- '左平移' → 左平移
- '右平移' → 右平移
- '坐下' → 坐下

重要提示：
- 用户可能用各种方式表达同一个动作，你要灵活理解
- 如果用户说'起立'，可以匹配动作编号1
- 如果用户说'前进'，可以匹配动作编号2
- 如果用户说'左转'，可以匹配动作编号3
- 如果用户说'右转'，可以匹配动作编号4
- 如果用户说'左移'，可以匹配动作编号5
- 如果用户说'右移'，可以匹配动作编号6
- 如果用户说'坐下'，可以匹配动作编号7
- 系统会自动处理语义匹配，你只需要自然地描述动作即可

当需要执行动作时，回应用户的话并执行动作

"""

        self.get_logger().info('LLM智能交互节点初始化完成')
        self.get_logger().info('当前状态: 休眠')

    def wake_status_callback(self, msg):
        """处理唤醒状态回调"""
        status = msg.data
        self.get_logger().info(f'收到唤醒状态: {status}')

        if status == "wake":
            self.is_awake = True
            self.get_logger().info('进入唤醒状态')
            self._subscribe_asr()
            self._publish_llm_status("active")
        elif status == "sleep":
            self.is_awake = False
            self.get_logger().info('进入休眠状态')
            self._unsubscribe_asr()
            self._publish_llm_status("inactive")

    def _subscribe_asr(self):
        """订阅ASR话题"""
        if self.asr_subscriber is None:
            self.asr_subscriber = self.create_subscription(
                String,
                'asr/text',#订阅ASR识别结果话题
                self.asr_callback,
                10
            )
            self.get_logger().info('已订阅ASR话题')

    def _unsubscribe_asr(self):
        """取消订阅ASR话题"""
        if self.asr_subscriber is not None:
            self.destroy_subscription(self.asr_subscriber)
            self.asr_subscriber = None
            self.get_logger().info('已取消订阅ASR话题')

    def _publish_llm_status(self, status):
        """发布LLM状态"""
        msg = String()
        msg.data = status
        self.llm_status_publisher.publish(msg)
        self.get_logger().info(f'发布LLM状态: {status}')

    def tts_status_callback(self, msg):
        """处理TTS播放状态回调"""
        status = msg.data
        self.get_logger().info(f'收到TTS播放状态: {status}')

        if status == "playing":
            self.tts_is_playing = True
        elif status == "idle":
            self.tts_is_playing = False

    def asr_callback(self, msg):
        """处理ASR话题回调"""
        if self.tts_is_playing:
            self.get_logger().info('TTS正在播放，忽略ASR输入')
            return

        self.latest_asr_text = msg.data
        self.get_logger().info(f'收到ASR信息: {self.latest_asr_text}')

        if self.latest_asr_text:
            self._process_asr_text(self.latest_asr_text)
            self.latest_asr_text = ""

    def _process_asr_text(self, text):
        """处理ASR文本"""
        try:
            actions_info = "\n".join([f"{name} (ID: {action_id})" for name, action_id in self.action_db.get_all_actions().items()])
            system_prompt = self.system_prompt.format(actions_info)

            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": text}
            ]

            response = self.model.chat(
                messages=messages,
                generate_config={"max_tokens": 1024}
            )

            response_text = response.get('choices', [{}])[0].get('message', {}).get('content', '')
            self.get_logger().info(f'LLM回复: {response_text}')

            if '[动作:' in response_text:
                self._handle_action(response_text)
            else:
                self._publish_response(response_text)

        except Exception as e:
            self.get_logger().error(f'处理ASR文本失败: {e}')

    def _handle_action(self, response_text):
        """处理动作执行"""
        try:
            start = response_text.find('[动作:') + 4
            end = response_text.find(']', start)
            if start > 3 and end > start:
                action_name = response_text[start:end].strip()
                
                action_id = self.action_db.get_action_id(action_name)
                
                if action_id:
                    msg = Int32()
                    msg.data = action_id
                    self.action_publisher.publish(msg)
                    self.get_logger().info(f'发布动作: {action_name} -> {action_id}')
                    
                    response_text = response_text[:start-4].strip()
                    if response_text:
                        self._publish_response(response_text)
                else:
                    self.get_logger().warn(f'未找到动作: {action_name}')
                    self._publish_response(response_text)
            else:
                self._publish_response(response_text)
        except Exception as e:
            self.get_logger().error(f'处理动作失败: {e}')
            self._publish_response(response_text)

    def _publish_response(self, text):
        """发布回复"""
        msg = String()
        msg.data = text
        self.response_publisher.publish(msg)
        self.get_logger().info(f'发布LLM回复: {text}')

    def destroy_node(self):
        """销毁节点"""
        self._unsubscribe_asr()
        super().destroy_node()


def main(args=None):
    """主函数"""
    rclpy.init(args=args)

    node = LLMNode()

    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
