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

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


sys.path.append('/home/ubuntu/software/puppypi_control/') #22
from action_group_control import runActionGroup, stopActionGroup #23

class VLLMCamera(object): #25
    def __init__(self, name): #26
        rospy.init_node(name) #27
        self.image_queue = queue.Queue(maxsize=2) #28
        self.language = os.environ["ASR_LANGUAGE"] #29
        self.bridge = CvBridge() #30
        self.vllm_result = '' #31
        self.running = True #32
        self.puppy_control_node = PuppyControlNode() #33

        # Publisher （发布者） #35
        self.tts_text_pub = rospy.Publisher('/tts_node/tts_text', String, queue_size=1) #36
        # Subscriber （订阅者） #37
        rospy.Subscriber('/usb_cam/image_raw', Image, self.image_callback) #38
        rospy.Subscriber('/agent_process/result', String, self.vllm_result_callback) #39
        rospy.Subscriber('/tts_node/play_finish', Bool, self.play_audio_finish_callback) #40
        # Service（服务） #41
        self.cli = rospy.ServiceProxy('/puppy_control/go_home', Empty) #42
        self.awake_client = rospy.ServiceProxy('/vocal_detect/enable_wakeup', SetBool) #43
        self.set_model_client = rospy.ServiceProxy('/agent_process/set_model', SetModel)         #44
        self.run_action_group_srv = rospy.ServiceProxy('/puppy_control/runActionGroup', SetRunActionName) #45
        
        rospy.wait_for_service('/puppy_control/runActionGroup') #47
        
        # Initialization process （初始化过程） #49
        self.init_process() #50

    def get_node_state(self, req): #52
        resp = Trigger.Response() #53
        resp.success = True #54
        return resp #55

    def init_process(self): #57
        """Initialization process （初始化过程）""" #58
        if os.environ['ASR_LANGUAGE'] == 'Chinese': #59
            self.set_model_client(stepfun_vllm_model, 'vllm', stepfun_api_key, stepfun_base_url) #60
        else: #61
            self.set_model_client(vllm_model, 'vllm', vllm_api_key, vllm_base_url) #62

        # Play start audio （播放启动音频） #64
        speech.play_audio(start_audio_path) #65

        # Start image processing thread （启动图像处理线程） #67
        threading.Thread(target=self.process, daemon=True).start() #68
        rospy.loginfo('启动成功') #69
        rospy.Service('~/init_finish', Trigger, self.get_node_state) #70

    def play_audio_finish_callback(self, msg): #72
        if msg.data: #73
            self.awake_client(True) #74

    def vllm_result_callback(self, msg): #76
        """Callback to handle results from VLLM （处理来自 VLLM 的结果回调）""" #77
        self.vllm_result = msg.data   #78
        rospy.loginfo("接收到 VLLM 结果回调: %s" % self.vllm_result) #79
        self.handle_vllm_result()   #80

    def handle_vllm_result(self): #82
        """Directly play VLLM result （直接播放VLLM结果）""" #83
        if not self.vllm_result: #84
            rospy.logwarn("VLLM 结果为空，跳过处理") #85
            return #86
        try: #87
            # Directly play the received text （直接播放接收到的文本） #88
            response_msg = String() #89
            response_msg.data = self.vllm_result #90
            self.tts_text_pub.publish(response_msg) #91
            rospy.loginfo("发布 TTS 消息: %s" % self.vllm_result) #92

        except Exception as e: #94
            rospy.logerr("处理 VLLM 结果出错: %s" % e) #95
        finally: #96
            self.vllm_result = ''  # Clear results to avoid repeated processing （清空结果，避免重复处理） #97

    def process(self): #99
        """Image processing thread （图像处理线程）""" #100
        rospy.loginfo("图像处理线程已启动。") #101
        while not rospy.is_shutdown() and self.running: #102
            try: #103
                image = self.image_queue.get(block=True, timeout=1) #104
                
                # Display image （显示图像） #106
                cv2.imshow('image', image) #107
                if cv2.waitKey(1) & 0xFF == ord('q'): #108
                    self.running = False #109
                    rospy.loginfo("检测到 'q' 按键，关闭图像窗口。") #110
                    break #111
            except queue.Empty: #112
                continue #113
            except Exception as e: #114
                rospy.logerr("处理图像时出错: %s" % e) #115
        cv2.destroyAllWindows() #116
        rospy.loginfo("图像处理线程已结束。") #117

    def image_callback(self, ros_image): #119
        """Image callback function: convert ROS Image to OpenCV image and enqueue it （图像回调函数，将 ROS Image 转换为 OpenCV 图像并放入队列）""" #120
        try: #121
            # Use cv_bridge to convert ROS Image to OpenCV format （使用 cv_bridge 将 ROS Image 转换为 OpenCV 图像） #122
            rgb_image = self.bridge.imgmsg_to_cv2(ros_image, desired_encoding='bgr8') #123
            if self.image_queue.full(): #124
                # If the queue is full, discard the oldest image （如果队列已满，丢弃最旧的图像） #125
                discarded = self.image_queue.get() #126
                #rospy.logwarn("图像队列已满，丢弃最旧的图像") #127
            # Add the image to the queue （将图像放入队列） #128
            self.image_queue.put(rgb_image) #129
        except Exception as e: #130
            rospy.logerr("处理图像时出错: %s" % e) #131

def main(): #133
    try: #134
        node = VLLMCamera('vllm_camera') #135
        rospy.spin() #136
    except KeyboardInterrupt: #137
        rospy.loginfo('关闭节点') #138
    finally: #139
        if 'node' in locals() and node.running: #140
            node.running = False #141
        cv2.destroyAllWindows() #142
        rospy.loginfo("程序结束。") #143

if __name__ == "__main__": #145
    main() #146
