rustytorch/crates/specialized/rtx-neuro-db/src/database.rs

//! Async database wrapper using tokio-rusqlite

use std::path::{Path, PathBuf};
use tokio_rusqlite::Connection;

use crate::error::Result;
use crate::protocol::Protocol;
use crate::schema::{check_schema_version, create_schema};
use crate::subject::Subject;

#[cfg(test)]
use crate::schema::SCHEMA_VERSION;

/// Async database handle for protocol persistence
pub struct NeuroDatabase {
    conn: Connection,
    db_path: PathBuf,
}

impl NeuroDatabase {
    /// Open or create a database at the specified path
    pub async fn open(path: impl AsRef<Path>) -> Result<Self> {
        let db_path = path.as_ref().to_path_buf();
        let conn = Connection::open(&db_path).await?;

        // Check and initialize schema
        let needs_init = conn
            .call(|conn| {
                let version = check_schema_version(conn);
                Ok(version.ok().flatten().is_none())
            })
            .await?;

        if needs_init {
            conn.call(|conn| {
                Ok(create_schema(conn)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?)
            })
            .await?;
        }

        Ok(Self { conn, db_path })
    }

    /// Open an in-memory database (for testing)
    pub async fn open_in_memory() -> Result<Self> {
        let conn = Connection::open_in_memory().await?;

        conn.call(|conn| {
            Ok(create_schema(conn)
                .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?)
        })
        .await?;

        Ok(Self {
            conn,
            db_path: PathBuf::from(":memory:"),
        })
    }

    /// Get the database file path
    pub fn path(&self) -> &Path {
        &self.db_path
    }

    /// Get the current schema version
    pub async fn schema_version(&self) -> Result<Option<i32>> {
        let version = self
            .conn
            .call(|conn| Ok(check_schema_version(conn).ok().flatten()))
            .await?;
        Ok(version)
    }

    // =========================================================================
    // Protocol operations
    // =========================================================================

    /// Create a new protocol
    pub async fn create_protocol(&self, name: &str, path: PathBuf) -> Result<Protocol> {
        let protocol = Protocol::new(name, path);
        let protocol_clone = protocol.clone();

        self.conn
            .call(move |conn| {
                Ok(protocol_clone
                    .insert(conn)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?)
            })
            .await?;

        Ok(protocol)
    }

    /// Load a protocol by ID
    pub async fn get_protocol(&self, id: &str) -> Result<Protocol> {
        let id = id.to_string();
        let protocol = self
            .conn
            .call(move |conn| {
                Ok(Protocol::load(conn, &id)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?)
            })
            .await?;
        Ok(protocol)
    }

    /// Update a protocol
    pub async fn update_protocol(&self, protocol: Protocol) -> Result<Protocol> {
        let mut protocol = protocol;
        let protocol_clone = protocol.clone();
        self.conn
            .call(move |conn| {
                let mut p = protocol_clone;
                Ok(p.update(conn)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?)
            })
            .await?;
        protocol.modified_at = chrono::Utc::now();
        Ok(protocol)
    }

    /// Delete a protocol
    pub async fn delete_protocol(&self, id: &str) -> Result<()> {
        let id = id.to_string();
        self.conn
            .call(move |conn| {
                Ok(Protocol::delete(conn, &id)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?)
            })
            .await?;
        Ok(())
    }

    /// List all protocols
    pub async fn list_protocols(&self) -> Result<Vec<Protocol>> {
        let protocols = self
            .conn
            .call(|conn| {
                Ok(Protocol::list(conn)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?)
            })
            .await?;
        Ok(protocols)
    }

    /// Check if a protocol exists
    pub async fn protocol_exists(&self, id: &str) -> Result<bool> {
        let id = id.to_string();
        let exists = self
            .conn
            .call(move |conn| {
                Ok(Protocol::exists(conn, &id)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?)
            })
            .await?;
        Ok(exists)
    }

    // =========================================================================
    // Subject operations
    // =========================================================================

    /// Create a new subject
    pub async fn create_subject(&self, protocol_id: &str, label: &str) -> Result<Subject> {
        let subject = Subject::new(protocol_id, label);
        let subject_clone = subject.clone();

        self.conn
            .call(move |conn| {
                Ok(subject_clone
                    .insert(conn)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?)
            })
            .await?;

        Ok(subject)
    }

    /// Load a subject by ID
    pub async fn get_subject(&self, id: &str) -> Result<Subject> {
        let id = id.to_string();
        let subject = self
            .conn
            .call(move |conn| {
                Ok(Subject::load(conn, &id)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?)
            })
            .await?;
        Ok(subject)
    }

    /// Update a subject
    pub async fn update_subject(&self, subject: Subject) -> Result<Subject> {
        let subject_clone = subject.clone();
        self.conn
            .call(move |conn| {
                Ok(subject_clone
                    .update(conn)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?)
            })
            .await?;

        Ok(subject)
    }

    /// Delete a subject
    pub async fn delete_subject(&self, id: &str) -> Result<()> {
        let id = id.to_string();
        self.conn
            .call(move |conn| {
                Ok(Subject::delete(conn, &id)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?)
            })
            .await?;
        Ok(())
    }

    /// List all subjects for a protocol
    pub async fn list_subjects(&self, protocol_id: &str) -> Result<Vec<Subject>> {
        let protocol_id = protocol_id.to_string();
        let subjects = self
            .conn
            .call(move |conn| {
                Ok(Subject::list_by_protocol(conn, &protocol_id)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?)
            })
            .await?;
        Ok(subjects)
    }

    /// Check if a subject exists
    pub async fn subject_exists(&self, id: &str) -> Result<bool> {
        let id = id.to_string();
        let exists = self
            .conn
            .call(move |conn| {
                Ok(Subject::exists(conn, &id)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?)
            })
            .await?;
        Ok(exists)
    }

    /// Set anatomy path for a subject
    pub async fn set_subject_anatomy(&self, id: &str, path: Option<String>) -> Result<Subject> {
        let id = id.to_string();
        let subject = self
            .conn
            .call(move |conn| {
                let mut subject = Subject::load(conn, &id)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;
                subject.set_anatomy_path(path);
                subject
                    .update(conn)
                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;
                Ok(subject)
            })
            .await?;
        Ok(subject)
    }

    // =========================================================================
    // Utility operations
    // =========================================================================

    /// Execute a raw SQL query (for advanced use cases)
    pub async fn execute_raw(&self, sql: &str) -> Result<usize> {
        let sql = sql.to_string();
        let count = self
            .conn
            .call(move |conn| Ok(conn.execute(&sql, [])?))
            .await?;
        Ok(count)
    }

    /// Close the database connection
    pub async fn close(self) -> Result<()> {
        self.conn.close().await?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_open_in_memory() {
        let db = NeuroDatabase::open_in_memory().await.unwrap();
        let version = db.schema_version().await.unwrap();
        assert_eq!(version, Some(SCHEMA_VERSION));
    }

    #[tokio::test]
    async fn test_protocol_crud() {
        let db = NeuroDatabase::open_in_memory().await.unwrap();

        // Create
        let protocol = db
            .create_protocol("Test Protocol", PathBuf::from("/tmp/test"))
            .await
            .unwrap();

        assert!(db.protocol_exists(&protocol.id).await.unwrap());

        // Read
        let loaded = db.get_protocol(&protocol.id).await.unwrap();
        assert_eq!(loaded.name, "Test Protocol");

        // Update
        let mut updated = loaded;
        updated.name = "Updated Protocol".to_string();
        let updated = db.update_protocol(updated).await.unwrap();
        assert_eq!(updated.name, "Updated Protocol");

        // List
        let protocols = db.list_protocols().await.unwrap();
        assert_eq!(protocols.len(), 1);

        // Delete
        db.delete_protocol(&protocol.id).await.unwrap();
        assert!(!db.protocol_exists(&protocol.id).await.unwrap());
    }

    #[tokio::test]
    async fn test_subject_crud() {
        let db = NeuroDatabase::open_in_memory().await.unwrap();

        let protocol = db
            .create_protocol("Test", PathBuf::from("/tmp"))
            .await
            .unwrap();

        // Create subject
        let subject = db.create_subject(&protocol.id, "sub-01").await.unwrap();
        assert!(db.subject_exists(&subject.id).await.unwrap());

        // Set anatomy
        let updated = db
            .set_subject_anatomy(&subject.id, Some("/path/to/fs".to_string()))
            .await
            .unwrap();
        assert_eq!(updated.anatomy_path, Some("/path/to/fs".to_string()));

        // List subjects
        let subjects = db.list_subjects(&protocol.id).await.unwrap();
        assert_eq!(subjects.len(), 1);

        // Delete subject
        db.delete_subject(&subject.id).await.unwrap();
        assert!(!db.subject_exists(&subject.id).await.unwrap());
    }

    #[tokio::test]
    async fn test_cascade_delete() {
        let db = NeuroDatabase::open_in_memory().await.unwrap();

        let protocol = db
            .create_protocol("Test", PathBuf::from("/tmp"))
            .await
            .unwrap();

        let subject = db.create_subject(&protocol.id, "sub-01").await.unwrap();

        // Delete protocol should cascade to subjects
        db.delete_protocol(&protocol.id).await.unwrap();
        assert!(!db.subject_exists(&subject.id).await.unwrap());
    }
}