rs-bag 0.1.0

use futures::future::try_join_all;
use std::collections::HashSet;
use std::path::Path;
use std::time::Duration;

use crate::Error;
use crate::MessageReader;
use crate::MessageWriter;
use crate::Metadata;
use crate::Result;
use crate::Topic;
use crate::message::MessageReadMut;
use crate::storage::{ConnectionType, Db3File, SplittedRosbag, find_db3_files};
use ConnectionType::{Atomic, Splitted};

type Topics = HashSet<Topic>;

pub struct Rosbag {
    connections: ConnectionType,
    metadata_path: std::path::PathBuf,
    topics: HashSet<Topic>,
}

impl Rosbag {
    pub async fn open<P: AsRef<std::path::Path>>(path: P) -> Result<Self> {
        let path = path.as_ref();
        if !path.is_dir() {
            return Err(Error::BagDirNotFound(path.to_path_buf()));
        }

        let db3_path = find_db3_files(path)?
            .into_iter()
            .next()
            .ok_or_else(|| Error::NoDb3File(path.to_path_buf()))?;

        let db3 = Db3File::open(&db3_path).await?;
        let topics = db3.get_topics().await?;
        let conn = Atomic(db3);

        let metadata_path = {
            let mut path = path.to_path_buf();
            path.push("metadata.yaml");
            path
        };

        Ok(Self {
            connections: conn,
            metadata_path,
            topics,
        })
    }

    pub async fn open_splitted<P: AsRef<std::path::Path>, I: IntoIterator<Item = P>>(
        paths: I,
        metadata_path: P,
        split_after: Duration,
    ) -> Result<Self> {
        let connections = try_join_all(paths.into_iter().map(|p| Db3File::open(p))).await?;

        let mut topics: Topics = HashSet::with_capacity(50);
        for c in &connections {
            c.get_topics().await?.into_iter().for_each(|t| {
                topics.insert(t);
            });
        }

        Ok(Self {
            connections: Splitted(SplittedRosbag {
                connections,
                split_after,
            }),
            metadata_path: metadata_path.as_ref().to_path_buf(),
            topics,
        })
    }

    pub async fn build<P: AsRef<std::path::Path>>(path: P) -> BagBuilder<P> {
        BagBuilder::new(path)
    }

    pub fn read(&self) -> MessageReader<'_> {
        MessageReader::create(&self.connections)
    }

    pub fn write(&self) -> MessageWriter<'_> {
        MessageWriter::create(&self.connections)
    }

    pub fn read_mut(&self) -> MessageReadMut {
        todo!("Create a MessageReadMut reader for mutable access")
    }

    pub fn topics(&self) -> &HashSet<Topic> {
        &self.topics
    }

    pub fn get_metadata(&self) -> Result<Metadata> {
        Metadata::create(&self.topics, self.connections.iter(), &self.metadata_path)
    }
}

pub struct BagBuilder<P: AsRef<std::path::Path>> {
    path: P,
    split_after: Option<Duration>,
}

impl<P: AsRef<std::path::Path>> BagBuilder<P> {
    pub fn new(path: P) -> Self {
        Self {
            path,
            split_after: None,
        }
    }

    pub async fn create(self) -> Result<Rosbag> {
        let dir = self.path.as_ref();
        let split_after = self.split_after;
        let metadata_path = {
            let mut path = dir.to_path_buf();
            path.push("metadata.yaml");
            path
        };

        create_dir_if_needed(dir)?;
        ensure_empty(dir)?;

        let db3_name = initial_db3_file_name(dir)?;
        let db3_path = dir.join(&db3_name);
        let connection = Db3File::init(db3_path).await?;

        Ok(Rosbag {
            connections: match split_after {
                Some(split_after) => ConnectionType::Splitted(SplittedRosbag {
                    connections: vec![connection],
                    split_after,
                }),
                None => ConnectionType::Atomic(connection),
            },
            topics: std::collections::HashSet::new(),
            metadata_path,
        })
    }

    pub fn split_after(self, duration: Duration) -> Self {
        Self {
            path: self.path,
            split_after: Some(duration),
        }
    }
}

fn create_dir_if_needed(dir: &Path) -> Result<()> {
    std::fs::create_dir_all(dir)?;
    Ok(())
}

fn ensure_empty(dir: &Path) -> Result<()> {
    if dir.join("metadata.yaml").is_file() {
        return Err(Error::BagHasMetadata(dir.to_path_buf()));
    }
    if !find_db3_files(dir)?.is_empty() {
        return Err(Error::BagHasDb3Files(dir.to_path_buf()));
    }
    Ok(())
}

fn initial_db3_file_name(dir: &Path) -> Result<String> {
    let dir_name = dir
        .file_name()
        .and_then(|n| n.to_str())
        .ok_or_else(|| Error::NoDirectoryName(dir.to_path_buf()))?;
    Ok(format!("{dir_name}_0.db3"))
}

#[cfg(test)]
mod test {
    use std::env::current_dir;

    use crate::Rosbag;

    #[tokio::test]
    async fn open_rosbag() {
        let mut p = current_dir().expect("Could not get cwd");
        p.push("test_bag");
        Rosbag::open(p).await.expect("Could not open rosbag");
    }

    #[tokio::test]
    async fn topics() {
        let mut p = current_dir().expect("Could not get cwd");
        p.push("test_bag");
        let bag = Rosbag::open(p).await.expect("Could not open rosbag");

        assert_eq!(bag.topics().len(), 1, "Expected one topic in rosbag");
        assert_eq!(
            bag.topics().iter().next().unwrap().name,
            String::from("/image_raw"),
            "Wrong topic name"
        );
        assert_eq!(
            bag.topics().iter().next().unwrap().r#type,
            String::from("sensor_msgs/msg/Image"),
            "Wrong topic name"
        )
    }
}