'''
Description: 
Version: 1.0
Autor: Sun Yalun
Date: 2024-09-05 14:53:32
LastEditors: Sun Yalun
LastEditTime: 2024-09-05 15:46:54
'''
#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy
import rosbag
import sys
import os
import signal
import subprocess
from geometry_msgs.msg import PoseWithCovarianceStamped, Twist
from move_base_msgs.msg import MoveBaseActionGoal
from datetime import datetime


class RosbagRecorder:
    def __init__(self, bag_dir, count):
        rospy.init_node('rosbag_recorder', anonymous=True)
        rospy.set_param('/use_sim_time', False)
        self.goal_sub = rospy.Subscriber('/move_base/goal', MoveBaseActionGoal, self.goal_callback)
        self.pose_sub = rospy.Subscriber('/amcl_pose', PoseWithCovarianceStamped, self.pose_callback)
        self.cmd_vel_sub = rospy.Subscriber('/cmd_vel', Twist, self.cmd_vel_callback)

        self.goal_position = None
        self.bag_process = None
        self.recording = False
        self.pose_position = None
        self.cmd_vel = None
        self.bag_dir = bag_dir
        self.bag_count = 1
        self.total_count = int(count)

    def goal_callback(self, msg):
        # 如果当前已录制的 bag 文件数量大于等于传入的 count 参数，就不再录制
        if self.bag_count > self.total_count:
            rospy.loginfo("Recording limit reached. Exiting...")
            self.shutdown_node()
            return
        # 更新时间戳为当前时间
        msg.header.stamp = rospy.Time.now()
        self.goal_position = msg.goal.target_pose.pose.position
        rospy.loginfo("Goal received, starting rosbag recording...")

        bag_file_name = f"test-{self.bag_count}.bag"
        bag_file_path = os.path.join(self.bag_dir, bag_file_name)
        self.bag_path = bag_file_path  
        self.start_recording()
    def pose_callback(self, msg):
        # 更新时间戳为当前时间
        msg.header.stamp = rospy.Time.now()
        self.pose_position = msg.pose.pose.position
        self.check_stop_condition()

    def cmd_vel_callback(self, msg):
        # 更新时间戳为当前时间
        # msg.header.stamp = rospy.Time.now()
        self.cmd_vel = msg
        self.check_stop_condition()

    def start_recording(self):
        if not self.recording:
            self.bag_process = subprocess.Popen(['rosbag', 'record', '-O', self.bag_path, '/amcl_pose', '/imu', '/obstacle_detection', '/odom', '/sys_info'], preexec_fn=os.setsid)
            self.recording = True

    def stop_recording(self):
        if self.recording and self.bag_process:
            rospy.loginfo("Stopping rosbag recording...")
            # os.killpg(os.getpgid(self.bag_process.pid), signal.SIGTERM)  # 首先发送SIGTERM信号
            os.killpg(os.getpgid(self.bag_process.pid), signal.SIGINT)  # 首先发送SIGINT信号
            try:
                self.bag_process.wait(timeout=10)  # 等待进程终止，如果超过10秒，进行强制终止
            except subprocess.TimeoutExpired:
                rospy.logwarn("Rosbag process did not terminate, sending SIGKILL...")
                os.killpg(os.getpgid(self.bag_process.pid), signal.SIGKILL)  # 强制终止
                self.bag_process.wait()  # 再次等待确保进程已经被终止                      
            self.recording = False
            self.bag_count += 1

            # 如果录制的包数量达到总数，停止整个节点
            if self.bag_count > self.total_count:
                rospy.loginfo("All recordings completed. Exiting...")
                rospy.sleep(2)  # 等待一秒，确保rosbag_recorder节点已经退出
                self.shutdown_node()

    def check_stop_condition(self):
        if self.goal_position and self.pose_position and self.cmd_vel:
            distance = self.calculate_distance(self.pose_position, self.goal_position)
            if distance < 0.3 and self.cmd_vel.linear.x == 0 and self.cmd_vel.linear.y == 0:
                self.stop_recording()

    def calculate_distance(self, pos1, pos2):
        return ((pos1.x - pos2.x) ** 2 + (pos1.y - pos2.y) ** 2) ** 0.5

    def shutdown_node(self):
        """停止ROS节点和当前进程"""
        rospy.loginfo("Shutting down ROS node and killing the process.")
        rospy.sleep(2)  # 等待2秒，确保rosbag_recorder节点已经退出
        rospy.signal_shutdown("Completed all recordings.")
        sys.exit(0)  # 终止Python脚本

if __name__ == '__main__':
    try:
        bag_folder_path = sys.argv[1]
        count = sys.argv[2]
        recorder = RosbagRecorder(bag_folder_path, count)
        rospy.spin()
    except rospy.ROSInterruptException:
        pass