# coding: utf-8

import sys

import os
import rospy
import rosbag
import csv
import cv2
from cv_bridge import CvBridge, CvBridgeError
import argparse
import struct
import subprocess
from subprocess import Popen, PIPE

bridge = CvBridge()


def parsebag(f, output_dir, target_topic):
    rosbag_name = f[f.rindex("/"):-4]
    output_path = output_dir + rosbag_name

    with rosbag.Bag(f, 'r') as bag:
        count = 0
        for key, val in bag.get_type_and_topic_info()[1].items():
            if val[0] == 'sensor_msgs/Image':

                for topic, msg, t in bag.read_messages(key):
                    if topic== target_topic:
                        img = bridge.imgmsg_to_cv2(msg, 'bgr8')
                        # img = bridge.compressed_imgmsg_to_cv2(msg, 'bgr8')
                        timestr = msg.header.stamp.to_nsec()
                        image_name = str(timestr) + ".jpg"
                        if not os.path.exists(output_path + '_' + 'image'):
                            os.makedirs(output_path + '_' + 'image')
                        output_path_img = os.path.join(output_path + '_' + 'image', image_name)
                        cv2.imwrite(output_path_img, img)
                        # print(count)
                        count = count + 1
    return count

def main():
    input_dir = sys.argv[1]
    #input_dir='/media/dell/BFAC-F22B/data/pujin_datareturn/pujin_robot-01_data_merge_2023-12-20-02-22-18_obstacledetection_10.bag'
    output_dir = sys.argv[2]
    #output_dir = '/media/dell/BFAC-F22B/data/pujin_datareturn'
    topic = '/camera/color/image_raw'  # /camera/color/image_raw

    num_count = parsebag(input_dir, output_dir, topic)
    print("解析完成")

    dirt = os.path.join(output_dir, input_dir.split('/')[-1].split('.')[0] + '_' + 'image')
    # print(dirt)

    bagtime = int(input_dir.split('/')[-1].split('.')[0].split('_')[-1])
    hz = str(float(num_count / bagtime))
    print(hz)

    command = ['ffmpeg', '-f', 'image2', '-r', hz, '-pattern_type', 'glob', '-i'] + [dirt + '/*.jpg'] + ['-tag:v',
                                                                                                         'avc1',
                                                                                                         '-y'] + [
                  dirt + '/camera.mp4']
    p = Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    p.communicate()
    list1 = os.listdir(dirt)
    for i in list1:
        if i.split('.')[-1] == 'jpg':
            os.remove(os.path.join(dirt, i))


if __name__ == '__main__':
    main()


def parse(input_bag_file, output_mp4_dir):
    input_dir = input_bag_file
    # input_dir = '/media/dell/BFAC-F22B/rapid_accel_2023-11-15-10-33-50.bag'
    output_dir = output_mp4_dir
    # output_dir = '/media/dell/BFAC-F22B/'
    topic = '/camera/color/image_raw'

    num_count = parsebag(input_dir, output_dir, topic)
    print("解析完成")

    dirt = os.path.join(output_dir, input_dir.split('/')[-1].split('.')[0] + '_' + 'image')
    # print(dirt)

    bagtime = int(input_dir.split('/')[-1].split('.')[0].split('_')[-1])
    hz = str(float(num_count / bagtime))
    print(hz)

    command = ['ffmpeg', '-f', 'image2', '-r', hz, '-pattern_type', 'glob', '-i'] + [dirt + '/*.jpg'] + ['-tag:v',
                                                                                                         'avc1',
                                                                                                         '-y'] + [
                  dirt + '/camera.mp4']
    p = Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    p.communicate()
    list1 = os.listdir(dirt)
    for i in list1:
        if i.split('.')[-1] == 'jpg':
            os.remove(os.path.join(dirt, i))
    return dirt