sochdb_query/

memory_compaction.rs

// Copyright 2025 Sushanth (https://github.com/sushanthpy)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Hierarchical Memory Compaction (Task 5)
//!
//! This module implements semantic memory compaction inspired by LSM-trees.
//! It manages tiered storage where older memories are summarized to maintain
//! bounded context while preserving semantic continuity.
//!
//! ## Architecture
//!
//! ```text
//! L0: Raw Episodes (recent, full detail)
//!     │
//!     ▼ Summarization
//! L1: Summaries (older, compressed)
//!     │
//!     ▼ Abstraction
//! L2: Abstractions (oldest, highly compressed)
//! ```
//!
//! ## Compaction Strategy
//!
//! - Episodes older than tier threshold are grouped by semantic similarity
//! - Each group is summarized (via LLM or extractive methods)
//! - Summaries are re-embedded for retrieval
//! - Growth is O(log_c T) where c = compaction ratio, T = total events

use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, RwLock};
use std::time::{SystemTime, UNIX_EPOCH};

// ============================================================================
// Configuration
// ============================================================================

/// Configuration for memory compaction
#[derive(Debug, Clone)]
pub struct MemoryCompactionConfig {
    /// Maximum episodes in L0 before compaction
    pub l0_max_episodes: usize,
    
    /// Maximum summaries in L1 before compaction
    pub l1_max_summaries: usize,
    
    /// Age threshold for L0 → L1 compaction (seconds)
    pub l0_age_threshold_secs: u64,
    
    /// Age threshold for L1 → L2 compaction (seconds)
    pub l1_age_threshold_secs: u64,
    
    /// Number of episodes to group for summarization
    pub group_size: usize,
    
    /// Similarity threshold for grouping (0.0 to 1.0)
    pub similarity_threshold: f32,
    
    /// Maximum tokens per summary
    pub max_summary_tokens: usize,
    
    /// Whether to re-embed summaries for retrieval
    pub reembed_summaries: bool,
    
    /// Compaction check interval (seconds)
    pub check_interval_secs: u64,
}

impl Default for MemoryCompactionConfig {
    fn default() -> Self {
        Self {
            l0_max_episodes: 1000,
            l1_max_summaries: 100,
            l0_age_threshold_secs: 3600,      // 1 hour
            l1_age_threshold_secs: 86400 * 7, // 1 week
            group_size: 10,
            similarity_threshold: 0.7,
            max_summary_tokens: 200,
            reembed_summaries: true,
            check_interval_secs: 300, // 5 minutes
        }
    }
}

impl MemoryCompactionConfig {
    /// Create config for aggressive compaction (testing/demos)
    pub fn aggressive() -> Self {
        Self {
            l0_max_episodes: 100,
            l1_max_summaries: 20,
            l0_age_threshold_secs: 60,
            l1_age_threshold_secs: 3600,
            group_size: 5,
            ..Default::default()
        }
    }
    
    /// Create config for long-running agents
    pub fn long_running() -> Self {
        Self {
            l0_max_episodes: 5000,
            l1_max_summaries: 500,
            l0_age_threshold_secs: 3600 * 6,      // 6 hours
            l1_age_threshold_secs: 86400 * 30,    // 30 days
            group_size: 20,
            ..Default::default()
        }
    }
}

// ============================================================================
// Memory Types
// ============================================================================

/// A raw episode (L0)
#[derive(Debug, Clone)]
pub struct Episode {
    /// Unique identifier
    pub id: String,
    
    /// Timestamp (seconds since epoch)
    pub timestamp: f64,
    
    /// Episode content (e.g., user message, tool call)
    pub content: String,
    
    /// Episode type
    pub episode_type: EpisodeType,
    
    /// Associated metadata
    pub metadata: HashMap<String, String>,
    
    /// Embedding vector (for similarity grouping)
    pub embedding: Option<Vec<f32>>,
    
    /// Token count (estimated or exact)
    pub token_count: usize,
}

/// Episode types
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum EpisodeType {
    /// User message
    UserMessage,
    /// Assistant response
    AssistantResponse,
    /// Tool call
    ToolCall,
    /// Tool result
    ToolResult,
    /// System event
    SystemEvent,
    /// Observation
    Observation,
}

/// A summary (L1)
#[derive(Debug, Clone)]
pub struct Summary {
    /// Unique identifier
    pub id: String,
    
    /// Summarized content
    pub content: String,
    
    /// IDs of episodes that were summarized
    pub source_episode_ids: Vec<String>,
    
    /// Time range covered
    pub time_range: (f64, f64),
    
    /// Summary embedding
    pub embedding: Option<Vec<f32>>,
    
    /// Token count
    pub token_count: usize,
    
    /// When this summary was created
    pub created_at: f64,
    
    /// Topics/themes extracted
    pub topics: Vec<String>,
}

/// An abstraction (L2)
#[derive(Debug, Clone)]
pub struct Abstraction {
    /// Unique identifier
    pub id: String,
    
    /// High-level abstraction content
    pub content: String,
    
    /// IDs of summaries that were abstracted
    pub source_summary_ids: Vec<String>,
    
    /// Time range covered
    pub time_range: (f64, f64),
    
    /// Abstraction embedding
    pub embedding: Option<Vec<f32>>,
    
    /// Token count
    pub token_count: usize,
    
    /// When this abstraction was created
    pub created_at: f64,
    
    /// Key insights
    pub insights: Vec<String>,
}

// ============================================================================
// Summarizer Trait
// ============================================================================

/// Trait for summarization backends
pub trait Summarizer: Send + Sync {
    /// Summarize a group of episodes into a single summary
    fn summarize_episodes(&self, episodes: &[Episode]) -> Result<String, CompactionError>;
    
    /// Summarize a group of summaries into an abstraction
    fn abstract_summaries(&self, summaries: &[Summary]) -> Result<String, CompactionError>;
    
    /// Extract topics/themes from content
    fn extract_topics(&self, content: &str) -> Vec<String>;
}

/// Extractive summarizer (no LLM required)
pub struct ExtractiveSummarizer {
    /// Maximum sentences to include
    pub max_sentences: usize,
    
    /// Whether to include timestamps
    pub include_timestamps: bool,
}

impl Default for ExtractiveSummarizer {
    fn default() -> Self {
        Self {
            max_sentences: 5,
            include_timestamps: true,
        }
    }
}

impl Summarizer for ExtractiveSummarizer {
    fn summarize_episodes(&self, episodes: &[Episode]) -> Result<String, CompactionError> {
        if episodes.is_empty() {
            return Ok(String::new());
        }
        
        let mut summary_parts = Vec::new();
        
        // Time range
        let first_ts = episodes.first().map(|e| e.timestamp).unwrap_or(0.0);
        let last_ts = episodes.last().map(|e| e.timestamp).unwrap_or(0.0);
        
        if self.include_timestamps {
            summary_parts.push(format!(
                "[{} episodes over {:.0} seconds]",
                episodes.len(),
                last_ts - first_ts
            ));
        }
        
        // Group by type and summarize
        let mut by_type: HashMap<EpisodeType, Vec<&Episode>> = HashMap::new();
        for episode in episodes {
            by_type.entry(episode.episode_type).or_default().push(episode);
        }
        
        // Extract key content from each type
        for (ep_type, eps) in by_type {
            let type_name = match ep_type {
                EpisodeType::UserMessage => "User messages",
                EpisodeType::AssistantResponse => "Responses",
                EpisodeType::ToolCall => "Tool calls",
                EpisodeType::ToolResult => "Tool results",
                EpisodeType::SystemEvent => "Events",
                EpisodeType::Observation => "Observations",
            };
            
            // Take first sentence from each, up to max_sentences
            let sentences: Vec<String> = eps
                .iter()
                .take(self.max_sentences)
                .filter_map(|e| e.content.split('.').next().map(|s| s.trim().to_string()))
                .filter(|s| !s.is_empty())
                .collect();
            
            if !sentences.is_empty() {
                summary_parts.push(format!("{}: {}", type_name, sentences.join("; ")));
            }
        }
        
        Ok(summary_parts.join("\n"))
    }
    
    fn abstract_summaries(&self, summaries: &[Summary]) -> Result<String, CompactionError> {
        if summaries.is_empty() {
            return Ok(String::new());
        }
        
        let mut abstraction_parts = Vec::new();
        
        // Time range
        let first_ts = summaries.iter().map(|s| s.time_range.0).fold(f64::MAX, f64::min);
        let last_ts = summaries.iter().map(|s| s.time_range.1).fold(f64::MIN, f64::max);
        
        abstraction_parts.push(format!(
            "[{} summaries, {:.1} hours span]",
            summaries.len(),
            (last_ts - first_ts) / 3600.0
        ));
        
        // Collect all topics
        let all_topics: Vec<&str> = summaries
            .iter()
            .flat_map(|s| s.topics.iter().map(|t| t.as_str()))
            .collect();
        
        if !all_topics.is_empty() {
            let unique_topics: Vec<_> = all_topics.iter()
                .cloned()
                .collect::<std::collections::HashSet<_>>()
                .into_iter()
                .take(10)
                .collect();
            abstraction_parts.push(format!("Topics: {}", unique_topics.join(", ")));
        }
        
        // Take first line of each summary
        let key_points: Vec<String> = summaries
            .iter()
            .take(5)
            .filter_map(|s| s.content.lines().next().map(|l| l.to_string()))
            .collect();
        
        if !key_points.is_empty() {
            abstraction_parts.push(format!("Key points:\n- {}", key_points.join("\n- ")));
        }
        
        Ok(abstraction_parts.join("\n"))
    }
    
    fn extract_topics(&self, content: &str) -> Vec<String> {
        // Simple keyword extraction (would use NLP in production)
        let stopwords = ["the", "a", "an", "is", "are", "was", "were", "to", "from", "in", "on", "at", "for", "and", "or"];
        
        let words: Vec<&str> = content
            .split_whitespace()
            .filter(|w| w.len() > 3)
            .filter(|w| !stopwords.contains(&w.to_lowercase().as_str()))
            .collect();
        
        // Count word frequencies
        let mut freq: HashMap<String, usize> = HashMap::new();
        for word in words {
            let normalized = word.to_lowercase().trim_matches(|c: char| !c.is_alphanumeric()).to_string();
            if normalized.len() > 3 {
                *freq.entry(normalized).or_insert(0) += 1;
            }
        }
        
        // Return top 5 by frequency
        let mut sorted: Vec<_> = freq.into_iter().collect();
        sorted.sort_by(|a, b| b.1.cmp(&a.1));
        
        sorted.into_iter().take(5).map(|(w, _)| w).collect()
    }
}

// ============================================================================
// Memory Store
// ============================================================================

/// Hierarchical memory store with compaction
pub struct HierarchicalMemory<S: Summarizer> {
    /// Configuration
    config: MemoryCompactionConfig,
    
    /// L0: Raw episodes
    l0_episodes: RwLock<VecDeque<Episode>>,
    
    /// L1: Summaries
    l1_summaries: RwLock<VecDeque<Summary>>,
    
    /// L2: Abstractions
    l2_abstractions: RwLock<VecDeque<Abstraction>>,
    
    /// Summarizer backend
    summarizer: Arc<S>,
    
    /// Compaction statistics
    stats: RwLock<CompactionStats>,
    
    /// ID counter
    next_id: std::sync::atomic::AtomicU64,
}

/// Compaction statistics
#[derive(Debug, Clone, Default)]
pub struct CompactionStats {
    /// Total episodes added
    pub total_episodes: usize,
    
    /// Total summaries created
    pub total_summaries: usize,
    
    /// Total abstractions created
    pub total_abstractions: usize,
    
    /// Episodes compacted (removed from L0)
    pub episodes_compacted: usize,
    
    /// Summaries compacted (removed from L1)
    pub summaries_compacted: usize,
    
    /// Last compaction time
    pub last_compaction: Option<f64>,
    
    /// Total token savings (estimated)
    pub token_savings: usize,
}

impl<S: Summarizer> HierarchicalMemory<S> {
    /// Create a new hierarchical memory store
    pub fn new(config: MemoryCompactionConfig, summarizer: Arc<S>) -> Self {
        Self {
            config,
            l0_episodes: RwLock::new(VecDeque::new()),
            l1_summaries: RwLock::new(VecDeque::new()),
            l2_abstractions: RwLock::new(VecDeque::new()),
            summarizer,
            stats: RwLock::new(CompactionStats::default()),
            next_id: std::sync::atomic::AtomicU64::new(1),
        }
    }
    
    /// Generate next ID
    fn next_id(&self) -> String {
        let id = self.next_id.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        format!("mem_{}", id)
    }
    
    /// Add an episode to L0
    pub fn add_episode(&self, content: String, episode_type: EpisodeType) -> String {
        let id = self.next_id();
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs_f64();
        
        let token_count = content.len() / 4; // Rough estimate
        
        let episode = Episode {
            id: id.clone(),
            timestamp,
            content,
            episode_type,
            metadata: HashMap::new(),
            embedding: None,
            token_count,
        };
        
        {
            let mut l0 = self.l0_episodes.write().unwrap();
            l0.push_back(episode);
        }
        
        {
            let mut stats = self.stats.write().unwrap();
            stats.total_episodes += 1;
        }
        
        id
    }
    
    /// Add episode with embedding
    pub fn add_episode_with_embedding(
        &self,
        content: String,
        episode_type: EpisodeType,
        embedding: Vec<f32>,
    ) -> String {
        let id = self.next_id();
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs_f64();
        
        let token_count = content.len() / 4;
        
        let episode = Episode {
            id: id.clone(),
            timestamp,
            content,
            episode_type,
            metadata: HashMap::new(),
            embedding: Some(embedding),
            token_count,
        };
        
        {
            let mut l0 = self.l0_episodes.write().unwrap();
            l0.push_back(episode);
        }
        
        {
            let mut stats = self.stats.write().unwrap();
            stats.total_episodes += 1;
        }
        
        id
    }
    
    /// Check if compaction is needed and run if so
    pub fn maybe_compact(&self) -> Result<bool, CompactionError> {
        let needs_l0 = {
            let l0 = self.l0_episodes.read().unwrap();
            l0.len() >= self.config.l0_max_episodes
        };
        
        let needs_l1 = {
            let l1 = self.l1_summaries.read().unwrap();
            l1.len() >= self.config.l1_max_summaries
        };
        
        if needs_l0 || needs_l1 {
            self.run_compaction()?;
            return Ok(true);
        }
        
        Ok(false)
    }
    
    /// Run compaction cycle
    pub fn run_compaction(&self) -> Result<(), CompactionError> {
        // L0 → L1 compaction
        self.compact_l0_to_l1()?;
        
        // L1 → L2 compaction
        self.compact_l1_to_l2()?;
        
        // Update stats
        {
            let mut stats = self.stats.write().unwrap();
            stats.last_compaction = Some(
                SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_secs_f64()
            );
        }
        
        Ok(())
    }
    
    /// Compact L0 episodes to L1 summaries
    fn compact_l0_to_l1(&self) -> Result<(), CompactionError> {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs_f64();
        
        let age_threshold = now - self.config.l0_age_threshold_secs as f64;
        
        // Collect episodes to compact
        let to_compact: Vec<Episode> = {
            let l0 = self.l0_episodes.read().unwrap();
            l0.iter()
                .filter(|e| e.timestamp < age_threshold)
                .cloned()
                .collect()
        };
        
        if to_compact.is_empty() {
            return Ok(());
        }
        
        // Group episodes by similarity or time
        let groups = self.group_episodes(&to_compact);
        
        // Summarize each group
        for group in groups {
            if group.is_empty() {
                continue;
            }
            
            let content = self.summarizer.summarize_episodes(&group)?;
            let topics = self.summarizer.extract_topics(&content);
            
            let first_ts = group.iter().map(|e| e.timestamp).fold(f64::MAX, f64::min);
            let last_ts = group.iter().map(|e| e.timestamp).fold(f64::MIN, f64::max);
            
            let episode_ids: Vec<String> = group.iter().map(|e| e.id.clone()).collect();
            let original_tokens: usize = group.iter().map(|e| e.token_count).sum();
            let summary_tokens = content.len() / 4;
            
            let summary = Summary {
                id: self.next_id(),
                content,
                source_episode_ids: episode_ids,
                time_range: (first_ts, last_ts),
                embedding: None, // Would be generated if reembed_summaries is true
                token_count: summary_tokens,
                created_at: now,
                topics,
            };
            
            // Add summary to L1
            {
                let mut l1 = self.l1_summaries.write().unwrap();
                l1.push_back(summary);
            }
            
            // Update stats
            {
                let mut stats = self.stats.write().unwrap();
                stats.total_summaries += 1;
                stats.episodes_compacted += group.len();
                stats.token_savings += original_tokens.saturating_sub(summary_tokens);
            }
        }
        
        // Remove compacted episodes from L0
        {
            let mut l0 = self.l0_episodes.write().unwrap();
            l0.retain(|e| e.timestamp >= age_threshold);
        }
        
        Ok(())
    }
    
    /// Compact L1 summaries to L2 abstractions
    fn compact_l1_to_l2(&self) -> Result<(), CompactionError> {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs_f64();
        
        let age_threshold = now - self.config.l1_age_threshold_secs as f64;
        
        // Collect summaries to compact
        let to_compact: Vec<Summary> = {
            let l1 = self.l1_summaries.read().unwrap();
            l1.iter()
                .filter(|s| s.created_at < age_threshold)
                .cloned()
                .collect()
        };
        
        if to_compact.len() < self.config.group_size {
            return Ok(());
        }
        
        // Group summaries
        let groups = self.group_summaries(&to_compact);
        
        for group in groups {
            if group.is_empty() {
                continue;
            }
            
            let content = self.summarizer.abstract_summaries(&group)?;
            
            let first_ts = group.iter().map(|s| s.time_range.0).fold(f64::MAX, f64::min);
            let last_ts = group.iter().map(|s| s.time_range.1).fold(f64::MIN, f64::max);
            
            let summary_ids: Vec<String> = group.iter().map(|s| s.id.clone()).collect();
            let original_tokens: usize = group.iter().map(|s| s.token_count).sum();
            let abstraction_tokens = content.len() / 4;
            
            // Extract insights from topics
            let insights: Vec<String> = group
                .iter()
                .flat_map(|s| s.topics.clone())
                .collect::<std::collections::HashSet<_>>()
                .into_iter()
                .take(5)
                .collect();
            
            let abstraction = Abstraction {
                id: self.next_id(),
                content,
                source_summary_ids: summary_ids,
                time_range: (first_ts, last_ts),
                embedding: None,
                token_count: abstraction_tokens,
                created_at: now,
                insights,
            };
            
            // Add abstraction to L2
            {
                let mut l2 = self.l2_abstractions.write().unwrap();
                l2.push_back(abstraction);
            }
            
            // Update stats
            {
                let mut stats = self.stats.write().unwrap();
                stats.total_abstractions += 1;
                stats.summaries_compacted += group.len();
                stats.token_savings += original_tokens.saturating_sub(abstraction_tokens);
            }
        }
        
        // Remove compacted summaries from L1
        {
            let mut l1 = self.l1_summaries.write().unwrap();
            l1.retain(|s| s.created_at >= age_threshold);
        }
        
        Ok(())
    }
    
    /// Group episodes by time windows (simplified)
    fn group_episodes(&self, episodes: &[Episode]) -> Vec<Vec<Episode>> {
        // Simple grouping by fixed size
        episodes
            .chunks(self.config.group_size)
            .map(|chunk| chunk.to_vec())
            .collect()
    }
    
    /// Group summaries by time windows
    fn group_summaries(&self, summaries: &[Summary]) -> Vec<Vec<Summary>> {
        summaries
            .chunks(self.config.group_size)
            .map(|chunk| chunk.to_vec())
            .collect()
    }
    
    /// Get total token count across all tiers
    pub fn total_tokens(&self) -> usize {
        let l0: usize = self.l0_episodes.read().unwrap().iter().map(|e| e.token_count).sum();
        let l1: usize = self.l1_summaries.read().unwrap().iter().map(|s| s.token_count).sum();
        let l2: usize = self.l2_abstractions.read().unwrap().iter().map(|a| a.token_count).sum();
        
        l0 + l1 + l2
    }
    
    /// Get memory for context assembly (most recent first)
    pub fn get_context(&self, max_tokens: usize) -> Vec<MemoryEntry> {
        let mut entries = Vec::new();
        let mut tokens_used = 0;
        
        // Start with L0 (most recent)
        let l0 = self.l0_episodes.read().unwrap();
        for episode in l0.iter().rev() {
            if tokens_used + episode.token_count > max_tokens {
                break;
            }
            entries.push(MemoryEntry::Episode(episode.clone()));
            tokens_used += episode.token_count;
        }
        
        // Add L1 summaries if space
        let l1 = self.l1_summaries.read().unwrap();
        for summary in l1.iter().rev() {
            if tokens_used + summary.token_count > max_tokens {
                break;
            }
            entries.push(MemoryEntry::Summary(summary.clone()));
            tokens_used += summary.token_count;
        }
        
        // Add L2 abstractions if space
        let l2 = self.l2_abstractions.read().unwrap();
        for abstraction in l2.iter().rev() {
            if tokens_used + abstraction.token_count > max_tokens {
                break;
            }
            entries.push(MemoryEntry::Abstraction(abstraction.clone()));
            tokens_used += abstraction.token_count;
        }
        
        entries
    }
    
    /// Get statistics
    pub fn stats(&self) -> CompactionStats {
        self.stats.read().unwrap().clone()
    }
    
    /// Get tier counts
    pub fn tier_counts(&self) -> (usize, usize, usize) {
        let l0 = self.l0_episodes.read().unwrap().len();
        let l1 = self.l1_summaries.read().unwrap().len();
        let l2 = self.l2_abstractions.read().unwrap().len();
        (l0, l1, l2)
    }
}

/// Entry from hierarchical memory
#[derive(Debug, Clone)]
pub enum MemoryEntry {
    Episode(Episode),
    Summary(Summary),
    Abstraction(Abstraction),
}

impl MemoryEntry {
    /// Get content
    pub fn content(&self) -> &str {
        match self {
            Self::Episode(e) => &e.content,
            Self::Summary(s) => &s.content,
            Self::Abstraction(a) => &a.content,
        }
    }
    
    /// Get token count
    pub fn token_count(&self) -> usize {
        match self {
            Self::Episode(e) => e.token_count,
            Self::Summary(s) => s.token_count,
            Self::Abstraction(a) => a.token_count,
        }
    }
    
    /// Get tier level
    pub fn tier(&self) -> usize {
        match self {
            Self::Episode(_) => 0,
            Self::Summary(_) => 1,
            Self::Abstraction(_) => 2,
        }
    }
}

// ============================================================================
// Errors
// ============================================================================

/// Compaction error
#[derive(Debug, Clone)]
pub enum CompactionError {
    /// Summarization failed
    SummarizationFailed(String),
    /// Embedding failed
    EmbeddingFailed(String),
    /// Storage error
    StorageError(String),
}

impl std::fmt::Display for CompactionError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::SummarizationFailed(msg) => write!(f, "Summarization failed: {}", msg),
            Self::EmbeddingFailed(msg) => write!(f, "Embedding failed: {}", msg),
            Self::StorageError(msg) => write!(f, "Storage error: {}", msg),
        }
    }
}

impl std::error::Error for CompactionError {}

// ============================================================================
// Convenience Functions
// ============================================================================

/// Create a hierarchical memory with extractive summarizer
pub fn create_hierarchical_memory() -> HierarchicalMemory<ExtractiveSummarizer> {
    HierarchicalMemory::new(
        MemoryCompactionConfig::default(),
        Arc::new(ExtractiveSummarizer::default()),
    )
}

/// Create with aggressive compaction for testing
pub fn create_test_memory() -> HierarchicalMemory<ExtractiveSummarizer> {
    HierarchicalMemory::new(
        MemoryCompactionConfig::aggressive(),
        Arc::new(ExtractiveSummarizer::default()),
    )
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_add_episode() {
        let memory = create_test_memory();
        
        let id = memory.add_episode(
            "User asked about weather".to_string(),
            EpisodeType::UserMessage,
        );
        
        assert!(id.starts_with("mem_"));
        
        let (l0, l1, l2) = memory.tier_counts();
        assert_eq!(l0, 1);
        assert_eq!(l1, 0);
        assert_eq!(l2, 0);
    }
    
    #[test]
    fn test_extractive_summarizer() {
        let summarizer = ExtractiveSummarizer::default();
        
        let episodes = vec![
            Episode {
                id: "1".to_string(),
                timestamp: 0.0,
                content: "User asked about the weather forecast.".to_string(),
                episode_type: EpisodeType::UserMessage,
                metadata: HashMap::new(),
                embedding: None,
                token_count: 10,
            },
            Episode {
                id: "2".to_string(),
                timestamp: 1.0,
                content: "Assistant provided weather information for NYC.".to_string(),
                episode_type: EpisodeType::AssistantResponse,
                metadata: HashMap::new(),
                embedding: None,
                token_count: 12,
            },
        ];
        
        let summary = summarizer.summarize_episodes(&episodes).unwrap();
        
        assert!(!summary.is_empty());
        assert!(summary.contains("episodes") || summary.contains("User") || summary.contains("Responses"));
    }
    
    #[test]
    fn test_topic_extraction() {
        let summarizer = ExtractiveSummarizer::default();
        
        let content = "The weather forecast shows sunny conditions with temperatures around 75 degrees. Tomorrow expects rain and thunderstorms across the region.";
        
        let topics = summarizer.extract_topics(content);
        
        assert!(!topics.is_empty());
        // Should extract meaningful words like "weather", "forecast", "temperatures", etc.
    }
    
    #[test]
    fn test_memory_context_retrieval() {
        let memory = create_test_memory();
        
        // Add some episodes
        for i in 0..5 {
            memory.add_episode(
                format!("Episode {} content here with some text.", i),
                EpisodeType::UserMessage,
            );
        }
        
        let context = memory.get_context(1000);
        
        assert!(!context.is_empty());
        
        // All should be L0 episodes
        for entry in &context {
            assert_eq!(entry.tier(), 0);
        }
    }
    
    #[test]
    fn test_token_tracking() {
        let memory = create_test_memory();
        
        memory.add_episode(
            "Short message".to_string(),
            EpisodeType::UserMessage,
        );
        
        memory.add_episode(
            "A much longer message with more content that should have more tokens estimated".to_string(),
            EpisodeType::AssistantResponse,
        );
        
        let total = memory.total_tokens();
        assert!(total > 0);
    }
    
    #[test]
    fn test_stats_tracking() {
        let memory = create_test_memory();
        
        for _ in 0..10 {
            memory.add_episode("Test episode".to_string(), EpisodeType::UserMessage);
        }
        
        let stats = memory.stats();
        assert_eq!(stats.total_episodes, 10);
    }
}