metabolism.rs

use crate::ports::{BufferPort, MemoryPort, LlmPort, EmbeddingPort, EngramPort, HolographicPort, SanitizerPort};
use crate::error::Result;
use crate::{MemoryNode, NodeType};
use std::sync::Arc;
use chrono::Utc;
use uuid::Uuid;
use crate::logic::system3::{NeuroSymbolicEngine, ReasoningGraph, LogicNode, LogicNodeType};

/// Metabolism: The process of digesting short-term buffer into long-term memory.
///
/// This struct implements the core logic for the "Metabolism (Buffer → Long-term)" feature.
/// It consolidates interactions from a short-term buffer, summarizes them using an LLM,
/// and stores the result as a `MemoryNode` in long-term storage.
#[derive(Clone)]
pub struct Metabolism {
    buffer: Arc<dyn BufferPort>,
    memory: Arc<dyn MemoryPort>,
    llm: Arc<dyn LlmPort>,
    embedder: Arc<dyn EmbeddingPort>,
    engram: Option<Arc<dyn EngramPort>>,
    holographic: Arc<dyn HolographicPort>,
    sanitizer: Arc<dyn SanitizerPort>,
    threshold: usize,
    neuro_symbolic: Option<Arc<NeuroSymbolicEngine>>,
}

impl Metabolism {
    pub fn new(
        buffer: Arc<dyn BufferPort>,
        memory: Arc<dyn MemoryPort>,
        llm: Arc<dyn LlmPort>,
        embedder: Arc<dyn EmbeddingPort>,
        holographic: Arc<dyn HolographicPort>,
        sanitizer: Arc<dyn SanitizerPort>,
    ) -> Self {
        Self {
            buffer,
            memory,
            llm,
            embedder,
            engram: None,
            holographic,
            sanitizer,
            threshold: 10, // Default threshold
            neuro_symbolic: None,
        }
    }

    /// Enable System 3 validation using the Neuro-Symbolic Engine
    pub fn with_system3(mut self, engine: Arc<NeuroSymbolicEngine>) -> Self {
        self.neuro_symbolic = Some(engine);
        self
    }

    /// Push a new interaction to the short-term buffer.
    pub async fn push_interaction(&self, interaction: crate::Interaction) -> Result<()> {
        self.buffer.push(interaction).await
    }

    /// Peek into the buffer (for UI history).
    pub async fn peek_buffer(&self, limit: usize) -> Result<Vec<crate::Interaction>> {
        self.buffer.peek(limit).await
    }

    /// Digest the buffer if it exceeds the threshold.
    pub async fn digest(&self) -> Result<usize> {
        let count = self.buffer.len().await?;
        if count < self.threshold {
            return Ok(0);
        }

        // 1. Retrieve the Last Neurogram (if any) to close the loop
        let previous_context = if let Ok(Some(last_node)) = self.memory.fetch_last_node().await {
            if let Some(neurogram) = &last_node.holographic_data {
                // Decompress the past
                 match self.holographic.decompress_context(neurogram).await {
                    Ok(ctx) => Some(ctx),
                    Err(_) => None,
                }
            } else {
                None
            }
        } else {
            None
        };

        // 2. Pop batch from buffer
        let interactions = self.buffer.pop_batch(self.threshold).await?;
        if interactions.is_empty() {
            return Ok(0);
        }

        // 3. Sanitize and combine content
        let mut combined_text = String::new();

        // Inject previous context first (The Loop)
        if let Some(prev) = previous_context {
            combined_text.push_str(&format!("--- Previous Context ---\n{}\n--- End Previous ---\n", prev));
        }

        for interaction in &interactions {
            let sanitized_user = self.sanitizer.sanitize(&interaction.user_input)?;
            let sanitized_ai = self.sanitizer.sanitize(&interaction.ai_response)?;
            combined_text.push_str(&format!("User: {}\nAI: {}\n", sanitized_user, sanitized_ai));
        }

        // 4. Summarize via LLM
        let summary = self.llm.summarize(&combined_text).await?;

        // System 3 Validation: Ensure the summary reasoning does not contain hallucinations
        if let Some(engine) = &self.neuro_symbolic {
            // Very simplified: Generate a mock Reasoning Graph for the summary logic
            // In a real scenario, the LLM should output the ReasoningGraph directly.
            let mut graph = ReasoningGraph::new();
            graph.add_node(LogicNode::new(
                "1".to_string(),
                LogicNodeType::Conclusion,
                summary.clone(),
                Some(summary.clone()) // Passing the summary directly as a translation to verify
            ));

            let is_sound = engine.evaluate_graph(&mut graph).await.unwrap_or(false);
            if !is_sound {
                // For MVP: Return error or silently drop/flag the invalid memory
                // We'll log and skip digestion if the logic is flawed
                println!("System 3 Validation Failed: Hallucination detected in memory digestion.");
                return Ok(0);
            }

            // 5. Generate embedding
            let embedding = self.embedder.embed(&summary).await?;

            // 6. Diffuse validated logic into Holographic Memory
            // Convert the discrete reasoning into a continuous latent representation (Diffusion / Entropy Injection)
            let serialized_graph = serde_json::to_string(&graph).unwrap_or_default();
            let neurogram = match self.holographic.diffuse_logic(&serialized_graph).await {
                Ok(data) => Some(data),
                Err(_) => None,
            };

            // 7. Store in Long-Term Memory
            let memory = MemoryNode {
                id: Uuid::new_v4().to_string(),
                content: summary,
                layer: 0,
                node_type: crate::NodeType::Summary,
                created_at: Utc::now().timestamp(),
                updated_at: Utc::now().timestamp(),
                embedding,
                metadata: std::collections::HashMap::new(),
                namespace: "default".to_string(),
                source: "metabolism".to_string(),
                holographic_data: neurogram,
                loop_level: 0,
            };

            self.memory.store(memory).await?;
            self.buffer.clear().await?;
            return Ok(count);
        }

        // Fallback for non-System3 path (Continuous Mode)
        // 5. Generate embedding
        let embedding = self.embedder.embed(&summary).await?;

        // 6. Compress into Neurogram (Holographic Memory)
        // This is the "Infinite Context" mechanism - compressing the history into a visual token.
        let neurogram = match self.holographic.compress_context(&combined_text).await {
            Ok(data) => Some(data),
            Err(_) => None, // Fallback if compression fails (or log warning)
        };

        // 7. Engram upsert (optional) - must happen before summary is moved
        if let Some(engram) = &self.engram {
            let ngrams = Self::extract_ngrams(&combined_text, 2, 3);
            if !ngrams.is_empty() {
                engram.upsert_ngrams(&ngrams, &summary).await?;
            }
        }

        // 8. Create MemoryNode (Layer 0)
        let node = MemoryNode {
            id: Uuid::new_v4().to_string(),
            content: summary,
            layer: 0,
            node_type: NodeType::Summary,
            created_at: Utc::now().timestamp(),
            updated_at: Utc::now().timestamp(),
            embedding,
            metadata: std::collections::HashMap::new(),
            namespace: "default".to_string(),
            source: "metabolism".to_string(),
            holographic_data: neurogram,
            loop_level: 1, // Promoted from Buffer(0) to ShortTerm(1)
        };

        // 9. Store in long-term memory
        self.memory.store(node).await?;

        Ok(interactions.len())
    }

    /// Set custom threshold for testing.
    pub fn with_threshold(mut self, threshold: usize) -> Self {
        self.threshold = threshold;
        self
    }

    /// Enable Engram memory upserts for this metabolism instance.
    pub fn with_engram(mut self, engram: Arc<dyn EngramPort>) -> Self {
        self.engram = Some(engram);
        self
    }

    fn extract_ngrams(text: &str, min_n: usize, max_n: usize) -> Vec<String> {
        let tokens: Vec<String> = text
            .split_whitespace()
            .map(|token| token.trim_matches(|c: char| !c.is_alphanumeric()).to_lowercase())
            .filter(|token| !token.is_empty())
            .collect();

        let mut ngrams = Vec::new();
        if tokens.is_empty() {
            return ngrams;
        }

        let max_n = max_n.min(tokens.len());
        for n in min_n..=max_n {
            for i in 0..=tokens.len().saturating_sub(n) {
                let gram = tokens[i..i + n].join(" ");
                ngrams.push(gram);
            }
        }

        ngrams
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Interaction;
    use async_trait::async_trait;
    use crate::ports::{SearchResult, SanitizerPort};

    // === Mock Sanitizer ===
    struct MockSanitizer;
    impl SanitizerPort for MockSanitizer {
        fn sanitize(&self, text: &str) -> Result<String> {
            Ok(text.replace("PII", "REDACTED"))
        }
    }

    // === Mock Buffer ===
    struct MockBuffer {
        items: tokio::sync::Mutex<Vec<Interaction>>,
    }

    impl MockBuffer {
        fn new() -> Self {
            Self {
                items: tokio::sync::Mutex::new(Vec::new()),
            }
        }
    }

    #[async_trait]
    impl BufferPort for MockBuffer {
        async fn push(&self, interaction: Interaction) -> Result<()> {
            self.items.lock().await.push(interaction);
            Ok(())
        }

        async fn pop_batch(&self, size: usize) -> Result<Vec<Interaction>> {
            let mut items = self.items.lock().await;
            let drain_count = size.min(items.len());
            Ok(items.drain(..drain_count).collect())
        }

        async fn peek(&self, size: usize) -> Result<Vec<Interaction>> {
            let items = self.items.lock().await;
            Ok(items.iter().take(size).cloned().collect())
        }

        async fn len(&self) -> Result<usize> {
            Ok(self.items.lock().await.len())
        }

        async fn is_empty(&self) -> Result<bool> {
            Ok(self.items.lock().await.is_empty())
        }

        async fn clear(&self) -> Result<()> {
            self.items.lock().await.clear();
            Ok(())
        }
    }

    // === Mock Memory ===
    struct MockMemory {
        nodes: tokio::sync::Mutex<Vec<MemoryNode>>,
    }

    impl MockMemory {
        fn new() -> Self {
            Self {
                nodes: tokio::sync::Mutex::new(Vec::new()),
            }
        }
    }

    #[async_trait]
    impl MemoryPort for MockMemory {
        async fn fetch_last_node(&self) -> Result<Option<MemoryNode>> {
            let nodes = self.nodes.lock().await;
            Ok(nodes.last().cloned())
        }

        async fn store(&self, node: MemoryNode) -> Result<String> {
            let id = node.id.clone();
            self.nodes.lock().await.push(node);
            Ok(id)
        }

        async fn search(&self, _embedding: &[f32], _top_k: usize) -> Result<Vec<SearchResult>> {
            Ok(vec![])
        }

        async fn search_layer(&self, _embedding: &[f32], _layer: u8, _top_k: usize) -> Result<Vec<SearchResult>> {
            Ok(vec![])
        }

        async fn search_namespace(&self, _embedding: &[f32], _namespace: &str, _top_k: usize) -> Result<Vec<SearchResult>> {
            Ok(vec![])
        }

        async fn get_by_id(&self, id: &str) -> Result<Option<MemoryNode>> {
            Ok(self.nodes.lock().await.iter().find(|n| n.id == id).cloned())
        }

        async fn get_by_layer(&self, layer: u8) -> Result<Vec<MemoryNode>> {
            Ok(self.nodes.lock().await.iter().filter(|n| n.layer == layer).cloned().collect())
        }

        async fn update(&self, _node: MemoryNode) -> Result<()> {
            Ok(())
        }

        async fn delete(&self, _id: &str) -> Result<()> {
            Ok(())
        }

        async fn count(&self) -> Result<usize> {
            Ok(self.nodes.lock().await.len())
        }

        async fn add_relationship(&self, _from_id: &str, _relation: &str, _to_id: &str) -> Result<()> {
            Ok(())
        }

        async fn count_by_layer(&self, layer: u8) -> Result<usize> {
            Ok(self.nodes.lock().await.iter().filter(|n| n.layer == layer).count())
        }

        async fn search_all(&self, limit: usize) -> Result<Vec<MemoryNode>> {
            Ok(self.nodes.lock().await.iter().take(limit).cloned().collect())
        }

        async fn get_by_hypertoken(&self, _hypertoken: &str) -> Result<Option<MemoryNode>> {
            Ok(None)
        }
    }

    // === Mock LLM ===
    struct MockLlm;

    #[async_trait]
    impl LlmPort for MockLlm {
        async fn generate(&self, _prompt: &str, _max_tokens: usize) -> Result<String> {
            Ok("Mock summary of conversation".to_string())
        }

        async fn generate_with_params(&self, _prompt: &str, _max_tokens: usize, _temp: f32, _top_p: f32) -> Result<String> {
            Ok("Mock summary".to_string())
        }

        async fn summarize(&self, text: &str) -> Result<String> {
            Ok(text.to_string())
        }
    }

    // === Mock Embedder ===
    struct MockEmbedder;

    #[async_trait]
    impl EmbeddingPort for MockEmbedder {
        async fn embed(&self, _text: &str) -> Result<Vec<f32>> {
            Ok(vec![0.1, 0.2, 0.3])
        }

        fn dimension(&self) -> usize {
            3
        }

        fn provider_name(&self) -> &str {
            "mock"
        }
    }

    // === Mock Holographic ===
    struct MockHolographic;

    #[async_trait]
    impl HolographicPort for MockHolographic {
        async fn compress_context(&self, _text: &str) -> Result<Vec<u8>> {
            Ok(vec![1, 2, 3, 4]) // Mock compressed data
        }

        async fn decompress_context(&self, _hologram: &[u8]) -> Result<String> {
            Ok("Mock decompressed context".to_string())
        }
    }

    #[tokio::test]
    async fn test_digest_under_threshold() {
        let buffer = Arc::new(MockBuffer::new());
        let memory = Arc::new(MockMemory::new());
        let llm = Arc::new(MockLlm);
        let embedder = Arc::new(MockEmbedder);
        let holographic = Arc::new(MockHolographic);
        let sanitizer = Arc::new(MockSanitizer);

        let metabolism = Metabolism::new(
            buffer.clone(),
            memory.clone(),
            llm,
            embedder,
            holographic,
            sanitizer,
        ).with_threshold(5);

        // Add only 3 interactions (below threshold of 5)
        for i in 0..3 {
            buffer.push(Interaction::new(
                format!("Question {}", i),
                format!("Answer {}", i),
            )).await.unwrap();
        }

        let digested = metabolism.digest().await.unwrap();
        assert_eq!(digested, 0, "Should not digest when below threshold");
        assert_eq!(buffer.len().await.unwrap(), 3, "Buffer should remain unchanged");
        assert_eq!(memory.count().await.unwrap(), 0, "No memory should be created");
    }

    #[tokio::test]
    async fn test_digest_creates_memory_node() {
        let buffer = Arc::new(MockBuffer::new());
        let memory = Arc::new(MockMemory::new());
        let llm = Arc::new(MockLlm);
        let embedder = Arc::new(MockEmbedder);
        let holographic = Arc::new(MockHolographic);
        let sanitizer = Arc::new(MockSanitizer);

        let metabolism = Metabolism::new(
            buffer.clone(),
            memory.clone(),
            llm,
            embedder,
            holographic,
            sanitizer,
        ).with_threshold(3);

        // Add 5 interactions (exceeds threshold of 3)
        for i in 0..5 {
            buffer.push(Interaction::new(
                format!("Question {}", i),
                format!("Answer {}", i),
            )).await.unwrap();
        }

        let digested = metabolism.digest().await.unwrap();

        assert_eq!(digested, 3, "Should digest threshold amount");
        assert_eq!(buffer.len().await.unwrap(), 2, "Buffer should have remaining items");
        assert_eq!(memory.count().await.unwrap(), 1, "One memory node should be created");

        // Verify the created node
        let nodes = memory.get_by_layer(0).await.unwrap();
        assert_eq!(nodes.len(), 1);
        assert_eq!(nodes[0].node_type, NodeType::Summary);
        assert_eq!(nodes[0].source, "metabolism");
        assert!(!nodes[0].embedding.is_empty());
    }

    #[tokio::test]
    async fn test_digest_full_pipeline() {
        let buffer = Arc::new(MockBuffer::new());
        let memory = Arc::new(MockMemory::new());
        let llm = Arc::new(MockLlm);
        let embedder = Arc::new(MockEmbedder);
        let holographic = Arc::new(MockHolographic);
        let sanitizer = Arc::new(MockSanitizer);

        let metabolism = Metabolism::new(
            buffer.clone(),
            memory.clone(),
            llm,
            embedder,
            holographic,
            sanitizer,
        ).with_threshold(2);

        // Fill buffer and digest multiple times
        for i in 0..6 {
            buffer.push(Interaction::new(
                format!("Q{}", i),
                format!("A{}", i),
            )).await.unwrap();
        }

        // First digest
        let d1 = metabolism.digest().await.unwrap();
        assert_eq!(d1, 2);

        // Second digest
        let d2 = metabolism.digest().await.unwrap();
        assert_eq!(d2, 2);

        // Third digest
        let d3 = metabolism.digest().await.unwrap();
        assert_eq!(d3, 2);

        // Buffer should be empty, 3 memory nodes created
        assert_eq!(buffer.len().await.unwrap(), 0);
        assert_eq!(memory.count().await.unwrap(), 3);
    }

    #[tokio::test]
    async fn test_infinite_loop_mixing() {
        let buffer = Arc::new(MockBuffer::new());
        let memory = Arc::new(MockMemory::new());
        let llm = Arc::new(MockLlm);
        let embedder = Arc::new(MockEmbedder);
        let holographic = Arc::new(MockHolographic);
        let sanitizer = Arc::new(MockSanitizer);

        let metabolism = Metabolism::new(
            buffer.clone(),
            memory.clone(),
            llm,
            embedder,
            holographic,
            sanitizer,
        ).with_threshold(1);

        // Digest 1: First context
        buffer.push(Interaction::new("Hello".to_string(), "World".to_string())).await.unwrap();
        metabolism.digest().await.unwrap();

        // Check it was stored with a neurogram
        let last = memory.fetch_last_node().await.unwrap().unwrap();
        assert!(last.holographic_data.is_some());

        // Digest 2: Should retrieve previous context
        buffer.push(Interaction::new("New".to_string(), "Data".to_string())).await.unwrap();
        metabolism.digest().await.unwrap();

        // This test mostly verifies the code path doesn't crash,
        // as MockHolographic returns constant data.
        // In a real scenario, we'd check if the content summary changed.
    }

    #[tokio::test]
    async fn test_sanitization_is_applied() {
        let buffer = Arc::new(MockBuffer::new());
        let memory = Arc::new(MockMemory::new());
        let llm = Arc::new(MockLlm);
        let embedder = Arc::new(MockEmbedder);
        let holographic = Arc::new(MockHolographic);
        let sanitizer = Arc::new(MockSanitizer);

        let metabolism = Metabolism::new(
            buffer.clone(),
            memory.clone(),
            llm,
            embedder,
            holographic,
            sanitizer,
        ).with_threshold(1);

        buffer.push(Interaction::new(
            "My email is PII".to_string(),
            "Thanks for the PII".to_string(),
        )).await.unwrap();

        metabolism.digest().await.unwrap();

        let node = memory.fetch_last_node().await.unwrap().unwrap();
        // The mock LLM now returns its input, so we can directly check
        // if the sanitized text is being processed.
        assert!(node.content.contains("REDACTED"));
    }
}