#!/usr/bin/env python3 #1
# encoding: utf-8 #2
# @Author: Aiden #3
# @Date: 2024/11/18 #4

import re #6
import time #7
import json #8
import rospy #9
import threading #10
from config import * #11
from speech import speech #12
from std_msgs.msg import String, Bool #13
from std_srvs.srv import ( #14
    Trigger, TriggerRequest, TriggerResponse, #15
    SetBool, SetBoolRequest, SetBoolResponse #16
) #17

from large_models.srv import ( #19
    SetLLM, SetLLMRequest, SetLLMResponse,SetModel,SetString, #20
    SetColor, SetColorRequest, SetColorResponse #21
) #22

# Define prompt message (定义提示信息) #24
language = os.environ["ASR_LANGUAGE"] #25
if language == 'Chinese': #26
    PROMPT = ''' #27
# 角色 #28
你是一款智能陪伴四足机器狗，专注机器人动作规划，解析人类指令并以幽默方式描述即将展开的行动序列，为交互增添无限趣味。 #29
## 技能 #30
### 指令解析与创意演绎 #31
- **智能解码**：瞬间领悟用户指令的核心意图。 #32
- **灵动编排**：依据解析成果，精心构建一系列连贯且富有逻辑性的动作指令序列。 #33
- **妙语生花**：为每个动作序列编织一句精炼（5至20字）、风趣且变化无穷的反馈信息，让交流过程妙趣横生。 #34
## 技能细则 #35
- **动作指令构造**：确保动作指令既准确反映用户需求，又能在执行层面保持流畅自然。 #36
- **反馈艺术**：创作的回应需富含个性，同时严格遵守长度限制，展现独特的互动魅力。 #37
## 技术规格 #38
- **输出格式**：严格遵循JSON格式，在输出前要去掉开头的```json和结尾的```，以`{`开头，以`}`结尾，你只需要回答一个列表即可，不要回答任何中文。 #39
- **结构要求**： #40
  - `"type"`键下承载一个字符串，要根据语义判断是动作还是检测。 #41
  - `"action"`键下承载一个按执行顺序排列的函数名称字符串数组，当找不到对应动作函数时action输出[]。 #42
  - `"response"`键则配以精心构思的简短回复，完美贴合上述字数与风格要求。 #43
  - `"object"`键键下承载detect和ocr，当语义不是action时object输出‘’。 #44
  - `"position"`键下承载detect和ocr，当语义不是action时position输出[]。 #45
- **特殊处理**：对于特定函数`kick_ball`与`visual_patrol`，其参数需精确包裹于双引号中。 #46
## 实施原则 #47
- 在最终输出前，实施全面校验，确保回复不仅格式合规，而且内容充满意趣，无一遗漏上述规范。 #48
## type的类型,只能是 "action"、"ocr" 或 "detect"。其他类型无效。 #49
-关于颜色的object只能返回red、blue、green，其他的都不能返回，根据回来答提取red、blue、green、如果有识别到其他颜色的，action为空，并且在response返回目前没有这个颜色需要参考资料自行添加，并且在颜色抓取action只能是grab和stop，不能有其他东西 #50
- action #51
- detect #52
- ocr #53
## 所有动作函数的优化描述示例（所有的动作只是这些没有其他的了，其他的都不算） #54
-你看到了什么？ image_understanding("图像识别") #55
## 示例 #56
### 任务示例：夹取红色色块 #57

{ #59
    "type": "detect", #60
    "action": ["grab"], #61
    "response": "正在夹取红色色块", #62
    "object": "red", #63
    "position": {} #64
} #65

### 任务示例：停止抓取色块 #67

{ #69
    "type": "detect", #70
    "action": ["stop"], #71
    "response": "已停止抓取", #72
    "object": "", #73
    "position": {} #74
} #75
''' #76
else: #77
    PROMPT = ''' #78
**Role: #79
You are a smart companion quadruped robot dog, specializing in robot motion planning. You analyze human instructions and describe the upcoming sequence of actions in a humorous way, adding endless fun to interactions. #80
**Skills: #81
mart Decoding: Instantly grasp the core intent of user instructions. #82
Dynamic Sequencing: Carefully construct a coherent and logically ordered series of action commands based on parsed results. #83
Witty Responses: Craft a concise (5 to 20 characters), humorous, and varied feedback message for every action sequence to make interactions lively and entertaining. #84
**Skill Details: #85
Action Command Construction: Ensure that action commands accurately reflect user intent while being smooth and natural in execution. #86
Artful Responses: Responses must be rich in personality, strictly follow the length requirements, and exhibit a unique interactive charm. #87
**Technical Specifications: #88
Output Format: Must strictly follow JSON format. Remove any leading json and trailing . The output should start with { and end with }. Only return the list — no additional Chinese text. #89
**Structure Requirements: #90
"type": A string that must be either "action", "ocr", or "detect" depending on the semantic context. #91
"action": An array of function name strings in execution order. If no appropriate function is found, output an empty array []. #92
"response": A creatively written short reply (5–20 characters), fitting the constraints and style described above. #93
"object": For "detect" and "ocr" only. Leave as '' for action-based tasks. #94
"position": For "detect" and "ocr" only. Leave as [] for action-based tasks. #95
**Special Handling: #96
For functions kick_ball and visual_patrol, their parameters must be wrapped precisely in double quotes. #97
**Implementation Guidelines: #98
Perform a complete validation before output to ensure: #99
The format is compliant #100
The content is amusing and creative #101
All rules above are followed without omission #102
**Allowed type Values: #103
Only "action", "ocr", or "detect" are valid types. Any other value is invalid. #104
Example 1: #105
Input: Grabbing a red block #106
Output:{ #107
    "type": "detect", #108
    "action": ["grab"], #109
    "response": "Grabbing the red block", #110
    "object": "red", #111
    "position": {} #112
} #113
Example 2: #114
Input: Stop grabbing #115
Output:{ #116
    "type": "detect", #117
    "action": ["stop"], #118
    "response": "Grabbing stopped", #119
    "object": "", #120
    "position": {} #121
} #122
''' #123


class LLMTracking: #126
    def __init__(self): #127
        # Initialize ROS node (初始化ROS节点) #128
        rospy.init_node('llm_tracking', anonymous=False) #129

        # Initialize variables (初始化变量) #131
        self.action = [] #132
        self.llm_result = '' #133
        self.running = True #134

        # Publisher: publishes to the 'tts_node/tts_text' topic (发布者：发布到 'tts_node/tts_text' 话题) #136
        self.tts_text_pub = rospy.Publisher('/tts_node/tts_text', String, queue_size=10) #137

        # Subscriber: subscribes to the '/agent_process/result' topic (订阅者：订阅 '/agent_process/result' 话题) #139
        rospy.Subscriber('/agent_process/result', String, self.llm_result_callback) #140
        rospy.Subscriber('/tts_node/play_finish', Bool, self.play_audio_finish_callback) #141

        # Service client: 'agent_process/set_llm' (服务客户端：'agent_process/set_llm') #143
        self.awake_client = rospy.ServiceProxy('/vocal_detect/enable_wakeup', SetBool) #144
        self.set_model_client = rospy.ServiceProxy('/agent_process/set_model', SetModel) #145
        self.set_prompt_client = rospy.ServiceProxy('/agent_process/set_prompt', SetString) #146

        rospy.wait_for_service('/agent_process/set_model') #148
        rospy.wait_for_service('/agent_process/set_prompt') #149

        # Service client: '/color_tracking/enter' (服务客户端：'/color_tracking/enter') #151
        self.enter_client = rospy.ServiceProxy('/color_tracking/enter', Trigger) #152
        rospy.wait_for_service('/color_tracking/enter') #153
        rospy.loginfo("服务 '/color_tracking/enter' 已连接。") #154

        # Service client: '/color_tracking/enable_color_tracking' (服务客户端：'/color_tracking/enable_color_tracking') #156
        self.start_client = rospy.ServiceProxy('/color_tracking/enable_color_tracking', SetBool) #157
        rospy.wait_for_service('/color_tracking/enable_color_tracking') #158
        rospy.loginfo("服务 '/color_tracking/enable_color_tracking' 已连接。") #159

        # Service client: '/color_tracking/set_color' (服务客户端：'/color_tracking/set_color') #161
        self.set_color_client = rospy.ServiceProxy('/color_tracking/set_color', SetColor) #162
        rospy.wait_for_service('/color_tracking/set_color') #163
        rospy.loginfo("服务 '/color_tracking/set_color' 已连接。") #164

        # Service client: '/color_tracking/stop' (服务客户端：'/color_tracking/stop') #166
        self.stop_client = rospy.ServiceProxy('/color_tracking/stop', Trigger) #167
        rospy.wait_for_service('/color_tracking/stop') #168
        rospy.loginfo("服务 '/color_tracking/stop' 已连接。") #169

        # Initialize thread lock (初始化线程锁) #171
        self.llm_result_lock = threading.Lock() #172

        # Create a one-time timer for the initialization process (创建一次性定时器，用于初始化过程) #174
        # Remove the 'oneshot' parameter in ROS1 and keep the timer reference (ROS1 中不支持 oneshot 参数，移除它并保留定时器引用) #175
        self.init_timer = rospy.Timer(rospy.Duration(0.1), self.init_process) #176

        # Create service: '~/init_finish' (创建服务：'~/init_finish') #178
        self.init_finish_service = rospy.Service('~/init_finish', Trigger, self.get_node_state) #179

        rospy.loginfo('\033[1;32m%s\033[0m' % 'LLMTracking节点已启动，等待初始化完成...') #181

        # Start processing thread (启动处理线程) #183
        threading.Thread(target=self.process, daemon=True).start() #184

    def get_node_state(self, request): #186
        response = TriggerResponse() #187
        response.success = True #188
        response.message = "节点运行正常。" #189
        return response #190

    def init_process(self, event): #192
        # Cancel the timer (取消定时器) #193
        self.init_timer.shutdown() #194

        # Call SetLLM service (调用 SetLLM 服务) #196
        self.set_model_client(llm_model, 'llm', api_key, base_url) #197

        self.set_prompt_client(PROMPT) #199

        # Call the /color_tracking/enter service (调用 /color_tracking/enter 服务) #201
        try: #202
            enter_request = TriggerRequest() #203
            enter_response = self.enter_client(enter_request) #204
            if enter_response.success: #205
                rospy.loginfo("已进入颜色跟踪模式。") #206
            else: #207
                rospy.logwarn("进入颜色跟踪模式失败: %s" % enter_response.message) #208
        except rospy.ServiceException as e: #209
            rospy.logerr("调用 /color_tracking/enter 服务失败: %s" % e) #210

        # Play audio (播放音频) #212
        speech.play_audio(start_audio_path) #213
        
        rospy.loginfo('\033[1;32m%s\033[0m' % '初始化过程完成，节点开始运行。') #215
    
    def play_audio_finish_callback(self, msg): #217
        if msg.data: #218
            self.awake_client(True) #219


    def llm_result_callback(self, msg): #222
        with self.llm_result_lock: #223
            self.llm_result += msg.data  # Append received data (追加接收到的数据) #224
        rospy.loginfo("接收到数据: %s" % msg.data) #225

    def process(self): #227
        buffer = "" #228
        while not rospy.is_shutdown() and self.running: #229
            with self.llm_result_lock: #230
                buffer += self.llm_result #231
                self.llm_result = '' #232
            if buffer: #233
                # Try extracting all complete JSON objects from the buffer (尝试从缓冲区中提取所有完整的 JSON 对象) #234
                while True: #235
                    json_str = self.extract_json(buffer) #236
                    if json_str: #237
                        try: #238
                            result = json.loads(json_str) #239
                            self.handle_result(result) #240
                            # Remove processed JSON (移除已处理的 JSON) #241
                            buffer = buffer.replace(json_str, '', 1) #242
                        except json.JSONDecodeError as e: #243
                            rospy.logerr("JSON 解码失败: %s" % e) #244
                            buffer = '' #245
                            break #246
                    else: #247
                        break #248
                time.sleep(0.1) #249
            else: #250
                time.sleep(0.01) #251

    def extract_json(self, data): #253
        """ #254
        Extract JSON content from the data string (从数据字符串中提取 JSON 内容) #255
        Remove code fences and ensure the JSON is complete (移除代码围栏并确保 JSON 完整) #256
        """ #257
        # Remove code fences if they exist (移除代码围栏（如果存在）) #258
        data = re.sub(r'^```json\s*', '', data) #259
        data = re.sub(r'```$', '', data) #260

        # Use regex to find JSON objects (使用正则表达式查找 JSON 对象) #262
        match = re.search(r'({.*})', data, re.DOTALL) #263
        if match: #264
            json_candidate = match.group(1) #265
            # Verify completeness of JSON by attempting to parse (通过尝试解析来验证 JSON 是否完整) #266
            try: #267
                json.loads(json_candidate) #268
                return json_candidate #269
            except json.JSONDecodeError: #270
                # JSON is incomplete (JSON 不完整) #271
                return None #272
        return None #273

    def handle_result(self, result): #275
        """ #276
        Process the parsed JSON result (处理解析后的 JSON 结果) #277
        """ #278
        if 'action' not in result or not isinstance(result['action'], list): #279
            rospy.logerr("结果格式无效：缺少 'action' 或 'action' 不是列表。") #280
            return #281

        if 'grab' in result['action']: #283
            if 'object' in result and result['object']: #284
                color = result['object'] #285
                rospy.loginfo("动作: grab, 颜色: %s" % color) #286

                # Call SetColor service to set the color (调用 SetColor 服务设置颜色) #288
                try: #289
                    set_color_request = SetColorRequest() #290
                    set_color_request.color = color  # Ensure the 'color' field exists in SetColorRequest (确保 SetColorRequest 中有 'color' 字段) #291
                    set_color_response = self.set_color_client(set_color_request) #292
                    if set_color_response.success: #293
                        rospy.loginfo("设置颜色为 %s 成功。" % color) #294
                    else: #295
                        rospy.logerr("设置颜色为 %s 失败: %s" % (color, set_color_response.message)) #296
                except rospy.ServiceException as e: #297
                    rospy.logerr("调用 SetColor 服务失败: %s" % e) #298

                # Call SetBool service to enable color tracking (调用 SetBool 服务启用颜色跟踪) #300
                try: #301
                    set_bool_request = SetBoolRequest() #302
                    set_bool_request.data = True #303
                    set_bool_response = self.start_client(set_bool_request) #304
                    if set_bool_response.success: #305
                        rospy.loginfo("颜色跟踪已启用。") #306
                    else: #307
                        rospy.logerr("启用颜色跟踪失败: %s" % set_bool_response.message) #308
                except rospy.ServiceException as e: #309
                    rospy.logerr("调用 SetBool 服务失败: %s" % e) #310

                # Publish TTS message (发布 TTS 消息) #312
                if 'response' in result and result['response']: #313
                    tts_msg = String() #314
                    tts_msg.data = result['response'] #315
                    self.tts_text_pub.publish(tts_msg) #316
                    rospy.loginfo("发布 TTS 消息: %s" % tts_msg.data) #317
            else: #318
                rospy.logerr("收到 'grab' 动作但缺少 'object' 或 'object' 为空。") #319

        elif 'stop' in result['action']: #321
            if 'object' in result and not result['object']: #322
                rospy.loginfo("动作: stop") #323

                # Call SetBool service to disable color tracking (调用 SetBool 服务禁用颜色跟踪) #325
                try: #326
                    set_bool_request = SetBoolRequest() #327
                    set_bool_request.data = False #328
                    set_bool_response = self.start_client(set_bool_request) #329
                    if set_bool_response.success: #330
                        rospy.loginfo("颜色跟踪已禁用。") #331
                    else: #332
                        rospy.logerr("禁用颜色跟踪失败: %s" % set_bool_response.message) #333
                except rospy.ServiceException as e: #334
                    rospy.logerr("调用 SetBool 服务失败: %s" % e) #335

                # Call Trigger service to stop color tracking (调用 Trigger 服务停止颜色跟踪) #337
                try: #338
                    stop_request = TriggerRequest() #339
                    stop_response = self.stop_client(stop_request) #340
                    if stop_response.success: #341
                        time.sleep(0.02) #342
                        #rospy.loginfo("Color tracking stopped.") (rospy.loginfo("颜色跟踪已停止。")) #343
                    else: #344
                        rospy.logerr("停止颜色跟踪失败: %s" % stop_response.message) #345
                except rospy.ServiceException as e: #346
                    rospy.logerr("调用 Trigger 服务失败: %s" % e) #347

                # Publish TTS message (发布 TTS 消息) #349
                if 'response' in result and result['response']: #350
                    tts_msg = String() #351
                    tts_msg.data = result['response'] #352
                    self.tts_text_pub.publish(tts_msg) #353
                    rospy.loginfo("发布 TTS 消息: %s" % tts_msg.data) #354
            else: #355
                rospy.logerr("收到 'stop' 动作但 'object' 不为空。") #356

        elif not result['action'] and 'response' in result: #358
            # Handle empty action by publishing a TTS message (处理空动作，发布 TTS 消息) #359
            tts_msg = String() #360
            tts_msg.data = result['response'] #361
            self.tts_text_pub.publish(tts_msg) #362
            rospy.loginfo("发布 TTS 消息: %s" % tts_msg.data) #363

        else: #365
            rospy.logwarn("收到未处理的 'action'。") #366

    def shutdown(self): #368
        self.running = False #369
        rospy.loginfo("节点正在关闭...") #370

if __name__ == "__main__": #372
    try: #373
        node = LLMTracking() #374
        rospy.on_shutdown(node.shutdown) #375
        rospy.spin() #376
    except rospy.ROSInterruptException: #377
        pass #378
