use chrono::{ TimeZone, Utc, }; use rusqlite::{ Connection, OptionalExtension, Result, Row, params, }; use crate::{ db::schema::{ deserialize_embedding, serialize_embedding, }, models::*, }; /// Insert a new message into the database pub fn insert_message(conn: &Connection, msg: &lib::Message) -> Result { let timestamp = msg.date.map(|dt| dt.timestamp()); let created_at = Utc::now().timestamp(); conn.execute( "INSERT INTO messages (chat_db_rowid, text, timestamp, is_from_me, created_at) VALUES (?1, ?2, ?3, ?4, ?5) ON CONFLICT(chat_db_rowid) DO NOTHING", params![msg.rowid, msg.text, timestamp, msg.is_from_me, created_at], )?; Ok(conn.last_insert_rowid()) } /// Get message ID by chat.db rowid pub fn get_message_id_by_chat_rowid(conn: &Connection, chat_db_rowid: i64) -> Result> { conn.query_row( "SELECT id FROM messages WHERE chat_db_rowid = ?1", params![chat_db_rowid], |row| row.get(0), ) .optional() } /// Get message by ID pub fn get_message(conn: &Connection, id: i64) -> Result { conn.query_row( "SELECT id, chat_db_rowid, text, timestamp, is_from_me, created_at FROM messages WHERE id = ?1", params![id], map_message_row, ) } fn map_message_row(row: &Row) -> Result { let timestamp: Option = row.get(3)?; let created_at: i64 = row.get(5)?; Ok(Message { id: row.get(0)?, chat_db_rowid: row.get(1)?, text: row.get(2)?, timestamp: timestamp.map(|ts| Utc.timestamp_opt(ts, 0).unwrap()), is_from_me: row.get(4)?, created_at: Utc.timestamp_opt(created_at, 0).unwrap(), }) } /// Insert message embedding pub fn insert_message_embedding( conn: &Connection, message_id: i64, embedding: &[f32], model_name: &str, ) -> Result { let embedding_bytes = serialize_embedding(embedding); let created_at = Utc::now().timestamp(); conn.execute( "INSERT INTO message_embeddings (message_id, embedding, model_name, created_at) VALUES (?1, ?2, ?3, ?4)", params![message_id, embedding_bytes, model_name, created_at], )?; Ok(conn.last_insert_rowid()) } /// Get message embedding pub fn get_message_embedding( conn: &Connection, message_id: i64, ) -> Result> { conn.query_row( "SELECT id, message_id, embedding, model_name, created_at FROM message_embeddings WHERE message_id = ?1", params![message_id], |row| { let embedding_bytes: Vec = row.get(2)?; let created_at: i64 = row.get(4)?; Ok(MessageEmbedding { id: row.get(0)?, message_id: row.get(1)?, embedding: deserialize_embedding(&embedding_bytes), model_name: row.get(3)?, created_at: Utc.timestamp_opt(created_at, 0).unwrap(), }) }, ) .optional() } /// Insert or get word embedding pub fn insert_word_embedding( conn: &Connection, word: &str, embedding: &[f32], model_name: &str, ) -> Result { let embedding_bytes = serialize_embedding(embedding); let created_at = Utc::now().timestamp(); conn.execute( "INSERT INTO word_embeddings (word, embedding, model_name, created_at) VALUES (?1, ?2, ?3, ?4) ON CONFLICT(word) DO NOTHING", params![word, embedding_bytes, model_name, created_at], )?; Ok(conn.last_insert_rowid()) } /// Get word embedding pub fn get_word_embedding(conn: &Connection, word: &str) -> Result> { conn.query_row( "SELECT id, word, embedding, model_name, created_at FROM word_embeddings WHERE word = ?1", params![word], |row| { let embedding_bytes: Vec = row.get(2)?; let created_at: i64 = row.get(4)?; Ok(WordEmbedding { id: row.get(0)?, word: row.get(1)?, embedding: deserialize_embedding(&embedding_bytes), model_name: row.get(3)?, created_at: Utc.timestamp_opt(created_at, 0).unwrap(), }) }, ) .optional() } /// Insert emotion classification pub fn insert_emotion( conn: &Connection, message_id: i64, emotion: &str, confidence: f64, model_version: &str, ) -> Result { let now = Utc::now().timestamp(); conn.execute( "INSERT INTO emotions (message_id, emotion, confidence, model_version, created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6)", params![message_id, emotion, confidence, model_version, now, now], )?; Ok(conn.last_insert_rowid()) } /// Update emotion classification pub fn update_emotion( conn: &Connection, message_id: i64, emotion: &str, confidence: f64, ) -> Result<()> { let updated_at = Utc::now().timestamp(); conn.execute( "UPDATE emotions SET emotion = ?1, confidence = ?2, updated_at = ?3 WHERE message_id = ?4", params![emotion, confidence, updated_at, message_id], )?; Ok(()) } /// Get emotion by message ID pub fn get_emotion_by_message_id(conn: &Connection, message_id: i64) -> Result> { conn.query_row( "SELECT id, message_id, emotion, confidence, model_version, created_at, updated_at FROM emotions WHERE message_id = ?1", params![message_id], map_emotion_row, ) .optional() } /// Get emotion by ID pub fn get_emotion_by_id(conn: &Connection, id: i64) -> Result> { conn.query_row( "SELECT id, message_id, emotion, confidence, model_version, created_at, updated_at FROM emotions WHERE id = ?1", params![id], map_emotion_row, ) .optional() } /// List all emotions with optional filters pub fn list_emotions( conn: &Connection, emotion_filter: Option<&str>, min_confidence: Option, limit: Option, offset: Option, ) -> Result> { let mut query = String::from( "SELECT id, message_id, emotion, confidence, model_version, created_at, updated_at FROM emotions WHERE 1=1", ); if emotion_filter.is_some() { query.push_str(" AND emotion = ?1"); } if min_confidence.is_some() { query.push_str(" AND confidence >= ?2"); } query.push_str(" ORDER BY created_at DESC"); if limit.is_some() { query.push_str(" LIMIT ?3"); } if offset.is_some() { query.push_str(" OFFSET ?4"); } let mut stmt = conn.prepare(&query)?; let emotions = stmt .query_map( params![ emotion_filter.unwrap_or(""), min_confidence.unwrap_or(0.0), limit.unwrap_or(1000), offset.unwrap_or(0), ], map_emotion_row, )? .collect::>>()?; Ok(emotions) } /// Delete emotion by ID pub fn delete_emotion(conn: &Connection, id: i64) -> Result<()> { conn.execute("DELETE FROM emotions WHERE id = ?1", params![id])?; Ok(()) } /// Count total emotions pub fn count_emotions(conn: &Connection) -> Result { conn.query_row("SELECT COUNT(*) FROM emotions", [], |row| row.get(0)) } fn map_emotion_row(row: &Row) -> Result { let created_at: i64 = row.get(5)?; let updated_at: i64 = row.get(6)?; Ok(Emotion { id: row.get(0)?, message_id: row.get(1)?, emotion: row.get(2)?, confidence: row.get(3)?, model_version: row.get(4)?, created_at: Utc.timestamp_opt(created_at, 0).unwrap(), updated_at: Utc.timestamp_opt(updated_at, 0).unwrap(), }) } /// Insert emotion training sample pub fn insert_training_sample( conn: &Connection, message_id: Option, text: &str, expected_emotion: &str, ) -> Result { let now = Utc::now().timestamp(); conn.execute( "INSERT INTO emotions_training_set (message_id, text, expected_emotion, created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5)", params![message_id, text, expected_emotion, now, now], )?; Ok(conn.last_insert_rowid()) } /// Get state value from daemon_state table pub fn get_state(conn: &Connection, key: &str) -> Result> { conn.query_row( "SELECT value FROM daemon_state WHERE key = ?1", params![key], |row| row.get(0), ) .optional() } /// Set state value in daemon_state table pub fn set_state(conn: &Connection, key: &str, value: &str) -> Result<()> { let updated_at = Utc::now().timestamp(); conn.execute( "INSERT INTO daemon_state (key, value, updated_at) VALUES (?1, ?2, ?3) ON CONFLICT(key) DO UPDATE SET value = ?2, updated_at = ?3", params![key, value, updated_at], )?; Ok(()) } /// Get last processed chat.db rowid from database or return 0 pub fn get_last_processed_rowid(conn: &Connection) -> Result { Ok(get_state(conn, "last_processed_rowid")? .and_then(|s| s.parse().ok()) .unwrap_or(0)) } /// Save last processed chat.db rowid to database pub fn save_last_processed_rowid(conn: &Connection, rowid: i64) -> Result<()> { set_state(conn, "last_processed_rowid", &rowid.to_string()) }