use async_trait::async_trait;
use chrono::{DateTime, Utc};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::core::error::StorageError;
use crate::core::session::Session;
use crate::storage::{Storage, Message, StorageStatistics};
pub struct MemoryStorage {
sessions: Arc<RwLock<HashMap<String, Session>>>,
messages: Arc<RwLock<HashMap<String, Vec<Message>>>>,
}
impl MemoryStorage {
pub fn new() -> Self {
Self {
sessions: Arc::new(RwLock::new(HashMap::new())),
messages: Arc::new(RwLock::new(HashMap::new())),
}
}
}
impl Default for MemoryStorage {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl Storage for MemoryStorage {
async fn create_session(&self, session: &Session) -> Result<String, StorageError> {
let mut sessions = self.sessions.write().await;
sessions.insert(session.id.clone(), session.clone());
Ok(session.id.clone())
}
async fn get_session(&self, id: &str) -> Result<Option<Session>, StorageError> {
let sessions = self.sessions.read().await;
Ok(sessions.get(id).cloned())
}
async fn update_session(&self, session: &Session) -> Result<(), StorageError> {
let mut sessions = self.sessions.write().await;
sessions.insert(session.id.clone(), session.clone());
Ok(())
}
async fn delete_session(&self, id: &str) -> Result<(), StorageError> {
let mut sessions = self.sessions.write().await;
let mut messages = self.messages.write().await;
sessions.remove(id);
messages.remove(id);
Ok(())
}
async fn list_sessions(&self) -> Result<Vec<Session>, StorageError> {
let sessions = self.sessions.read().await;
Ok(sessions.values().cloned().collect())
}
async fn save_message(&self, message: &Message) -> Result<String, StorageError> {
let mut messages = self.messages.write().await;
let session_messages = messages.entry(message.session_id.clone()).or_insert_with(Vec::new);
session_messages.push(message.clone());
Ok(message.id.clone())
}
async fn get_messages(&self, session_id: &str) -> Result<Vec<Message>, StorageError> {
let messages = self.messages.read().await;
Ok(messages.get(session_id).cloned().unwrap_or_default())
}
async fn delete_message(&self, id: &str) -> Result<(), StorageError> {
let mut messages = self.messages.write().await;
for session_messages in messages.values_mut() {
session_messages.retain(|msg| msg.id != id);
}
Ok(())
}
async fn search_messages(&self, query: &str) -> Result<Vec<Message>, StorageError> {
let messages = self.messages.read().await;
let mut results = Vec::new();
for session_messages in messages.values() {
for message in session_messages {
let content_str = serde_json::to_string(&message.content).unwrap_or_default();
if content_str.to_lowercase().contains(&query.to_lowercase()) {
results.push(message.clone());
}
}
}
Ok(results)
}
async fn get_statistics(&self) -> Result<StorageStatistics, StorageError> {
let sessions = self.sessions.read().await;
let messages = self.messages.read().await;
let total_sessions = sessions.len() as u64;
let total_messages: u64 = messages.values().map(|msgs| msgs.len() as u64).sum();
let oldest_session = sessions.values()
.map(|s| s.created_at)
.min();
let newest_session = sessions.values()
.map(|s| s.updated_at)
.max();
let avg_messages_per_session = if total_sessions > 0 {
total_messages as f64 / total_sessions as f64
} else {
0.0
};
Ok(StorageStatistics {
total_sessions,
total_messages,
storage_size: 0, oldest_session,
newest_session,
avg_messages_per_session,
})
}
async fn cleanup(&self, older_than: DateTime<Utc>) -> Result<u64, StorageError> {
let mut sessions = self.sessions.write().await;
let mut messages = self.messages.write().await;
let mut removed_count = 0;
let sessions_to_remove: Vec<String> = sessions
.iter()
.filter(|(_, session)| session.updated_at < older_than)
.map(|(id, _)| id.clone())
.collect();
for session_id in sessions_to_remove {
sessions.remove(&session_id);
messages.remove(&session_id);
removed_count += 1;
}
Ok(removed_count)
}
}