import rospy
import rosbag
import argparse
import os

from tqdm import tqdm


def build_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('input', help='path to input .bag file')
    parser.add_argument('output', help='path to output .bag file')
    return parser


def main(args):
    inbag = rosbag.Bag(args.input)
    # os.makedirs(os.path.dirname(args.output))

    with tqdm(total=inbag.get_message_count()) as pbar:
        with rosbag.Bag(args.output, mode='w') as outbag:
            for topic, msg, t in inbag.read_messages():
                pbar.update()

                if topic == '/camera/camera_info':
                    msg.header.frame_id = 'camera_left'
                    outbag.write('/camera/color/camera_info', msg, t)
                    outbag.write('/camera/depth/camera_info', msg, t)
                elif topic == '/camera/color/image_raw':
                    msg.header.frame_id = 'camera_left'
                    outbag.write('/camera/color/image_rect', msg, t)
                elif topic == '/camera/depth/image_raw':
                    msg.header.frame_id = 'camera_left'
                    outbag.write('/camera/depth/image_rect', msg, t)
                elif topic == '/tf':
                    outbag.write(topic, msg, t)
                else:
                    outbag.write(topic, msg, t)


if __name__ == '__main__':
    parser = build_parser()
    args = parser.parse_args()
    main(args)