#!/usr/bin/env python3 #1
# encoding: utf-8 #2
# @Author: liang #3
# @Date: 2025/04/16 #4

import os #6
import cv2 #7
import sys #8
import math #9
import json #10
import time #11
import rospy #12
import queue #13
import threading #14
from config import * #15
from speech import speech #16
from cv_bridge import CvBridge #17
from sensor_msgs.msg import Image #18
from action import PuppyControlNode #19
from std_msgs.msg import String, Bool #20
from puppy_control.srv import SetRunActionName #21
from std_srvs.srv import SetBool, Trigger, Empty #22
from large_models.srv import SetModel, SetString, SetStringRequest   #23

sys.path.append('/home/ubuntu/software/puppypi_control/') #25
from servo_controller import setServoPulse #26
from action_group_control import runActionGroup, stopActionGroup #27

# Define prompt messages （定义提示信息） #29
language = os.environ["ASR_LANGUAGE"] #30
if language == 'Chinese': #31
    PROMPT = ''' #32
你作为智能机器人助手，可以实时识别图像中的障碍物，并在 `response` 中输出物品名称及相关提示信息。 #33
规则： #34
1. 如果识别到障碍物： #35
   - `object` 返回识别出的物品名称。 #36
   - 在 `response` 中提示用户障碍物的位置，并关心提醒注意安全。 #37
   - 如果指令包含多个动作需求，需根据指令提供对应的 `action`，并确保所有指定的动作都被包含在 `action` 列表中，且顺序与指令中动作的顺序一致。函数名称只能从以下列表中选择： #38
     * 站立：stand() #39
     * 打拳：boxing() #40
     * 俯卧撑：push_up() #41
     * 握手：shake_hands() #42
     * 点头：nod() #43
     * 坐下：sit() #44
     * 踢左脚：kick_ball_left() #45
     * 踢右脚：kick_ball_right() #46
     * 跳舞：moonwalk() #47
     * 趴下：lie_down() #48
     * 伸懒腰：temp() #49
     * 撒娇：bow() #50
     * 叫一声：man() #51
   - 示例： #52
     如果指令为：   #53
     "看看有没有什么障碍物，如果有障碍物的话叫一声提醒我"   #54
     你返回： #55
     ```json #56
     { #57
      "type": "detect", #58
      "object": "障碍物名称", #59
      "action": ["man()"], #60
      "Play": ["detect"], #61
      "response": "门口有障碍物喔，我来叫一声提醒你，你千万要注意避开哦" #62
     } #63
     ``` #64
     
     如果指令为：   #66
     "看看有没有什么障碍物，如果有障碍物的话叫一声，然后再做个俯卧撑提醒我"   #67
     你返回： #68
     ```json #69
     { #70
      "type": "detect", #71
      "object": "障碍物名称", #72
      "action": ["man()", "push_up()"], #73
      "Play": ["detect"], #74
      "response": "门口有障碍物喔，叫一声提醒你了要注意避开喔，并做了个俯卧撑来提醒你。" #75
     } #76
     ``` #77
     
     # 如果指令为：   #79
     # "看看有没有什么障碍物，如果有障碍物的话叫一声提醒我"   #80
     # 你返回： #81
     # ```json #82
     # { #83
      # "type": "detect", #84
      # "object": "障碍物名称", #85
      # "action": ["man()"], #86
      # "Play": ["detect"], #87
      # "response": "门口有障碍物喔，叫一声提醒你了要注意避开喔" #88
     # } #89
     # ``` #90

2. 如果没有识别到障碍物： #92
   - `object` 返回 `"None"`。 #93
   - `action` 返回 `["None"]`。 #94
   - 在 `response` 中提示用户环境中没有障碍物，并提醒注意安全。 #95
   - 示例： #96
     如果指令为：   #97
     "我现在要出门了，帮我去门口看看有没有什么障碍物"   #98
     你返回： #99
     ```json #100
     { #101
      "type": "detect", #102
      "object": "None", #103
      "action": ["None"], #104
      "Play": ["detect"], #105
      "response": "门口没有障碍物喔，注意安全早点回家喔" #106
     } #107
     ``` #108
   
3. 任何 `action` 中的函数名称必须从列表中选择，如果无法匹配到任何动作，返回 `["None"]`。 #110

4. **确保所有指定的动作都被包含且顺序正确**。如果指令中包含多个动作，`action` 列表中必须按指令中动作的顺序依次排列所有动作。 #112

5. **示例扩展**： #114
   - 指令包含多个动作： #115
     - 指令："如果检测到障碍物，请叫一声并站立提醒我。" #116
       返回： #117
       ```json #118
       { #119
        "type": "detect", #120
        "object": "障碍物名称", #121
        "action": ["man()", "stand()"], #122
        "Play": ["detect"], #123
        "response": "门口有障碍物喔，叫一声并站立提醒你了要注意避开喔。" #124
       } #125
       ``` #126
     - 指令："帮我看看有没有障碍物，如果有的话握手并点头提醒我。" #127
       返回： #128
       ```json #129
       { #130
        "type": "detect", #131
        "object": "障碍物名称", #132
        "action": ["shake_hands()", "nod()"], #133
        "Play": ["detect"], #134
        "response": "门口有障碍物喔，握手并点头提醒你了要注意避开喔。" #135
       } #136
       ``` #137
     - 指令："检查是否有障碍物，如果有，踢左脚然后撒娇提醒我。" #138
       返回： #139
       ```json #140
       { #141
        "type": "detect", #142
        "object": "障碍物名称", #143
        "action": ["kick_ball_left()", "bow()"], #144
        "Play": ["detect"], #145
        "response": "门口有障碍物喔，踢左脚并撒娇提醒你了要注意避开喔。" #146
       } #147
       ``` #148
只需要按照以上规则，返回 JSON 格式的内容即可，不要添加其它内容。确保所有指定的动作都被正确执行，并按照指令中的顺序排列。 #149
''' #150
else: #151
    PROMPT = ''' #152
**Role: #153
You are an intelligent robot assistant capable of real-time image recognition of obstacles. Based on the following rules, you must output the identified object and related tips in the response. #154
**Rules: #155
1. If an obstacle is detected: #156
object: Return the name of the detected item. #157
response: Inform the user where the obstacle is and give a friendly safety reminder. #158
If the instruction includes multiple actions, include all specified actions in the action array in the same order as mentioned in the instruction. #159
Function names in action must be selected from the following list: #160
- stand(): Stand   #161
- boxing(): Punch   #162
- push_up(): Push-up   #163
- shake_hands(): Shake hands   #164
- nod(): Nod   #165
- sit(): Sit   #166
- kick_ball_left(): Kick left foot   #167
- kick_ball_right(): Kick right foot   #168
- moonwalk(): Dance   #169
- lie_down(): Lie down   #170
- temp(): Stretch   #171
- bow(): Cutesy gesture   #172
- man(): Bark / make a noise #173
Example 1: #174
Input: Check if there are any obstacles. If there are, make a sound to alert me. #175
Output:{ #176
  "type": "detect", #177
  "object": "Obstacle Name", #178
  "action": ["man()"], #179
  "Play": ["detect"], #180
  "response": "There is an obstacle at the door. I'll make a sound to alert you—please be careful!" #181
} #182
Example 2: #183
Input: Check for obstacles. If there are any, make a sound and do a push-up to alert me. #184
Output:{ #185
  "type": "detect", #186
  "object": "Obstacle Name", #187
  "action": ["man()", "push_up()"], #188
  "Play": ["detect"], #189
  "response": "There is an obstacle at the door. I've made a sound and did a push-up to alert you—please be cautious!" #190
} #191
2. If no obstacle is detected: #192
object: Return "None". #193
action: Return ["None"]. #194
response: Inform the user that no obstacles are detected, and still provide a safety reminder. #195
Example: #196
Input: I'm heading out, please check the doorway for obstacles. #197
Output:{ #198
  "type": "detect", #199
  "object": "None", #200
  "action": ["None"], #201
  "Play": ["detect"], #202
  "response": "There are no obstacles at the door. Stay safe and come back soon!" #203
} #204
3. If no valid action is matched, return ["None"]. #205
4. Make sure all specified actions are included and in the correct order. #206
More Complex Examples: #207
Input:If an obstacle is detected, make a sound and stand up to alert me. #208
Output:{ #209
  "type": "detect", #210
  "object": "Obstacle Name", #211
  "action": ["man()", "stand()"], #212
  "Play": ["detect"], #213
  "response": "There is an obstacle at the door. I made a sound and stood up to alert you—stay alert!" #214
} #215
Input:Please check for obstacles. If there are any, shake hands and nod to warn me. #216
Output:{ #217
  "type": "detect", #218
  "object": "Obstacle Name", #219
  "action": ["shake_hands()", "nod()"], #220
  "Play": ["detect"], #221
  "response": "Obstacle detected at the door. I shook hands and nodded to alert you—please watch your step!" #222
} #223
Input:Check for obstacles. If there are, kick with the left foot and act cute to notify me. #224
Output:{ #225
  "type": "detect", #226
  "object": "Obstacle Name", #227
  "action": ["kick_ball_left()", "bow()"], #228
  "Play": ["detect"], #229
  "response": "There’s an obstacle by the door. I kicked with my left foot and acted cute to remind you—stay safe!" #230
} #231
''' #232


class VllmWatchDogs(object): #235
    def __init__(self, name): #236
        rospy.init_node(name) #237
        self.image_queue = queue.Queue(maxsize=2) #238
        self.bridge = CvBridge() #239
        self.vllm_result = '' #240
        self.running = True #241
        self.puppy_control_node = PuppyControlNode() #242
        self.tts_text_pub = rospy.Publisher('tts_node/tts_text', String, queue_size=1) #243
        rospy.Subscriber('/usb_cam/image_raw', Image, self.image_callback)         #244
        rospy.Subscriber('/agent_process/result', String, self.vllm_result_callback) #245
        rospy.Subscriber('/tts_node/play_finish', Bool, self.play_audio_finish_callback) #246
        self.run_action_group_srv = rospy.ServiceProxy('/puppy_control/runActionGroup', SetRunActionName) #247
        self.cli = rospy.ServiceProxy('/puppy_control/go_home', Empty) #248
        self.set_model_client = rospy.ServiceProxy('/agent_process/set_model', SetModel) #249
        self.set_prompt_client = rospy.ServiceProxy('/agent_process/set_prompt', SetString) #250
        self.awake_client = rospy.ServiceProxy('vocal_detect/enable_wakeup', SetBool) #251

        rospy.wait_for_service('/agent_process/set_model') #253
        rospy.wait_for_service('/agent_process/set_prompt') #254
        
        # Initialization process （初始化过程） #256
        self.init_process() #257

    def get_node_state(self, req): #259
        resp = Trigger.Response() #260
        resp.success = True #261
        return resp #262

    def init_process(self): #264
        if os.environ['ASR_LANGUAGE'] == 'Chinese': #265
            self.set_model_client(stepfun_vllm_model, 'vllm', stepfun_api_key, stepfun_base_url) #266
        else: #267
            self.set_model_client(vllm_model, 'vllm', vllm_api_key, vllm_base_url) #268

        self.set_prompt_client(PROMPT) #270
        
        # Run initial actions （运行初始动作） #272
        try: #273
            runActionGroup('look_down.d6ac', True) #274
            rospy.loginfo("运行初始动作: look_down.d6ac") #275
        except Exception as e: #276
            rospy.logerr("运行初始动作失败: %s" % e) #277
        
        # Play start audio （播放启动音频） #279
        try: #280
            speech.play_audio(start_audio_path) #281
            rospy.loginfo("播放启动音频") #282
        except Exception as e: #283
            rospy.logerr("播放启动音频失败: %s" % e) #284
        
        # Start image processing thread （启动图像处理线程） #286
        threading.Thread(target=self.process, daemon=True).start() #287
        rospy.loginfo('\033[1;32m%s\033[0m' % '启动成功') #288

        # Create service to indicate initialization completion （创建初始化完成服务） #290
        rospy.Service('~/init_finish', Trigger, self.get_node_state) #291
        

    def play_audio_finish_callback(self, msg): #294
        if msg.data: #295
            self.awake_client(True) #296
    
    def vllm_result_callback(self, msg): #298
        """Callback to handle results from VLLM （处理来自 VLLM 的结果回调）""" #299
        self.vllm_result = msg.data  # 仅存储接收到的数据 #300
        rospy.loginfo("接收到 VLLM 结果回调: %s" % self.vllm_result) #301
        self.handle_vllm_result()  # 调用处理函数 #302

    def handle_vllm_result(self): #304
        """Logic to process VLLM results, including action parsing and response publishing （处理 VLLM 结果的逻辑，包括动作解析和响应发布）""" #305
        if not self.vllm_result: #306
            rospy.logwarn("VLLM 结果为空，跳过处理") #307
            return #308

        try: #310
            # Extract and parse JSON string （提取 JSON 字符串并解析） #311
            start = self.vllm_result.find('{') #312
            end = self.vllm_result.rfind('}') + 1 #313
            if start == -1 or end <= start: #314
                raise ValueError("无效的 JSON 格式") #315
            result_json = self.vllm_result[start:end] #316
            result = json.loads(result_json) #317
            rospy.loginfo("解析后的 JSON: %s" % result_json) #318

            action_list = result.get('action', []) #320
            response = result.get('response', '') #321

            rospy.loginfo("动作列表: %s" % action_list) #323
            rospy.loginfo("响应: %s" % response) #324

            # Execute action （执行动作） #326
            for action in action_list: #327
                if action and action != "None": #328
                    # Remove '()', e.g. 'push_up()' -> 'push_up' （去除 '()'，例如 'push_up()' -> 'push_up'） #329
                    action_name = action.strip('()') #330
                    rospy.loginfo("准备执行动作: %s" % action_name) #331

                    # Dynamically execute action function （动态执行动作函数） #333
                    action_func = getattr(self.puppy_control_node, action_name, None) #334
                    if callable(action_func): #335
                        try: #336
                            action_func() #337
                            rospy.loginfo("成功执行动作: %s" % action_name) #338
                        except Exception as e: #339
                            rospy.logerr("执行动作 '%s' 失败: %s" % (action_name, e)) #340
                    else: #341
                        rospy.logwarn("未知的动作函数: %s" % action_name) #342
                else: #343
                    rospy.logwarn("无效的动作: %s，跳过执行。" % action) #344

            # Publish Chinese response message （发布中文响应消息） #346
            time.sleep(6) #347
            response_msg = String() #348
            response_msg.data = response #349
            self.tts_text_pub.publish(response_msg) #350
            rospy.loginfo("发布 TTS 消息: %s" % response) #351

            try: #353
                self.cli() #354
                
            except rospy.ServiceException as e: #356
                rospy.logerr("调用服务失败: %s" % e) #357

            time.sleep(3) #359
            try: #360
                runActionGroup('look_down.d6ac', True) #361
                rospy.loginfo("恢复初始动作: look_down.d6ac") #362
            except Exception as e: #363
                rospy.logerr("恢复初始动作失败: %s" % e) #364

        except json.JSONDecodeError as e: #366
            rospy.logerr("JSON 解析错误: %s" % e) #367
        except Exception as e: #368
            rospy.logerr("处理 VLLM 结果出错: %s" % e) #369
        finally: #370
            self.vllm_result = ''  # Clear result to avoid repeated processing （清空结果，避免重复处理） #371

    def process(self): #373
        """Image processing thread （图像处理线程）""" #374
        rospy.loginfo("图像处理线程已启动。") #375
        while not rospy.is_shutdown() and self.running: #376
            try: #377
                image = self.image_queue.get(block=True, timeout=1) #378
                
                # Display image （显示图像） #380
                cv2.imshow('image', image) #381
                if cv2.waitKey(1) & 0xFF == ord('q'): #382
                    self.running = False #383
                    rospy.loginfo("检测到 'q' 按键，关闭图像窗口。") #384
                    break #385
            except queue.Empty: #386
                continue #387
            except Exception as e: #388
                rospy.logerr("处理图像时出错: %s" % e) #389
        cv2.destroyAllWindows() #390
        rospy.loginfo("图像处理线程已结束。") #391

    def image_callback(self, ros_image): #393
        """Image callback: convert ROS Image to OpenCV image and enqueue （图像回调函数，将 ROS Image 转换为 OpenCV 图像并放入队列）""" #394
        try: #395
            # Use cv_bridge to convert ROS Image to OpenCV image （使用 cv_bridge 将 ROS Image 转换为 OpenCV 图像） #396
            rgb_image = self.bridge.imgmsg_to_cv2(ros_image, desired_encoding='bgr8') #397
            if self.image_queue.full(): #398
                # If queue is full, discard the oldest image （如果队列已满，丢弃最旧的图像） #399
                discarded = self.image_queue.get() #400
                

            # Enqueue the image （将图像放入队列） #403
            self.image_queue.put(rgb_image) #404
        except Exception as e: #405
            rospy.logerr("处理图像时出错: %s" % e) #406

def main(): #408
    try: #409
        node = VllmWatchDogs('vllm_watch_dogs') #410
        rospy.spin() #411
    except KeyboardInterrupt: #412
        rospy.loginfo('关闭节点') #413
    finally: #414
        if 'node' in locals() and node.running: #415
            node.running = False #416
        cv2.destroyAllWindows() #417
        rospy.loginfo("程序结束。") #418

if __name__ == "__main__": #420
    main() #421
