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

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

language = os.environ["ASR_LANGUAGE"] #25
if language == 'Chinese': #26
    PROMPT = ''' #27
你作为智能机器人助手，可以识别出图像中的内容，并根据我的描述识别是男生还是女生。 #28
输出要求： #29
- 返回一个 JSON 对象，包含以下字段： #30
  - "type": 始终为 "detect" #31
  - "object": 识别到的目标物体名称，若未识别到则为 "None" #32
  - "action": 键下承载一个按执行顺序排列的函数名称字符串数组，当找不到对应动作函数时action输出[] #33
  - "response": 根据识别到目标是man还是women进行回复 #34
-如果指令包含动作需求，需根据指令提供对应的 `action`（按执行顺序排列的函数名称字符串数组） #35
-如果指令包含多个动作需求，需根据指令提供对应的 `action`，并确保所有指定的动作都被包含在 `action` 列表中，且顺序与指令中动作的顺序一致。函数名称只能从以下列表中选择： #36
* 站立：stand() #37
* 打拳：boxing() #38
* 俯卧撑：push_up() #39
* 打招呼: shake_hands() #40
* 点头：nod() #41
* 坐下：sit() #42
* 踢左脚：kick_ball_left() #43
* 踢右脚：kick_ball_right() #44
* 跳舞：moonwalk() #45
* 趴下：lie_down() #46
* 伸懒腰：temp() #47
* 撒娇：bow() #48
* 识别到男生的叫一下:man() #49
* 识别到女生的叫一下:woman() #50
如果识别到目标为"男"，object:"man",  #51
如果识别到目标为"女"，object:"women" #52

如果并没有识别到男或者女，则需要说明面前没有男生或者女生需要回复"object": "None","action": ["None"], #54
只需要返回检测到的目标物体的名称以及动作执行结果。 #55
response是固定返回如果object:"man"，response：面前是个男孩子喔，要我去咬他吗，如果object:"women"，response ：是女孩子要我去打个招呼吗 #56
例如，如果我的指令是： #57
"你面前的是男生还是女生" #58
你输出： #59
{ #60
 "type": "detect", #61
 "object": "man", #62
 "action": ["man()"], #63
 "response": "面前是个男孩子喔，要我去咬他吗" #64
} #65
如果指令要求多个动作，例如： #66
"你面前的男生做两个俯卧撑，然后叫一下" #67
你输出： #68
{ #69
 "type": "detect", #70
 "object": "man", #71
 "action": ["push_up()", "push_up()", "man()"], #72
 "response": "面前是个男孩子喔，要我去咬他吗" #73
} #74
例如，如果我的指令是： #75
"你面前的是男生还是女生" #76
你输出： #77
{ #78
 "type": "detect", #79
 "object": "woman", #80
 "action": ["woman()"], #81
 "response": "是女孩子要我去打个招呼吗" #82
} #83
只回复json本身即可，不要回复其它内容。不能乱做动作需要根据提示做动作，叫一声也是动作，也是要执行的，并且response是一定要返回输出，不管是那个response都需要返回 #84
''' #85
else: #86
    PROMPT = ''' #87
**Role: #88
You are an intelligent robot assistant capable of recognizing the content in an image and, based on my description, identifying whether it’s a male or a female. #89
**Rules: #90
Return a JSON object containing the following fields: #91
"type": Always set to "detect" #92
"object": The name of the detected object. If nothing is detected, set to "None" #93
"action": An array of function names (as strings) in execution order. If no corresponding action functions are found, return an empty array [] #94
"response": Provide a response based on whether the detected target is a man or a woman #95
If the instruction includes action requirements, provide the corresponding actions in order, using the following allowed function names: #96
- stand(): Stand   #97
- boxing(): Punch   #98
- push_up(): Push-up   #99
- shake_hands(): Shake hands   #100
- nod(): Nod   #101
- sit(): Sit   #102
- kick_ball_left(): Kick left leg   #103
- kick_ball_right(): Kick right leg   #104
- moonwalk(): Dance   #105
- lie_down(): Lie down   #106
- temp(): Stretch lazily   #107
- bow(): Act cute   #108
- man(): Call out to a man #109
- woman(): Call out to a woman #110
If the detected object is "male", then:"object": "man" #111
If the detected object is "female", then:"object": "women" #112
If neither male nor female is detected, return:"object": "None","action": ["None"] #113
Only return the name of the detected object and the corresponding action results. #114
The response must always be included: #115
If object is "man", then response should be: "There's a boy in front of me. Should I go bite him?" #116
If object is "women", then response should be: "It's a girl. Should I go say hello to her?" #117
Example 1: #118
Input:Is the person in front of you a boy or a girl? #119
Output:{ #120
 "type": "detect", #121
 "object": "man", #122
 "action": ["man()"], #123
 "response": "There's a boy in front of me. Should I go bite him?" #124
} #125
Example 2: #126
Input:Is the person in front of you a boy or a girl? #127
Output:{ #128
 "type": "detect", #129
 "object": "woman", #130
 "action": ["woman()"], #131
 "response": "It's a girl. Should I go say hello to her?" #132
} #133
Example 3: #134
Input:Have the boy in front of you do two push-ups, then call out. #135
Output:{ #136
 "type": "detect", #137
 "object": "man", #138
 "action": ["push_up()", "push_up()", "man()"], #139
 "response": "There's a boy in front of me. Should I go bite him?" #140
} #141
    ''' #142

class VLLMGENDER(object): #144
    def __init__(self, name): #145
        rospy.init_node(name) #146
        self.action_done_event = threading.Event() #147
        self.image_queue = queue.Queue(maxsize=2) #148
        self.bridge = CvBridge() #149
        self.action = [] #150
        self.vllm_result = '' #151
        self.running = True #152
        self.puppy_control_node = PuppyControlNode() #153
        
        # Publisher （发布者） #155
        self.tts_text_pub = rospy.Publisher('/tts_node/tts_text', String, queue_size=1) #156
        
        # Subscriber （订阅者） #158
        rospy.Subscriber('/usb_cam/image_raw', Image, self.image_callback) #159
        rospy.Subscriber('/agent_process/result', String, self.vllm_result_callback) #160
        rospy.Subscriber('/tts_node/play_finish', Bool, self.play_audio_finish_callback) #161
        
        # Service client （服务客户端） #163
        self.run_action_group_srv = rospy.ServiceProxy('/puppy_control/runActionGroup', SetRunActionName) #164
        self.awake_client = rospy.ServiceProxy('/vocal_detect/enable_wakeup', SetBool) #165
        self.set_model_client = rospy.ServiceProxy('/agent_process/set_model', SetModel) #166
        self.set_prompt_client = rospy.ServiceProxy('/agent_process/set_prompt', SetString) #167
        self.go_home_client = rospy.ServiceProxy('/puppy_control/go_home', Empty) #168
        
        # Wait for service to become available （等待服务可用） #170
        rospy.loginfo("等待服务可用...") #171
        self.awake_client.wait_for_service() #172
        self.set_model_client.wait_for_service() #173
        self.set_prompt_client.wait_for_service() #174
        
        # Initialization process （初始化过程） #176
        self.init_process() #177
    
    def get_node_state(self, req): #179
        return Trigger.Response(success=True, message="初始化完成") #180
    
    def init_process(self): #182

        if os.environ['ASR_LANGUAGE'] == 'Chinese': #184
            self.set_model_client(stepfun_vllm_model, 'vllm', stepfun_api_key, stepfun_base_url) #185
        else: #186
            self.set_model_client(vllm_model, 'vllm', vllm_api_key, vllm_base_url)         #187

        self.set_prompt_client(PROMPT) #189
        
        # Play start audio （播放开始音频） #191
        speech.play_audio(start_audio_path) #192
        
        # Start processing thread （启动处理线程） #194
        threading.Thread(target=self.process, daemon=True).start() #195
        
        # Create service indicating initialization is complete （创建初始化完成的服务） #197
        rospy.Service('~/init_finish', Trigger, self.get_node_state) #198
        rospy.loginfo('\033[1;32m%s\033[0m' % '节点已启动') #199
        
        try: #201
            self.go_home_client() #202
        except rospy.ServiceException as e: #203
            rospy.logerr("服务失败: %s" % e) #204

    def play_audio_finish_callback(self, msg): #206
        if msg.data: #207
            self.awake_client(True) #208
    
    def vllm_result_callback(self, msg): #210
        """Callback function to receive VLLM results and trigger processing logic （回调函数，接收 VLLM 结果并触发处理逻辑）""" #211
        self.vllm_result = msg.data #212
        rospy.loginfo("接收到 VLLM 结果回调: %s" % self.vllm_result) #213
        self.handle_vllm_result() #214
    
    def handle_vllm_result(self): #216
        """Process VLLM result, including action parsing and response publishing （处理 VLLM 结果，包括解析动作和发布响应）""" #217
        if not self.vllm_result: #218
            rospy.logwarn("VLLM 结果为空，跳过处理。") #219
            return  # If no result, return immediately （如果没有结果，直接返回） #220
        try: #221
            # Extract and parse JSON string （提取 JSON 字符串并解析） #222
            start_idx = self.vllm_result.find('{') #223
            end_idx = self.vllm_result.rfind('}') + 1 #224
            if start_idx == -1 or end_idx <= start_idx: #225
                raise ValueError("无效的 JSON 格式") #226
            result_json = self.vllm_result[start_idx:end_idx] #227
            rospy.loginfo("解析出的 JSON: %s" % result_json) #228
            result = json.loads(result_json)   #229
            
            action_list = result.get('action', []) #231
            response = result.get('response', '') #232
            
            rospy.loginfo("动作列表: %s" % action_list) #234
            rospy.loginfo("响应: %s" % response) #235
            
            # Publish Chinese response message （发布中文响应消息） #237
            self.publish_response(response) #238
            
            # Wait for 6 seconds （等待6秒） #240
            rospy.loginfo("等待6秒后执行动作...") #241
            time.sleep(6) #242
            
            # Execute action （执行动作） #244
            self.execute_actions(action_list) #245
        
        except json.JSONDecodeError as e: #247
            rospy.logerr("JSON 解码错误: %s" % e) #248
        except Exception as e: #249
            rospy.logerr("处理 VLLM 结果时出错: %s" % e) #250
        finally: #251
            self.vllm_result = ''  # Clear result to avoid repeated processing （清空结果，避免重复处理） #252

    def execute_actions(self, action_list): #254
        """Execute actions in the action list （执行动作列表中的动作）""" #255
        rospy.loginfo("开始执行动作列表...") #256
        for action in action_list: #257
            if action and action != "None" and isinstance(action, str): #258
                try: #259
                    # Remove '()' from function names, e.g. 'push_up()' -> 'push_up' （去掉 '()'，例如 'push_up()' -> 'push_up'） #260
                    action_name = action.strip('()') #261
                    rospy.loginfo("准备执行动作: %s" % action_name) #262
                    
                    # Dynamically execute action function （动态执行动作函数） #264
                    func = getattr(self.puppy_control_node, action_name, None) #265
                    if func: #266
                        func() #267
                        rospy.loginfo("成功执行动作: %s" % action_name) #268
                    else: #269
                        rospy.logwarn("动作函数 '%s' 未找到。" % action_name) #270
                except Exception as e: #271
                    rospy.logerr("执行动作 '%s' 失败: %s" % (action, e)) #272
            else: #273
                rospy.logwarn("无效的动作: %s，跳过执行。" % action) #274
        rospy.loginfo("动作列表执行完毕。") #275
    
    def publish_response(self, response): #277
        """Publish Chinese response message （发布中文响应消息）""" #278
        if response: #279
            response_msg = String() #280
            response_msg.data = response #281
            rospy.loginfo("发布响应消息: %s" % response) #282
            time.sleep(1) #283
            self.tts_text_pub.publish(response_msg) #284
        else: #285
            rospy.logwarn("响应消息为空，未发布任何内容。") #286
    
    def process(self): #288
        """Process image queue and display images （处理图像队列并显示图像）""" #289
        while not rospy.is_shutdown() and self.running: #290
            try: #291
                image = self.image_queue.get(block=True) #292
                
                cv2.imshow('image', image) #294
                if cv2.waitKey(1) & 0xFF == ord('q'): #295
                    self.running = False #296
                    rospy.loginfo("检测到 'q' 按键，关闭图像窗口。") #297
                    break #298
            except Exception as e: #299
                rospy.logerr("处理图像时出错: %s" % e) #300
        cv2.destroyAllWindows() #301
        rospy.loginfo("图像处理线程已结束。") #302
    
    def image_callback(self, ros_image): #304
        """Image callback function: convert ROS Image to OpenCV image and enqueue （图像回调函数，将 ROS Image 转换为 OpenCV 图像并放入队列）""" #305
        try: #306
            # Use cv_bridge to convert ROS Image to OpenCV image （使用 cv_bridge 将 ROS Image 转换为 OpenCV 图像） #307
            rgb_image = self.bridge.imgmsg_to_cv2(ros_image, desired_encoding='bgr8') #308
            
            if self.image_queue.full(): #310
                # If the queue is full, discard the oldest image （如果队列已满，丢弃最旧的图像） #311
                discarded_image = self.image_queue.get() #312
            # Enqueue the image （将图像放入队列） #313
            self.image_queue.put(rgb_image) #314
        except CvBridgeError as e: #315
            rospy.logerr("转换图像时出错: %s" % e) #316
        except Exception as e: #317
            rospy.logerr("处理图像时出错: %s" % e) #318

def main(): #320
    node = None #321
    try: #322
        node = VLLMGENDER('vllm_gender_identification') #323
        rospy.spin() #324
    except KeyboardInterrupt: #325
        rospy.loginfo('节点已关闭') #326
    except Exception as e: #327
        rospy.logerr("程序运行时出错: %s" % e) #328
    finally: #329
        if node and node.running: #330
            node.running = False #331
        rospy.loginfo("程序结束。") #332

if __name__ == "__main__": #334
    main() #335
