rs-bag 0.1.0

use std::{
    collections::{HashMap, HashSet},
    path::Path,
    time::Duration,
};

use turso::{Connection, Rows};

use self::sql::*;
use crate::Error;
use crate::Result;
use crate::Topic;
use crate::message::filter::MessageFilter;

pub mod sql;

pub enum ConnectionType {
    Atomic(Db3File),
    Splitted(SplittedRosbag),
}

impl ConnectionType {
    pub fn iter<'a>(&'a self) -> std::slice::Iter<'a, Db3File> {
        match self {
            ConnectionType::Atomic(db3_file) => std::slice::from_ref(db3_file).iter(),
            ConnectionType::Splitted(splitted_rosbag) => splitted_rosbag.connections.iter(),
        }
    }

    pub fn as_slice(&self) -> &[Db3File] {
        match self {
            ConnectionType::Atomic(db3_file) => std::slice::from_ref(db3_file),
            ConnectionType::Splitted(splitted_rosbag) => &splitted_rosbag.connections,
        }
    }
}

pub struct SplittedRosbag {
    pub connections: Vec<Db3File>,
    pub split_after: Duration,
}

#[derive(Debug)]
pub struct Db3File {
    pub connection: Connection,
    pub path: std::path::PathBuf,
    pub topic_ids: HashMap<String, i64>,
}

impl Db3File {
    pub async fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
        let path = path.as_ref();
        let connection = Self::open_connection(path).await?;
        Ok(Db3File {
            path: path.to_path_buf(),
            topic_ids: Self::read_topics(&connection).await?,
            connection,
        })
    }

    async fn read_topics(connection: &Connection) -> Result<HashMap<String, i64>> {
        let mut rows = connection.query(SQL_STATEMENT_SELECT_TOPIC_IDS, ()).await?;
        let mut topics = HashMap::new();
        while let Some(row) = rows.next().await? {
            topics.insert(row.get(1)?, row.get(0)?);
        }
        Ok(topics)
    }

    async fn open_connection(path: &Path) -> Result<Connection> {
        if !path.is_file() {
            return Err(Error::Db3FileNotFound(path.to_path_buf()));
        }

        let path_str = path
            .to_str()
            .ok_or_else(|| Error::NonUtf8Path(path.to_path_buf()))?;
        let db = turso::Builder::new_local(path_str).build().await?;
        let connection = db.connect()?;

        Ok(connection)
    }

    pub async fn init<P: AsRef<Path>>(path: P) -> Result<Self> {
        let path = path.as_ref();
        if path.exists() {
            return Err(Error::Db3FileAlreadyExists(path.to_path_buf()));
        }

        let path_str = path
            .to_str()
            .ok_or_else(|| Error::NonUtf8Path(path.to_path_buf()))?;
        let db = turso::Builder::new_local(path_str).build().await?;
        let connection = db.connect()?;

        // >>> Begin connection init
        connection.execute(SQL_STATEMENT_CREATE_TABLE, ()).await?;
        connection
            .execute(SQL_STATEMENT_CREATE_METADATA, ())
            .await?;
        connection.execute(SQL_STATEMENT_CREATE_TOPICS, ()).await?;
        connection
            .execute(SQL_STATEMENT_CREATE_MSG_DEFINITIONS, ())
            .await?;
        connection.execute(SQL_STATEMENT_CREATE_MSGS, ()).await?;
        connection.execute(SQL_STATEMENT_CREATE_INDEX, ()).await?;
        connection.execute(SQL_STATEMENT_INSERT_SCHEMA, ()).await?;
        // <<< End connection init

        Ok(Db3File {
            connection,
            path: path.to_path_buf(),
            topic_ids: HashMap::new(),
        })
    }

    pub async fn get_topics(&self) -> Result<HashSet<Topic>> {
        let mut rows = self
            .connection
            .query(SQL_STATEMENT_SELECT_TOPICS, ())
            .await?;
        let mut topics = HashSet::new();
        while let Some(row) = rows.next().await? {
            topics.insert(Topic::try_from(row)?);
        }
        Ok(topics)
    }

    pub fn topic_id(&self, topic: &str) -> Result<i64> {
        self.topic_ids
            .get(topic)
            .copied()
            .ok_or(Error::TopicNotFound(topic.to_string()))
    }

    pub async fn query_messages(&self, filter: &MessageFilter<'_>) -> Result<Rows> {
        let topic_count = filter.topics.as_ref().map(|topics| topics.len());
        let sql =
            sql_statement_message_filter(topic_count, filter.begin.is_some(), filter.end.is_some());
        let params = sql_params_message_filter(filter, &self.topic_ids)?;
        let params = turso::params::params_from_iter(params);

        let mut statement = self.connection.prepare(&sql).await?;
        let rows = statement.query(params).await?;
        Ok(rows)
    }

    pub async fn insert_message(&self, topic_id: i64, timestamp: i64, data: Vec<u8>) -> Result<()> {
        self.connection
            .execute(
                SQL_STATEMENT_INSERT_MSG,
                (turso::Value::Integer(topic_id), timestamp, data),
            )
            .await?;

        Ok(())
    }

    pub async fn time_range(&self) -> Result<(Option<i64>, Option<i64>)> {
        let mut rows = self
            .connection
            .query(SQL_STATEMENT_SELECT_TIME_RANGE, ())
            .await?;
        let row = rows
            .next()
            .await?
            .ok_or(turso::Error::QueryReturnedNoRows)?;

        // NULL on an empty messages table -> as_integer() yields None
        let min = row.get_value(0)?.as_integer().copied();
        let max = row.get_value(1)?.as_integer().copied();
        Ok((min, max))
    }
}

impl SplittedRosbag {
    pub fn iter(&self) -> std::slice::Iter<'_, Db3File> {
        return self.connections.iter();
    }
}

pub(crate) fn find_db3_files<P: AsRef<std::path::Path>>(dir: P) -> Result<Vec<std::path::PathBuf>> {
    let paths = std::fs::read_dir(dir)?
        .into_iter()
        .filter_map(|e| e.ok())
        .map(|e| e.path())
        .filter(|p| p.is_file() && p.extension().is_some_and(|ext| ext == "db3"))
        .collect();
    Ok(paths)
}