selfware 0.6.3

//! Long-term store — persists consolidated records into JSON and indexes them
//! for hybrid semantic + causal-graph retrieval.

use anyhow::Result;
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, RwLock};
use tracing::{debug, info};

use super::temporal::TemporalRecord;
use crate::analysis::vector_store::{EmbeddingProvider, TfIdfEmbeddingProvider, VectorIndex};

/// Build a single text blob from a [`TemporalRecord`]'s compacted content
/// for embedding purposes.  Concatenates the summary with key facts,
/// entities, actions, outcomes, and insights so the embedding captures
/// the full semantic content of the record.
fn record_text(record: &TemporalRecord) -> String {
    let mut text = record.content.summary.clone();
    for fact in &record.content.key_facts {
        text.push(' ');
        text.push_str(fact);
    }
    for entity in &record.content.entities {
        text.push(' ');
        text.push_str(entity);
    }
    for action in &record.content.actions {
        text.push(' ');
        text.push_str(action);
    }
    for outcome in &record.content.outcomes {
        text.push(' ');
        text.push_str(outcome);
    }
    for insight in &record.content.insights {
        text.push(' ');
        text.push_str(insight);
    }
    text
}

/// Persists consolidated temporal records into long-term storage.
///
/// Records are persisted as JSON files (durable storage) and also indexed
/// in an in-memory vector index keyed by `record.id` for fast hybrid
/// retrieval combining:
///
/// 1. **Semantic vector search** — cosine similarity over embeddings of the
///    record's textual content (summary + key facts + entities + actions +
///    outcomes + insights), generated by an offline [`EmbeddingProvider`]
///    (default: [`TfIdfEmbeddingProvider`]).
/// 2. **Causal-graph traversal** — one-hop expansion via each record's
///    `causal_parents` and `causal_children` fields, so that retrieving a
///    semantically-relevant record also surfaces its immediate causal
///    context.
///
/// This is a "KG-lite" approach: the causal links already present on
/// [`TemporalRecord`] serve as the graph edges — no separate
/// `KnowledgeGraph` is needed.
pub struct LongTermStore {
    /// Collection name (kept for compatibility / metadata).
    collection_name: String,
    /// Directory for persisting records as JSON.
    storage_dir: std::path::PathBuf,
    /// Offline embedding provider (default: TfIdf).
    embedder: Arc<dyn EmbeddingProvider>,
    /// In-memory index: record_id -> embedding vector.
    embeddings: RwLock<HashMap<String, Vec<f32>>>,
    /// In-memory cache: record_id -> TemporalRecord.
    records: RwLock<HashMap<String, TemporalRecord>>,
}

impl LongTermStore {
    /// Create a new long-term store with the default offline TfIdf embedder.
    pub fn new(storage_dir: std::path::PathBuf) -> Self {
        Self::with_embedder(storage_dir, Arc::new(TfIdfEmbeddingProvider::default()))
    }

    /// Create a new long-term store with a custom embedding provider.
    ///
    /// Use this in tests with `MockEmbeddingProvider` or in production
    /// with `HttpEmbeddingProvider` for higher-quality embeddings.
    pub fn with_embedder(
        storage_dir: std::path::PathBuf,
        embedder: Arc<dyn EmbeddingProvider>,
    ) -> Self {
        Self {
            collection_name: "consolidated".into(),
            storage_dir,
            embedder,
            embeddings: RwLock::new(HashMap::new()),
            records: RwLock::new(HashMap::new()),
        }
    }

    /// Store consolidated records.
    ///
    /// 1. Persist each record as a JSON file (durable storage).
    /// 2. Embed each record's textual content and add it to the in-memory
    ///    vector index for semantic retrieval.
    pub async fn store(&self, records: &[TemporalRecord]) -> Result<StoreResult> {
        std::fs::create_dir_all(&self.storage_dir)?;

        let mut stored = 0;
        let mut errors = Vec::new();

        for record in records {
            let path = self.storage_dir.join(format!("{}.json", record.id));

            match serde_json::to_string_pretty(record) {
                Ok(json) => {
                    if let Err(e) = std::fs::write(&path, &json) {
                        errors.push(format!("Failed to write {}: {e}", record.id));
                    } else {
                        stored += 1;
                        debug!(
                            record_id = %record.id,
                            sources = record.source_ids.len(),
                            "Stored consolidated record"
                        );
                    }
                }
                Err(e) => {
                    errors.push(format!("Failed to serialize {}: {e}", record.id));
                }
            }
        }

        // Index records into the in-memory vector store for retrieval.
        // Embedding failures are non-fatal — the JSON is already persisted.
        for record in records {
            let text = record_text(record);
            match self.embedder.embed(&text).await {
                Ok(embedding) => {
                    let mut emb_map = self.embeddings.write().unwrap_or_else(|e| e.into_inner());
                    emb_map.insert(record.id.clone(), embedding);

                    let mut rec_map = self.records.write().unwrap_or_else(|e| e.into_inner());
                    rec_map.insert(record.id.clone(), record.clone());
                }
                Err(e) => {
                    errors.push(format!("Failed to embed {}: {e}", record.id));
                }
            }
        }

        info!(
            "Stored {stored}/{} records to {}",
            records.len(),
            self.storage_dir.display(),
        );

        Ok(StoreResult {
            stored,
            errors,
            collection: self.collection_name.clone(),
        })
    }

    /// Load all stored records from disk.
    pub fn load_all(&self) -> Result<Vec<TemporalRecord>> {
        let mut records = Vec::new();

        if !self.storage_dir.exists() {
            return Ok(records);
        }

        for entry in std::fs::read_dir(&self.storage_dir)? {
            let entry = entry?;
            let path = entry.path();
            if path.extension().and_then(|e| e.to_str()) == Some("json") {
                let content = std::fs::read_to_string(&path)?;
                match serde_json::from_str::<TemporalRecord>(&content) {
                    Ok(record) => records.push(record),
                    Err(e) => {
                        tracing::warn!("Failed to parse {}: {e}", path.display());
                    }
                }
            }
        }

        // Sort by sequence order
        records.sort_by_key(|r| r.sequence_order);

        info!("Loaded {} consolidated records", records.len());
        Ok(records)
    }

    /// Query records by tag.
    pub fn query_by_tag(&self, tag: &str) -> Result<Vec<TemporalRecord>> {
        let all = self.load_all()?;
        Ok(all
            .into_iter()
            .filter(|r| r.tags.iter().any(|t| t == tag))
            .collect())
    }

    /// Index all persisted records from disk into the in-memory vector index.
    ///
    /// Call this after creating a new `LongTermStore` (e.g. after a restart)
    /// to enable [`Self::retrieve`] without re-storing records.
    ///
    /// Returns the number of records indexed.
    pub async fn index_existing(&self) -> Result<usize> {
        let records = self.load_all()?;
        let count = records.len();

        for record in &records {
            let text = record_text(record);
            match self.embedder.embed(&text).await {
                Ok(embedding) => {
                    let mut emb_map = self.embeddings.write().unwrap_or_else(|e| e.into_inner());
                    emb_map.insert(record.id.clone(), embedding);

                    let mut rec_map = self.records.write().unwrap_or_else(|e| e.into_inner());
                    rec_map.insert(record.id.clone(), record.clone());
                }
                Err(e) => {
                    tracing::warn!("Failed to embed {}: {e}", record.id);
                }
            }
        }

        info!("Indexed {} existing records for retrieval", count);
        Ok(count)
    }

    /// Hybrid retrieval: semantic vector search + causal-graph expansion.
    ///
    /// This method combines two retrieval strategies:
    ///
    /// 1. **Semantic search**: The query is embedded using the same offline
    ///    [`EmbeddingProvider`] used for indexing. Cosine similarity is
    ///    computed against all indexed record embeddings, and the top-`k`
    ///    most similar records are selected as "semantic hits".
    ///
    /// 2. **Causal-graph expansion**: For each semantic hit, its 1-hop
    ///    causal neighbors (`causal_parents` + `causal_children`) are
    ///    looked up from the in-memory record cache. Neighbors not already
    ///    among the semantic hits are appended to the result list.
    ///
    /// The final result is ordered: semantic hits first (sorted by
    /// descending similarity), then causal neighbors. The total number of
    /// results is bounded at `k * 2` to keep the response size sensible.
    ///
    /// If the in-memory index is empty (e.g. `index_existing` was not
    /// called after a restart), this method automatically loads and indexes
    /// all persisted records from disk before proceeding.
    pub async fn retrieve(&self, query: &str, k: usize) -> Result<Vec<TemporalRecord>> {
        if k == 0 {
            return Ok(Vec::new());
        }

        // If the in-memory index is empty, populate from disk. Compute
        // emptiness in a tight scope so the read guard is released before the
        // await (never held across it).
        let needs_index = {
            let emb_map = self.embeddings.read().unwrap_or_else(|e| e.into_inner());
            emb_map.is_empty()
        };
        if needs_index {
            self.index_existing().await?;
        }

        // Embed the query with the same offline embedder.
        let query_embedding = self.embedder.embed(query).await?;

        // Compute cosine similarity for all indexed records.
        let emb_map = self.embeddings.read().unwrap_or_else(|e| e.into_inner());

        let mut scored: Vec<(String, f32)> = emb_map
            .iter()
            .map(|(id, emb)| {
                let sim = VectorIndex::cosine_similarity(&query_embedding, emb);
                (id.clone(), sim)
            })
            .collect();

        // Sort by descending similarity.
        scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        // Take top-k semantic hits.
        let semantic_hits: Vec<(String, f32)> = scored.into_iter().take(k).collect();

        drop(emb_map);

        // Collect semantic hit IDs (ordered by similarity).
        let mut result_ids: Vec<String> = semantic_hits.iter().map(|(id, _)| id.clone()).collect();
        let mut seen: HashSet<String> = result_ids.iter().cloned().collect();

        // Graph expansion: for each semantic hit, add 1-hop causal neighbors.
        let rec_map = self.records.read().unwrap_or_else(|e| e.into_inner());

        for (id, _) in &semantic_hits {
            if let Some(record) = rec_map.get(id) {
                for parent in &record.causal_parents {
                    if !seen.contains(parent) && rec_map.contains_key(parent) {
                        seen.insert(parent.clone());
                        result_ids.push(parent.clone());
                    }
                }
                for child in &record.causal_children {
                    if !seen.contains(child) && rec_map.contains_key(child) {
                        seen.insert(child.clone());
                        result_ids.push(child.clone());
                    }
                }
            }
        }

        // Cap the total at k * 2 to keep the result size sensible.
        let cap = k * 2;
        if result_ids.len() > cap {
            result_ids.truncate(cap);
        }

        // Load and return the TemporalRecords in order (semantic hits first,
        // then causal neighbors).
        let results: Vec<TemporalRecord> = result_ids
            .iter()
            .filter_map(|id| rec_map.get(id).cloned())
            .collect();

        drop(rec_map);
        Ok(results)
    }
}

/// Result of a store operation.
#[derive(Debug, Clone)]
pub struct StoreResult {
    pub stored: usize,
    pub errors: Vec<String>,
    pub collection: String,
}

#[cfg(test)]
#[path = "../../tests/unit/consolidation/store/store_test.rs"]
mod tests;