#!/usr/bin/env python3 #1
# coding=utf8 #2

import rospy #4
import threading #5
import time #6
import cv2 #7
from cv_bridge import CvBridge, CvBridgeError #8
from sensor_msgs.msg import Image #9
from std_msgs.msg import Bool #10
from std_srvs.srv import Trigger, TriggerResponse, SetBool, SetBoolResponse #11
from large_models.srv import SetColor, SetColorResponse #12
from puppy_controller import PuppyController, PuppyStatus #13

class ControlServiceNode: #15
    def __init__(self): #16
        self.controller = PuppyController(debug=False) #17

        self.bridge = CvBridge() #19
        rospy.Subscriber('/usb_cam/image_raw', Image, self.image_callback) #20
        # Subscribe to the wake-up topic, stop the interaction when receiving True （订阅唤醒话题，收到 True 时停止玩法） #21
        rospy.Subscriber('/vocal_detect/wakeup', Bool, self.wakeup_callback) #22

        self.latest_image = None #24
        self.image_lock = threading.Lock() #25

        rospy.Service('color_tracking/enter', Trigger, self.enter_callback) #27
        rospy.Service('color_tracking/enable_color_tracking', SetBool, self.enable_callback) #28
        rospy.Service('color_tracking/set_color', SetColor, self.set_color_callback) #29
        rospy.Service('color_tracking/stop', Trigger, self.stop_callback) #30

        rospy.loginfo('控制服务节点已启动，等待服务请求...') #32

    def image_callback(self, msg): #34
        try: #35
            cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8') #36
            with self.image_lock: #37
                self.latest_image = cv_image #38
        except CvBridgeError as e: #39
            rospy.logerr(f"转换图像失败: {e}") #40

    def get_latest_image(self): #42
        with self.image_lock: #43
            if self.latest_image is not None: #44
                return self.latest_image.copy() #45
            else: #46
                return None #47

    def wakeup_callback(self, msg): #49
        if msg.data: #50
            rospy.loginfo("收到唤醒信号，停止颜色跟踪玩法。") #51
            self.stop_play() #52

    def stop_play(self): #54
        with self.controller.status_lock: #55
            self.controller.color_tracking_enabled = False #56
            self.controller.status = PuppyStatus.STOP #57
        rospy.loginfo("颜色跟踪玩法已停止（通过唤醒信号）。") #58

    def enter_callback(self, req): #60
        resp = TriggerResponse() #61
        with self.controller.status_lock: #62
            if self.controller.status == PuppyStatus.STOP: #63
                resp.success = False #64
                resp.message = "机器人处于停止状态，无法进入颜色跟踪模式。" #65
                rospy.logwarn(resp.message) #66
                return resp #67

            if self.controller.target_color == 'None': #69
                resp.success = False #70
                resp.message = "未设置目标颜色，请先调用 /color_tracking/set_color。" #71
                rospy.logwarn(resp.message) #72
                return resp #73

            if self.controller.status not in [PuppyStatus.START, PuppyStatus.NORMAL, PuppyStatus.FOUND_TARGET]: #75
                self.controller.status = PuppyStatus.START #76
                resp.success = True #77
                resp.message = "进入颜色跟踪模式" #78
                rospy.loginfo("进入颜色跟踪模式") #79
            else: #80
                resp.success = False #81
                resp.message = "颜色跟踪模式已运行" #82
        return resp #83

    def enable_callback(self, req): #85
        resp = SetBoolResponse() #86
        with self.controller.status_lock: #87
            if self.controller.status == PuppyStatus.STOP: #88
                resp.success = False #89
                resp.message = "机器人处于停止状态，无法修改颜色跟踪。" #90
                rospy.logwarn(resp.message) #91
                return resp #92

            if req.data: #94
                if self.controller.target_color == 'None': #95
                    resp.success = False #96
                    resp.message = "未设置目标颜色，请先调用 /color_tracking/set_color。" #97
                    rospy.logwarn(resp.message) #98
                    return resp #99
                self.controller.color_tracking_enabled = True #100
                self.controller.status = PuppyStatus.START #101
                resp.success = True #102
                resp.message = "颜色跟踪已启用" #103
                rospy.loginfo(resp.message) #104
            else: #105
                self.controller.color_tracking_enabled = False #106
                self.controller.status = PuppyStatus.STOP #107
                resp.success = True #108
                resp.message = "颜色跟踪已禁用" #109
                rospy.loginfo(resp.message) #110
        return resp #111

    def set_color_callback(self, req): #113
        resp = SetColorResponse() #114
        if req.color in self.controller.lab_data: #115
            with self.controller.status_lock: #116
                self.controller.target_color = req.color #117
                self.controller.color_tracking_enabled = True #118
                self.controller.status = PuppyStatus.START #119
                resp.success = True #120
                resp.message = f"目标颜色设置为 {req.color}" #121
                rospy.loginfo(resp.message) #122
        else: #123
            resp.success = False #124
            resp.message = f"颜色 {req.color} 不在配置文件中" #125
            rospy.logwarn(resp.message) #126
        return resp #127

    def stop_callback(self, req): #129
        resp = TriggerResponse() #130
        self.stop_play() #131
        resp.success = True #132
        resp.message = "颜色跟踪已停止" #133
        rospy.loginfo(resp.message) #134
        return resp #135

def display_images(control_node): #137
    rospy.loginfo("图像显示线程启动") #138
    while not rospy.is_shutdown(): #139
        img = control_node.get_latest_image() #140
        if img is not None: #141
            processed = control_node.controller.process_image(img.copy()) #142
            cv2.imshow('Frame', processed) #143
            if cv2.waitKey(1) == 27:  # Press ESC to exit(ESC键退出) #144
                rospy.signal_shutdown('ESC pressed') #145
                break #146
        else: #147
            time.sleep(0.01) #148
    cv2.destroyAllWindows() #149
    rospy.loginfo("图像显示线程结束") #150

def main(): #152
    rospy.init_node('vllm_color_tracking') #153

    control_node = ControlServiceNode() #155

    display_thread = threading.Thread(target=display_images, args=(control_node,)) #157
    display_thread.daemon = True #158
    display_thread.start() #159

    rospy.spin() #161

    # Clean up on program exit （程序退出时清理） #163
    control_node.controller.running = False #164
    try: #165
        control_node.controller.puppy.move_stop(servo_run_time=500) #166
    except Exception as e: #167
        rospy.logerr(f"停止移动异常: {e}") #168
    control_node.controller.puppy.stop() #169
    rospy.loginfo("程序结束") #170

if __name__ == '__main__': #172
    main() #173
