import cv2
import datetime
import os

def save_rtsp_to_video(rtsp_url, output_dir='output', duration_minutes=10):
    """
    将RTSP流保存为视频文件
    
    参数:
        rtsp_url: RTSP流地址
        output_dir: 输出目录
        duration_minutes: 每个视频文件的时长(分钟)
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 打开RTSP流
    cap = cv2.VideoCapture(rtsp_url)
    
    if not cap.isOpened():
        print("无法打开RTSP流")
        return
    
    # 获取视频的帧率和尺寸
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    if fps <= 0:
        fps = 25  # 默认帧率
    
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    # 计算最大帧数(按duration_minutes分钟分割)
    max_frames = fps * 60 * duration_minutes
    
    frame_count = 0
    file_count = 1
    
    # 创建第一个视频文件
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_file = os.path.join(output_dir, f"video_{timestamp}_part{file_count}.mp4")
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 或者使用'avc1'
    out = cv2.VideoWriter(output_file, fourcc, fps, (frame_width, frame_height))
    
    print(f"开始录制，保存到: {output_file}")
    
    try:
        while True:
            ret, frame = cap.read()
            
            if not ret:
                print("无法获取帧，可能流已断开")
                break
            
            # 写入帧
            out.write(frame)
            frame_count += 1
            
            # 如果达到最大帧数，创建新文件
            if frame_count >= max_frames:
                out.release()
                file_count += 1
                frame_count = 0
                
                timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                output_file = os.path.join(output_dir, f"video_{timestamp}_part{file_count}.mp4")
                out = cv2.VideoWriter(output_file, fourcc, fps, (frame_width, frame_height))
                print(f"创建新文件: {output_file}")
            
            # # 按q键退出(需要显示窗口时才有效)
            # if cv2.waitKey(1) & 0xFF == ord('q'):
            #     break
                
    except KeyboardInterrupt:
        print("用户中断录制")
    finally:
        cap.release()
        out.release()
        # cv2.destroyAllWindows()
        print("录制结束")

if __name__ == "__main__":
    # 示例RTSP URL - 替换为你的实际RTSP地址
    rtsp_url = "rtsp://10.0.0.61/live/video6"
    
    # 调用函数开始录制
    save_rtsp_to_video(rtsp_url, duration_minutes=10)