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

//! Subject CRUD operations

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

use crate::error::{DatabaseError, Result};

/// Subject data structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Subject {
    /// Unique identifier
    pub id: String,
    /// Protocol this subject belongs to
    pub protocol_id: String,
    /// Subject label (e.g., "sub-01")
    pub label: String,
    /// Path to FreeSurfer anatomy directory
    pub anatomy_path: Option<String>,
    /// Demographics (age, sex, handedness, etc.)
    pub demographics: HashMap<String, serde_json::Value>,
    /// Creation timestamp
    pub created_at: DateTime<Utc>,
}

impl Subject {
    /// Create a new subject
    pub fn new(protocol_id: &str, label: &str) -> Self {
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            protocol_id: protocol_id.to_string(),
            label: label.to_string(),
            anatomy_path: None,
            demographics: HashMap::new(),
            created_at: Utc::now(),
        }
    }

    /// Create subject in database
    pub fn insert(&self, conn: &rusqlite::Connection) -> Result<()> {
        let demographics_json = serde_json::to_string(&self.demographics)?;

        conn.execute(
            "INSERT INTO subjects (id, protocol_id, label, anatomy_path, demographics_json, created_at)
             VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
            rusqlite::params![
                self.id,
                self.protocol_id,
                self.label,
                self.anatomy_path,
                demographics_json,
                self.created_at.to_rfc3339(),
            ],
        )?;

        Ok(())
    }

    /// Load subject from database by ID
    pub fn load(conn: &rusqlite::Connection, id: &str) -> Result<Self> {
        let mut stmt = conn.prepare(
            "SELECT id, protocol_id, label, anatomy_path, demographics_json, created_at
             FROM subjects WHERE id = ?1",
        )?;

        let subject = stmt
            .query_row([id], |row| {
                let demographics_json: String = row.get(4)?;
                let demographics: HashMap<String, serde_json::Value> =
                    serde_json::from_str(&demographics_json).unwrap_or_default();

                Ok(Subject {
                    id: row.get(0)?,
                    protocol_id: row.get(1)?,
                    label: row.get(2)?,
                    anatomy_path: row.get(3)?,
                    demographics,
                    created_at: DateTime::parse_from_rfc3339(&row.get::<_, String>(5)?)
                        .map_or_else(|_| Utc::now(), |dt| dt.with_timezone(&Utc)),
                })
            })
            .map_err(|_| DatabaseError::SubjectNotFound { id: id.to_string() })?;

        Ok(subject)
    }

    /// Update subject in database
    pub fn update(&self, conn: &rusqlite::Connection) -> Result<()> {
        let demographics_json = serde_json::to_string(&self.demographics)?;

        conn.execute(
            "UPDATE subjects
             SET label = ?2, anatomy_path = ?3, demographics_json = ?4
             WHERE id = ?1",
            rusqlite::params![self.id, self.label, self.anatomy_path, demographics_json,],
        )?;

        Ok(())
    }

    /// Delete subject from database
    pub fn delete(conn: &rusqlite::Connection, id: &str) -> Result<()> {
        conn.execute("DELETE FROM subjects WHERE id = ?1", [id])?;
        Ok(())
    }

    /// List all subjects for a protocol
    pub fn list_by_protocol(
        conn: &rusqlite::Connection,
        protocol_id: &str,
    ) -> Result<Vec<Subject>> {
        let mut stmt = conn.prepare(
            "SELECT id, protocol_id, label, anatomy_path, demographics_json, created_at
             FROM subjects WHERE protocol_id = ?1 ORDER BY label",
        )?;

        let subjects = stmt
            .query_map([protocol_id], |row| {
                let demographics_json: String = row.get(4)?;
                let demographics: HashMap<String, serde_json::Value> =
                    serde_json::from_str(&demographics_json).unwrap_or_default();

                Ok(Subject {
                    id: row.get(0)?,
                    protocol_id: row.get(1)?,
                    label: row.get(2)?,
                    anatomy_path: row.get(3)?,
                    demographics,
                    created_at: DateTime::parse_from_rfc3339(&row.get::<_, String>(5)?)
                        .map_or_else(|_| Utc::now(), |dt| dt.with_timezone(&Utc)),
                })
            })?
            .collect::<std::result::Result<Vec<_>, _>>()?;

        Ok(subjects)
    }

    /// Check if subject exists
    pub fn exists(conn: &rusqlite::Connection, id: &str) -> Result<bool> {
        let count: i32 =
            conn.query_row("SELECT COUNT(*) FROM subjects WHERE id = ?1", [id], |row| {
                row.get(0)
            })?;
        Ok(count > 0)
    }

    /// Set anatomy path for subject
    pub fn set_anatomy_path(&mut self, path: Option<String>) {
        self.anatomy_path = path;
    }

    /// Add demographic field
    pub fn set_demographic(&mut self, key: &str, value: serde_json::Value) {
        self.demographics.insert(key.to_string(), value);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::protocol::Protocol;
    use crate::schema::create_schema;
    use std::path::PathBuf;

    fn setup_db() -> rusqlite::Connection {
        let conn = rusqlite::Connection::open_in_memory().unwrap();
        create_schema(&conn).unwrap();
        conn
    }

    #[test]
    fn test_create_subject() {
        let conn = setup_db();
        let protocol = Protocol::new("Test", PathBuf::from("/tmp"));
        protocol.insert(&conn).unwrap();

        let subject = Subject::new(&protocol.id, "sub-01");
        subject.insert(&conn).unwrap();

        assert!(Subject::exists(&conn, &subject.id).unwrap());
    }

    #[test]
    fn test_load_subject() {
        let conn = setup_db();
        let protocol = Protocol::new("Test", PathBuf::from("/tmp"));
        protocol.insert(&conn).unwrap();

        let mut subject = Subject::new(&protocol.id, "sub-01");
        subject.anatomy_path = Some("/path/to/freesurfer".to_string());
        subject.insert(&conn).unwrap();

        let loaded = Subject::load(&conn, &subject.id).unwrap();
        assert_eq!(loaded.label, "sub-01");
        assert_eq!(loaded.anatomy_path, Some("/path/to/freesurfer".to_string()));
    }

    #[test]
    fn test_update_subject() {
        let conn = setup_db();
        let protocol = Protocol::new("Test", PathBuf::from("/tmp"));
        protocol.insert(&conn).unwrap();

        let mut subject = Subject::new(&protocol.id, "sub-01");
        subject.insert(&conn).unwrap();

        subject.label = "sub-02".to_string();
        subject.set_demographic("age", serde_json::json!(25));
        subject.update(&conn).unwrap();

        let loaded = Subject::load(&conn, &subject.id).unwrap();
        assert_eq!(loaded.label, "sub-02");
        assert_eq!(loaded.demographics.get("age"), Some(&serde_json::json!(25)));
    }

    #[test]
    fn test_list_subjects() {
        let conn = setup_db();
        let protocol = Protocol::new("Test", PathBuf::from("/tmp"));
        protocol.insert(&conn).unwrap();

        Subject::new(&protocol.id, "sub-01").insert(&conn).unwrap();
        Subject::new(&protocol.id, "sub-02").insert(&conn).unwrap();

        let subjects = Subject::list_by_protocol(&conn, &protocol.id).unwrap();
        assert_eq!(subjects.len(), 2);
    }

    #[test]
    fn test_cascade_delete() {
        let conn = setup_db();
        let protocol = Protocol::new("Test", PathBuf::from("/tmp"));
        protocol.insert(&conn).unwrap();

        let subject = Subject::new(&protocol.id, "sub-01");
        subject.insert(&conn).unwrap();

        // Delete protocol should cascade to subjects
        Protocol::delete(&conn, &protocol.id).unwrap();
        assert!(!Subject::exists(&conn, &subject.id).unwrap());
    }
}