import rosbag
import glob
import cv2
import os
from cv_bridge import CvBridge
import argparse
from sensor_msgs.msg import CompressedImage

def convert_to_compressed_msg(cv_image, encoding_format='jpeg'):
    result, compressed_image = cv2.imencode('.' + encoding_format, cv_image)
    if not result:
        raise Exception("Failed to compress image")

    compressed_msg = CompressedImage()
    compressed_msg.format = encoding_format
    compressed_msg.data = compressed_image.tobytes()
    return compressed_msg

def process_bag(bag_path, save_path):
    bridge = CvBridge()
    bag = rosbag.Bag(bag_path, "r")
    new_bag_path = os.path.join(save_path, os.path.basename(bag_path))
    new_bag = rosbag.Bag(new_bag_path, "w")

    print("Processing {} and saving to {}".format(bag_path, new_bag_path))

    for topic, msg, t in bag.read_messages():
        if "cam" in topic and "image" in topic:

            cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
            

            compressed_msg = convert_to_compressed_msg(cv_image)
            compressed_msg.header = msg.header


            new_topic = topic+"/compressed"
            new_bag.write(new_topic, compressed_msg, t)

    bag.close()
    new_bag.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Convert raw image messages to compressed messages (batch)")
    parser.add_argument("bag_path", help="path to rosbag files", type=str)
    parser.add_argument("save_path", help="Path to output", type=str)
    args = parser.parse_args()

    if not os.path.exists(args.save_path):
        os.makedirs(args.save_path)

    bag_files = glob.glob(os.path.join(args.bag_path, "*.bag"))
    for bag_file in bag_files:
        process_bag(bag_file, args.save_path)