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

import rospy #6
import queue #7
import numpy as np #8
from config import * #9
from speech import speech #10
from cv_bridge import CvBridge #11
from std_msgs.msg import String #12
from sensor_msgs.msg import Image #13
from std_srvs.srv import Trigger, TriggerResponse, SetBool, Empty #14
from large_models.srv import SetString, SetStringResponse,SetModel, SetModelResponse,  SetLLM, SetLLMResponse, SetContent #15

class AgentProcess(object): #17
    def __init__(self, name): #18
        # Initialize the ROS node (初始化ROS节点) #19
        rospy.init_node(name, anonymous=False) #20

        # Get parameters, use default values if not set (获取参数，如果没有设置则使用默认值) #22
        self.camera_topic = rospy.get_param('~camera_topic', '/usb_cam/image_raw') #23

        self.prompt = '' #25
        self.model = llm_model  #26
        self.chat_text = '' #27
        self.start_record_chat = False #28
        self.bridge = CvBridge() #29
        self.language = os.environ["ASR_LANGUAGE"] #30
        self.image_queue = queue.Queue(maxsize=2) #31

        self.client = speech.OpenAIAPI(api_key, base_url) #33
        # Create a publisher (创建发布者) #34
        self.result_pub = rospy.Publisher('~result', String, queue_size=1) #35

        # Create a subscriber (创建订阅者) #37
        rospy.Subscriber('/vocal_detect/asr_result', String, self.asr_callback) #38
        rospy.Subscriber(self.camera_topic, Image, self.image_callback) #39

        # Create a service (创建服务) #41
        rospy.Service('~set_model', SetModel, self.set_model_srv) #42
        rospy.Service('~set_prompt', SetString, self.set_prompt_srv) #43
        rospy.Service('~set_llm_content', SetContent, self.set_llm_content_srv) #44
        rospy.Service('~set_vllm_content', SetContent, self.set_vllm_content_srv) #45

        rospy.Service('/record_chat', SetBool, self.record_chat) #47
        rospy.Service('/get_chat', Trigger, self.get_chat) #48
        rospy.Service('/clear_chat', Empty, self.clear_chat) #49
        
        rospy.Service('~init_finish', Trigger, self.get_node_state) #51

        rospy.loginfo('\033[1;32m%s\033[0m' % 'AgentProcess 启动成功') #53

    def get_node_state(self, request, response): #55
        """ #56
        Handle the ~init_finish service request and return the node status(处理 ~init_finish 服务请求，返回节点状态。) #57
        """ #58
        response.success = True #59
        response.message = "节点初始化完成" #60
        return response #61

    def record_chat(self, request, response): #63
        rospy.loginfo('record chat') #64
        self.start_record_chat = request.data #65
        response.success = True #66

        return response #68
    
    def get_chat(self, request): #70
        rospy.loginfo('get chat') #71
        response = TriggerResponse()  #72
        response.message = self.chat_text.rstrip(",") #73
        response.success = True #74
        return response #75


    def clear_chat(self, request): #78
        rospy.loginfo('clear chat') #79
        self.chat_text = '' #80
        self.start_record_chat = False #81
        response.success = True #82
        return response #83

    def asr_callback(self, msg): #85
        """ #86
        Callback function to process voice recognition results(处理语音识别结果的回调函数。) #87
        """ #88
        # Process voice recognition results (处理语音识别结果) #89
        if msg.data != '': #90
            rospy.loginfo('\033[1;32m%s\033[0m' % 'thinking..') #91
            if self.start_record_chat: #92
                self.chat_text += msg.data + ',' #93
            
            res = '' #95
            if self.model_type == 'llm': #96
                res = self.client.llm(msg.data, self.prompt, model=self.model) #97
                rospy.loginfo('\033[1;32m%s\033[0m' % 'publish llm result:' + str(res)) #98
            elif self.model_type == 'vllm': #99
                try: #100
                    image = self.image_queue.get(block=True, timeout=5) #101
                except queue.Empty: #102
                    rospy.logwarn('图像队列为空，无法获取图像') #103
                    return #104
                res = self.client.vllm(msg.data, image, prompt=self.prompt, model=self.model) #105
                rospy.loginfo('\033[1;32m%s\033[0m' % 'publish vllm result:' + str(res)) #106
            
            result_msg = String() #108
            result_msg.data = res #109
            self.result_pub.publish(result_msg) #110
        
        else: #112
            rospy.loginfo('\033[1;32m%s\033[0m' % 'asr result none') #113

    def image_callback(self, ros_image): #115
        """ #116
        Callback function to process camera images(处理摄像头图像的回调函数。) #117
        """ #118
        try: #119
            cv_image = self.bridge.imgmsg_to_cv2(ros_image, "bgr8") #120
        except Exception as e: #121
            rospy.logerr('图像转换错误: %s' % str(e)) #122
            return #123
        bgr_image = np.array(cv_image, dtype=np.uint8) #124
        if self.image_queue.full(): #125
            # If the queue is full, discard the oldest image (如果队列已满，丢弃最旧的图像) #126
            try: #127
                self.image_queue.get_nowait() #128
            except queue.Empty: #129
                pass #130
        # Put the image into the queue (将图像放入队列) #131
        self.image_queue.put(bgr_image) #132

    def set_model_srv(self, request): #134
        """ #135
        Service handler for setting the model type to be called (vllm/llm)(设置调用的模型类型（vllm/llm）的服务处理函数。) #136
        """ #137
        rospy.loginfo('\033[1;32m%s\033[0m' % '设置模型类型') #138
        self.model = request.model #139
        self.model_type = request.model_type #140
        self.client = speech.OpenAIAPI(request.api_key, request.base_url) #141

        response = SetModelResponse() #143
        response.success = True #144
        response.message = "模型类型设置成功" #145
        return response #146

    def set_prompt_srv(self, request): #148
        # Set the prompt for the large model (设置大模型的prompt) #149
        rospy.loginfo('\033[1;32m%s\033[0m' % 'set prompt') #150
        self.prompt = request.data #151

        response = SetStringResponse() #153
        response.success = True #154
        return response #155

    def set_llm_content_srv(self, request, response): #157
        # Pass input text to the agent to generate a response (输入文本传给智能体让他来回答) #158
        rospy.loginfo('\033[1;32m%s\033[0m' % 'thinking...') #159
        client = speech.OpenAIAPI(request.api_key, request.base_url) #160

        response.message = client.llm(request.query, request.prompt, model=request.model)  #162
        response.success = True #163
        return response #164

    def set_vllm_content_srv(self, request, response): #166
        # Send prompt and input to the vision agent to get a response (输入提示词和文本，视觉智能体返回回答) #167
        image = self.image_queue.get(block=True) #168
        client = speech.OpenAIAPI(request.api_key, request.base_url) #169
        res = client.vllm(request.query, image, prompt=request.prompt, model=request.model)  #170
        response.message = res #171
        response.success = True #172
        return response #173

def main(): #175
    try: #176
        node = AgentProcess('agent_process') #177
        rospy.spin() #178
    except KeyboardInterrupt: #179
        rospy.loginfo('节点已关闭') #180
    except Exception as e: #181
        rospy.logerr('出现异常: %s' % str(e)) #182

if __name__ == "__main__": #184
    main() #185
