ccswarm 0.4.2

//! Agent memory system integrated with session management
//! Provides working, episodic, semantic, and procedural memory for agents

//use anyhow::Result;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};

const WORKING_MEMORY_CAPACITY: usize = 7; // Miller's 7±2 rule

/// Integrated memory system for agent sessions
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionMemory {
    pub session_id: String,
    pub agent_id: String,
    pub working_memory: WorkingMemory,
    pub episodic_memory: EpisodicMemory,
    pub semantic_memory: SemanticMemory,
    pub procedural_memory: ProceduralMemory,
    pub memory_stats: MemoryStats,
    pub last_consolidation: DateTime<Utc>,
}

/// Working memory for immediate task processing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkingMemory {
    pub current_items: VecDeque<WorkingMemoryItem>,
    pub capacity: usize,
    pub active_task_context: Option<TaskContext>,
    pub attention_focus: Vec<String>,
    pub cognitive_load: f32, // 0.0-1.0
}

/// Item in working memory
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkingMemoryItem {
    pub id: String,
    pub content: String,
    pub item_type: WorkingMemoryType,
    pub priority: f32,
    pub decay_rate: f32,
    pub created_at: DateTime<Utc>,
    pub last_accessed: DateTime<Utc>,
}

/// Types of working memory items
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum WorkingMemoryType {
    TaskInstructions,
    IntermediateResult,
    ErrorMessage,
    UserFeedback,
    ContextualClue,
    PlanningStep,
}

/// Current task context in working memory
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskContext {
    pub task_id: String,
    pub task_description: String,
    pub current_step: String,
    pub progress: f32,
    pub obstacles: Vec<String>,
    pub solutions_attempted: Vec<String>,
    pub time_spent: std::time::Duration,
    pub expected_completion: Option<DateTime<Utc>>,
}

/// Episodic memory for experiences and events
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EpisodicMemory {
    pub episodes: VecDeque<Episode>,
    pub max_episodes: usize,
    pub recent_experiences: Vec<String>,
    pub emotional_markers: HashMap<String, EmotionalMarker>,
}

/// Individual episode/experience
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Episode {
    pub id: String,
    pub timestamp: DateTime<Utc>,
    pub event_type: EpisodeType,
    pub description: String,
    pub context: HashMap<String, String>,
    pub outcome: EpisodeOutcome,
    pub emotional_valence: f32, // -1.0 to 1.0
    pub learning_value: f32,    // 0.0 to 1.0
    pub related_episodes: Vec<String>,
}

/// Types of episodes
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum EpisodeType {
    TaskCompletion,
    ProblemSolving,
    Collaboration,
    Learning,
    ErrorRecovery,
    Discovery,
    Routine,
}

/// Outcome of an episode
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum EpisodeOutcome {
    Success {
        metrics: HashMap<String, f32>,
    },
    Failure {
        reason: String,
        recovery_actions: Vec<String>,
    },
    Partial {
        progress: f32,
        next_steps: Vec<String>,
    },
    Cancelled {
        reason: String,
    },
}

/// Emotional marker for memory consolidation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmotionalMarker {
    pub event_id: String,
    pub emotion_type: EmotionType,
    pub intensity: f32,
    pub influence_on_learning: f32,
}

/// Types of emotions affecting memory
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum EmotionType {
    Satisfaction,
    Frustration,
    Curiosity,
    Confidence,
    Uncertainty,
    Excitement,
}

/// Semantic memory for concepts and knowledge
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SemanticMemory {
    pub concepts: HashMap<String, Concept>,
    pub relationships: Vec<ConceptRelationship>,
    pub knowledge_domains: HashMap<String, KnowledgeDomain>,
    pub fact_base: Vec<Fact>,
}

/// Individual concept in semantic memory
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Concept {
    pub name: String,
    pub definition: String,
    pub properties: HashMap<String, String>,
    pub examples: Vec<String>,
    pub confidence: f32,
    pub source_episodes: Vec<String>,
    pub last_updated: DateTime<Utc>,
    pub usage_frequency: u32,
}

/// Relationship between concepts
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConceptRelationship {
    pub from_concept: String,
    pub to_concept: String,
    pub relationship_type: RelationshipType,
    pub strength: f32,
    pub evidence: Vec<String>,
}

/// Types of concept relationships
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RelationshipType {
    IsA,           // TypeOf
    PartOf,        // Component
    Uses,          // Functional
    CausedBy,      // Causal
    SimilarTo,     // Analogical
    ConflictsWith, // Contradictory
}

/// Knowledge domain organization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KnowledgeDomain {
    pub name: String,
    pub core_concepts: Vec<String>,
    pub proficiency_level: ProficiencyLevel,
    pub learning_progress: f32,
    pub recent_activities: Vec<String>,
}

/// Proficiency levels in knowledge domains
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ProficiencyLevel {
    Novice,
    Beginner,
    Intermediate,
    Advanced,
    Expert,
}

/// Individual fact in knowledge base
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Fact {
    pub id: String,
    pub statement: String,
    pub confidence: f32,
    pub source: String,
    pub verified: bool,
    pub created_at: DateTime<Utc>,
    pub related_concepts: Vec<String>,
}

/// Procedural memory for skills and processes
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProceduralMemory {
    pub procedures: HashMap<String, Procedure>,
    pub skill_patterns: Vec<SkillPattern>,
    pub automation_levels: HashMap<String, f32>,
    pub execution_templates: Vec<ExecutionTemplate>,
}

/// Individual procedure/skill
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Procedure {
    pub name: String,
    pub steps: Vec<ProcedureStep>,
    pub prerequisites: Vec<String>,
    pub success_rate: f32,
    pub average_execution_time: std::time::Duration,
    pub complexity_level: ComplexityLevel,
    pub last_executed: Option<DateTime<Utc>>,
    pub optimization_history: Vec<Optimization>,
}

/// Step in a procedure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProcedureStep {
    pub step_number: u32,
    pub description: String,
    pub action_type: ActionType,
    pub required_inputs: Vec<String>,
    pub expected_outputs: Vec<String>,
    pub error_handling: Vec<ErrorHandler>,
    pub execution_time: Option<std::time::Duration>,
}

/// Types of actions in procedures
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ActionType {
    Cognitive,    // Thinking/reasoning
    Physical,     // File operations, commands
    Interactive,  // User/system interaction
    Verificative, // Checking/validation
    Creative,     // Generation/synthesis
}

/// Error handling strategy
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ErrorHandler {
    pub error_pattern: String,
    pub recovery_action: String,
    pub fallback_procedure: Option<String>,
    pub escalation_threshold: u32,
}

/// Complexity levels for procedures
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ComplexityLevel {
    Simple,
    Moderate,
    Complex,
    Expert,
}

/// Skill pattern recognition
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SkillPattern {
    pub pattern_name: String,
    pub trigger_conditions: Vec<String>,
    pub execution_sequence: Vec<String>,
    pub success_indicators: Vec<String>,
    pub adaptation_rules: Vec<String>,
}

/// Execution template for common patterns
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionTemplate {
    pub template_name: String,
    pub situation_type: String,
    pub steps: Vec<TemplateStep>,
    pub customization_points: Vec<String>,
}

/// Step in execution template
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TemplateStep {
    pub description: String,
    pub variables: Vec<String>,
    pub optional: bool,
}

/// Optimization record for procedures
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Optimization {
    pub timestamp: DateTime<Utc>,
    pub change_description: String,
    pub performance_impact: f32,
    pub success_rate_change: f32,
}

/// Memory system statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryStats {
    pub working_memory_utilization: f32,
    pub episodic_memory_size: usize,
    pub semantic_concepts_count: usize,
    pub procedural_skills_count: usize,
    pub consolidation_events: u32,
    pub memory_efficiency: f32,
    pub last_updated: DateTime<Utc>,
}

impl SessionMemory {
    /// Create new session memory
    pub fn new(session_id: String, agent_id: String) -> Self {
        Self {
            session_id,
            agent_id,
            working_memory: WorkingMemory::new(),
            episodic_memory: EpisodicMemory::new(),
            semantic_memory: SemanticMemory::new(),
            procedural_memory: ProceduralMemory::new(),
            memory_stats: MemoryStats::new(),
            last_consolidation: Utc::now(),
        }
    }

    /// Add item to working memory
    pub fn add_to_working_memory(
        &mut self,
        content: String,
        item_type: WorkingMemoryType,
        priority: f32,
    ) {
        let item = WorkingMemoryItem {
            id: uuid::Uuid::new_v4().to_string(),
            content,
            item_type,
            priority,
            decay_rate: 0.1,
            created_at: Utc::now(),
            last_accessed: Utc::now(),
        };

        self.working_memory.add_item(item);
        self.update_cognitive_load();
    }

    /// Set current task context
    pub fn set_task_context(&mut self, task_id: String, description: String) {
        self.working_memory.active_task_context = Some(TaskContext {
            task_id,
            task_description: description,
            current_step: "initialization".to_string(),
            progress: 0.0,
            obstacles: Vec::new(),
            solutions_attempted: Vec::new(),
            time_spent: std::time::Duration::from_secs(0),
            expected_completion: None,
        });
    }

    /// Add episode to episodic memory
    pub fn add_episode(
        &mut self,
        event_type: EpisodeType,
        description: String,
        context: HashMap<String, String>,
        outcome: EpisodeOutcome,
    ) {
        let episode = Episode {
            id: uuid::Uuid::new_v4().to_string(),
            timestamp: Utc::now(),
            event_type,
            description,
            context,
            outcome,
            emotional_valence: 0.0, // Will be set based on outcome
            learning_value: 0.5,    // Default learning value
            related_episodes: Vec::new(),
        };

        self.episodic_memory.add_episode(episode);
    }

    /// Add concept to semantic memory
    pub fn add_concept(
        &mut self,
        name: String,
        definition: String,
        properties: HashMap<String, String>,
    ) {
        let concept = Concept {
            name: name.clone(),
            definition,
            properties,
            examples: Vec::new(),
            confidence: 0.7,
            source_episodes: Vec::new(),
            last_updated: Utc::now(),
            usage_frequency: 1,
        };

        self.semantic_memory.add_concept(concept);
    }

    /// Add procedure to procedural memory
    pub fn add_procedure(&mut self, name: String, steps: Vec<ProcedureStep>) {
        let procedure = Procedure {
            name: name.clone(),
            steps,
            prerequisites: Vec::new(),
            success_rate: 0.8,
            average_execution_time: std::time::Duration::from_secs(60),
            complexity_level: ComplexityLevel::Moderate,
            last_executed: None,
            optimization_history: Vec::new(),
        };

        self.procedural_memory.add_procedure(procedure);
    }

    /// Consolidate memories (transfer from working to long-term)
    pub fn consolidate_memories(&mut self) {
        // Process working memory items for consolidation
        let items_to_consolidate: Vec<_> = self
            .working_memory
            .current_items
            .iter()
            .filter(|item| self.should_consolidate_item(item))
            .cloned()
            .collect();

        for item in items_to_consolidate {
            self.consolidate_working_memory_item(&item);
        }

        // Clean up old working memory items
        self.working_memory.cleanup_expired_items();

        // Update consolidation timestamp
        self.last_consolidation = Utc::now();
        self.memory_stats.consolidation_events += 1;
    }

    /// Check if working memory item should be consolidated
    fn should_consolidate_item(&self, item: &WorkingMemoryItem) -> bool {
        let age = Utc::now() - item.created_at;
        let significance = item.priority;

        // Consolidate if item is old enough and significant
        age.num_minutes() > 30 && significance > 0.7
    }

    /// Consolidate a working memory item into long-term memory
    fn consolidate_working_memory_item(&mut self, item: &WorkingMemoryItem) {
        match &item.item_type {
            WorkingMemoryType::TaskInstructions => {
                // Add to procedural memory as a pattern
                self.procedural_memory.skill_patterns.push(SkillPattern {
                    pattern_name: format!("Task pattern: {}", item.id),
                    trigger_conditions: vec![item.content.clone()],
                    execution_sequence: Vec::new(),
                    success_indicators: Vec::new(),
                    adaptation_rules: Vec::new(),
                });
            }
            WorkingMemoryType::IntermediateResult => {
                // Add as semantic knowledge
                self.semantic_memory.fact_base.push(Fact {
                    id: item.id.clone(),
                    statement: item.content.clone(),
                    confidence: item.priority,
                    source: "working_memory".to_string(),
                    verified: false,
                    created_at: item.created_at,
                    related_concepts: Vec::new(),
                });
            }
            WorkingMemoryType::ErrorMessage => {
                // Add to episodic memory as learning experience
                self.episodic_memory.add_episode(Episode {
                    id: uuid::Uuid::new_v4().to_string(),
                    timestamp: item.created_at,
                    event_type: EpisodeType::ErrorRecovery,
                    description: item.content.clone(),
                    context: HashMap::new(),
                    outcome: EpisodeOutcome::Partial {
                        progress: 0.5,
                        next_steps: vec!["Learn from error".to_string()],
                    },
                    emotional_valence: -0.3, // Negative but learning
                    learning_value: 0.8,     // High learning value
                    related_episodes: Vec::new(),
                });
            }
            _ => {
                // General consolidation to episodic memory
                self.episodic_memory.add_episode(Episode {
                    id: uuid::Uuid::new_v4().to_string(),
                    timestamp: item.created_at,
                    event_type: EpisodeType::Routine,
                    description: item.content.clone(),
                    context: HashMap::new(),
                    outcome: EpisodeOutcome::Success {
                        metrics: HashMap::new(),
                    },
                    emotional_valence: 0.1,
                    learning_value: item.priority,
                    related_episodes: Vec::new(),
                });
            }
        }
    }

    /// Update cognitive load based on working memory utilization
    fn update_cognitive_load(&mut self) {
        let utilization =
            self.working_memory.current_items.len() as f32 / WORKING_MEMORY_CAPACITY as f32;
        self.working_memory.cognitive_load = utilization.min(1.0);

        // Update memory stats
        self.memory_stats.working_memory_utilization = utilization;
        self.memory_stats.last_updated = Utc::now();
    }

    /// Retrieve relevant memories for current context
    pub fn retrieve_relevant_memories(&self, query: &str) -> RetrievalResult {
        let mut result = RetrievalResult::new();

        // Search working memory
        for item in &self.working_memory.current_items {
            if item.content.to_lowercase().contains(&query.to_lowercase()) {
                result.working_memory_items.push(item.clone());
            }
        }

        // Search episodic memory
        for episode in &self.episodic_memory.episodes {
            if episode
                .description
                .to_lowercase()
                .contains(&query.to_lowercase())
            {
                result.relevant_episodes.push(episode.clone());
            }
        }

        // Search semantic memory
        for concept in self.semantic_memory.concepts.values() {
            if concept.name.to_lowercase().contains(&query.to_lowercase())
                || concept
                    .definition
                    .to_lowercase()
                    .contains(&query.to_lowercase())
            {
                result.relevant_concepts.push(concept.clone());
            }
        }

        // Search procedural memory
        for procedure in self.procedural_memory.procedures.values() {
            if procedure
                .name
                .to_lowercase()
                .contains(&query.to_lowercase())
            {
                result.relevant_procedures.push(procedure.clone());
            }
        }

        result
    }

    /// Get memory summary
    pub fn get_summary(&self) -> MemorySummary {
        self.generate_memory_summary()
    }

    /// Generate memory summary for reporting
    pub fn generate_memory_summary(&self) -> MemorySummary {
        MemorySummary {
            session_id: self.session_id.clone(),
            agent_id: self.agent_id.clone(),
            working_memory_load: self.working_memory.cognitive_load,
            episodic_memory_size: self.episodic_memory.episodes.len(),
            semantic_concepts: self.semantic_memory.concepts.len(),
            procedural_skills: self.procedural_memory.procedures.len(),
            recent_episodes: self
                .episodic_memory
                .episodes
                .iter()
                .rev()
                .take(5)
                .map(|e| e.description.clone())
                .collect(),
            last_consolidation: self.last_consolidation,
            memory_efficiency: self.memory_stats.memory_efficiency,
        }
    }
}

/// Result of memory retrieval
#[derive(Debug, Clone)]
pub struct RetrievalResult {
    pub working_memory_items: Vec<WorkingMemoryItem>,
    pub relevant_episodes: Vec<Episode>,
    pub relevant_concepts: Vec<Concept>,
    pub relevant_procedures: Vec<Procedure>,
}

impl RetrievalResult {
    fn new() -> Self {
        Self {
            working_memory_items: Vec::new(),
            relevant_episodes: Vec::new(),
            relevant_concepts: Vec::new(),
            relevant_procedures: Vec::new(),
        }
    }
}

/// Memory summary for reporting
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemorySummary {
    pub session_id: String,
    pub agent_id: String,
    pub working_memory_load: f32,
    pub episodic_memory_size: usize,
    pub semantic_concepts: usize,
    pub procedural_skills: usize,
    pub recent_episodes: Vec<String>,
    pub last_consolidation: DateTime<Utc>,
    pub memory_efficiency: f32,
}

// Implementation for memory subsystems
impl WorkingMemory {
    fn new() -> Self {
        Self {
            current_items: VecDeque::with_capacity(WORKING_MEMORY_CAPACITY),
            capacity: WORKING_MEMORY_CAPACITY,
            active_task_context: None,
            attention_focus: Vec::new(),
            cognitive_load: 0.0,
        }
    }

    fn add_item(&mut self, item: WorkingMemoryItem) {
        // Remove oldest item if at capacity
        if self.current_items.len() >= self.capacity {
            self.current_items.pop_front();
        }
        self.current_items.push_back(item);
    }

    fn cleanup_expired_items(&mut self) {
        let now = Utc::now();
        self.current_items.retain(|item| {
            let age = now - item.created_at;
            let decay_threshold = item.decay_rate * age.num_minutes() as f32;
            decay_threshold < 1.0 // Keep if not fully decayed
        });
    }
}

impl EpisodicMemory {
    fn new() -> Self {
        Self {
            episodes: VecDeque::with_capacity(1000),
            max_episodes: 1000,
            recent_experiences: Vec::new(),
            emotional_markers: HashMap::new(),
        }
    }

    fn add_episode(&mut self, episode: Episode) {
        // Remove oldest episode if at capacity
        if self.episodes.len() >= self.max_episodes {
            if let Some(old_episode) = self.episodes.pop_front() {
                self.emotional_markers.remove(&old_episode.id);
            }
        }

        // Add emotional marker if significant
        if episode.emotional_valence.abs() > 0.5 {
            self.emotional_markers.insert(
                episode.id.clone(),
                EmotionalMarker {
                    event_id: episode.id.clone(),
                    emotion_type: EmotionType::Satisfaction, // Simplified
                    intensity: episode.emotional_valence.abs(),
                    influence_on_learning: episode.learning_value,
                },
            );
        }

        self.episodes.push_back(episode);
    }
}

impl SemanticMemory {
    fn new() -> Self {
        Self {
            concepts: HashMap::new(),
            relationships: Vec::new(),
            knowledge_domains: HashMap::new(),
            fact_base: Vec::new(),
        }
    }

    fn add_concept(&mut self, concept: Concept) {
        self.concepts.insert(concept.name.clone(), concept);
    }
}

impl ProceduralMemory {
    fn new() -> Self {
        Self {
            procedures: HashMap::new(),
            skill_patterns: Vec::new(),
            automation_levels: HashMap::new(),
            execution_templates: Vec::new(),
        }
    }

    fn add_procedure(&mut self, procedure: Procedure) {
        self.procedures.insert(procedure.name.clone(), procedure);
    }
}

impl MemoryStats {
    fn new() -> Self {
        Self {
            working_memory_utilization: 0.0,
            episodic_memory_size: 0,
            semantic_concepts_count: 0,
            procedural_skills_count: 0,
            consolidation_events: 0,
            memory_efficiency: 1.0,
            last_updated: Utc::now(),
        }
    }
}