semantic-memory 0.5.8

use chrono::Utc;
use serde::{Deserialize, Serialize};
use std::cmp::Ordering;
use std::collections::{HashMap, HashSet};

/// A connected reasoning subgraph discovered from stored graph edges.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ReasoningSubgraph {
    /// Representative node for the component.
    pub root: String,
    /// Nodes participating in the component.
    pub nodes: Vec<String>,
    /// Directed edges inside the component.
    pub edges: Vec<(String, String)>,
    /// Aggregate access count for the nodes in this component.
    pub access_count: usize,
    /// Most recent access timestamp across component nodes (ISO-8601).
    pub last_accessed: String,
}

/// Access metadata consumed when ranking components and pruning candidates.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AccessLog {
    /// Item identifier referenced by edge traversal or graph lookup.
    pub item_id: String,
    /// Number of accesses.
    pub access_count: usize,
    /// Last access timestamp (ISO-8601).
    pub last_accessed: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
struct InternalReceiptSubgraph {
    root: String,
    nodes: Vec<String>,
    edges: Vec<(String, String)>,
    access_count: usize,
    last_accessed: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
/// Retained metadata to reconstruct and audit a pruned reasoning subgraph.
pub struct PruningReceipt {
    /// Root node associated with the pruned subgraph.
    pub subgraph_root: String,
    /// Nodes removed by pruning.
    pub pruned_nodes: Vec<String>,
    /// Provenance chain preserved for later forensic inspection.
    pub preserved_provenance: Vec<String>,
    /// Contradiction history preserved for this subgraph.
    pub preserved_contradictions: Vec<(String, String)>,
    /// Serialized snapshot used to rebuild the subgraph.
    pub reconstruction_data: String,
    /// Timestamp when the receipt was produced (ISO-8601).
    pub timestamp: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
/// Summary report from maintenance/cleanup planning stages.
pub struct MaintenanceReport {
    /// Number of identified reasoning subgraphs.
    pub subgraphs_identified: usize,
    /// Number of subgraphs selected for pruning.
    pub subgraphs_pruned: usize,
    /// Receipts generated by pruning.
    pub receipts: Vec<PruningReceipt>,
    /// Human-readable summary line for operator logs.
    pub summary: String,
}

#[derive(Debug, Clone)]
struct DisjointSet {
    parent: Vec<usize>,
    rank: Vec<u8>,
}

impl DisjointSet {
    fn new() -> Self {
        Self {
            parent: Vec::new(),
            rank: Vec::new(),
        }
    }

    fn add_node(&mut self) -> usize {
        let id = self.parent.len();
        self.parent.push(id);
        self.rank.push(0);
        id
    }

    fn find(&mut self, x: usize) -> usize {
        let parent = self.parent[x];
        if parent != x {
            let root = self.find(parent);
            self.parent[x] = root;
        }
        self.parent[x]
    }

    fn union(&mut self, a: usize, b: usize) {
        let ra = self.find(a);
        let rb = self.find(b);
        if ra == rb {
            return;
        }

        let (root, child) = match self.rank[ra].cmp(&self.rank[rb]) {
            Ordering::Less => (rb, ra),
            _ => (ra, rb),
        };

        if self.rank[root] == self.rank[child] && root != child {
            self.rank[root] += 1;
        }
        self.parent[child] = root;
    }
}

fn build_node_index(edges: &[(String, String)]) -> (HashMap<String, usize>, Vec<String>, DisjointSet) {
    let mut index = HashMap::new();
    let mut reverse = Vec::new();
    let mut dsu = DisjointSet::new();

    for (left, right) in edges {
        if !index.contains_key(left) {
            let id = dsu.add_node();
            index.insert(left.clone(), id);
            reverse.push(left.clone());
        }
        if !index.contains_key(right) {
            let id = dsu.add_node();
            index.insert(right.clone(), id);
            reverse.push(right.clone());
        }
        let left_id = index[left];
        let right_id = index[right];
        dsu.union(left_id, right_id);
    }

    (index, reverse, dsu)
}

fn build_access_lookup(access_logs: &[AccessLog]) -> HashMap<String, (usize, String)> {
    let mut lookup = HashMap::new();

    for log in access_logs {
        lookup.insert(log.item_id.clone(), (log.access_count, log.last_accessed.clone()));
    }

    lookup
}

fn summarize_access_info(nodes: &[String], access_lookup: &HashMap<String, (usize, String)>) -> (usize, String) {
    let mut access_count = 0usize;
    let mut last_accessed = "1970-01-01T00:00:00Z".to_string();

    for node in nodes {
        if let Some((count, timestamp)) = access_lookup.get(node) {
            access_count += count;
            if timestamp > &last_accessed {
                last_accessed = timestamp.clone();
            }
        }
    }

    (access_count, last_accessed)
}

fn sort_if_needed(items: &mut Vec<(String, String)>) {
    items.sort_unstable_by(|(left_a, right_a), (left_b, right_b)| {
        left_a
            .cmp(left_b)
            .then_with(|| right_a.cmp(right_b))
    });
}

/// Identify connected reasoning subgraphs with more than one node.
pub fn identify_subgraphs(
    edges: &[(String, String)],
    access_logs: &[AccessLog],
) -> Vec<ReasoningSubgraph> {
    if edges.is_empty() {
        return Vec::new();
    }

    let (node_to_index, index_to_node, mut dsu) = build_node_index(edges);
    let access_lookup = build_access_lookup(access_logs);

    let mut component_nodes: HashMap<usize, Vec<usize>> = HashMap::new();
    for idx in 0..index_to_node.len() {
        let root = dsu.find(idx);
        component_nodes.entry(root).or_default().push(idx);
    }

    let mut component_edges: HashMap<usize, Vec<(String, String)>> = HashMap::new();
    for (left, right) in edges {
        if let (Some(&left_id), Some(&_right_id)) = (node_to_index.get(left), node_to_index.get(right)) {
            let root = dsu.find(left_id);
            component_edges
                .entry(root)
                .or_default()
                .push((left.clone(), right.clone()));
        }
    }

    let mut subgraphs = Vec::new();
    for (root_id, nodes) in component_nodes {
        if nodes.len() <= 1 {
            continue;
        }

        let mut node_names: Vec<String> = nodes
            .into_iter()
            .map(|idx| index_to_node[idx].clone())
            .collect();
        node_names.sort_unstable();

        let mut component_edges = component_edges.remove(&root_id).unwrap_or_default();
        sort_if_needed(&mut component_edges);

        let (access_count, last_accessed) = summarize_access_info(&node_names, &access_lookup);
        let root = node_names[0].clone();
        let subgraph = ReasoningSubgraph {
            root,
            nodes: node_names,
            edges: component_edges,
            access_count,
            last_accessed,
        };
        subgraphs.push(subgraph);
    }

    subgraphs.sort_unstable_by(|left, right| left.root.cmp(&right.root));
    subgraphs
}

/// Construct a prunable receipt that preserves enough context for reconstruction and audit.
pub fn prune_subgraph(
    subgraph: &ReasoningSubgraph,
    contradictions: &[(String, String)],
) -> PruningReceipt {
    let mut contradiction_set = HashSet::new();
    let mut preserved_contradictions = Vec::new();
    let mut pruned_nodes = subgraph.nodes.clone();

    pruned_nodes.sort_unstable();
    pruned_nodes.dedup();

    let pruned_node_lookup: HashSet<&str> = pruned_nodes.iter().map(String::as_str).collect();

    for contradiction in contradictions {
        if contradiction_set.insert(contradiction) {
            if pruned_node_lookup.contains(contradiction.0.as_str())
                || pruned_node_lookup.contains(contradiction.1.as_str())
            {
                preserved_contradictions.push(contradiction.clone());
            }
        }
    }

    let serialization = InternalReceiptSubgraph {
        root: subgraph.root.clone(),
        nodes: subgraph.nodes.clone(),
        edges: subgraph.edges.clone(),
        access_count: subgraph.access_count,
        last_accessed: subgraph.last_accessed.clone(),
    };

    let reconstruction_data = match serde_json::to_string(&serialization) {
        Ok(value) => value,
        Err(_) => String::new(),
    };

    PruningReceipt {
        subgraph_root: subgraph.root.clone(),
        pruned_nodes,
        preserved_provenance: subgraph.nodes.clone(),
        preserved_contradictions,
        reconstruction_data,
        timestamp: Utc::now().to_rfc3339(),
    }
}

/// Rebuild a previously pruned reasoning subgraph.
pub fn reconstruct_subgraph(receipt: &PruningReceipt) -> ReasoningSubgraph {
    match serde_json::from_str::<InternalReceiptSubgraph>(&receipt.reconstruction_data) {
        Ok(data) => ReasoningSubgraph {
            root: data.root,
            nodes: data.nodes,
            edges: data.edges,
            access_count: data.access_count,
            last_accessed: data.last_accessed,
        },
        Err(_) => ReasoningSubgraph {
            root: receipt.subgraph_root.clone(),
            nodes: receipt.pruned_nodes.clone(),
            edges: Vec::new(),
            access_count: 0,
            last_accessed: "1970-01-01T00:00:00Z".to_string(),
        },
    }
}

/// Return pruning order metadata (`root`, `access_count`) sorted least-accessed first.
pub fn pruning_priority(subgraphs: &[ReasoningSubgraph]) -> Vec<(String, usize)> {
    let mut priorities: Vec<(String, usize)> = subgraphs
        .iter()
        .map(|subgraph| (subgraph.root.clone(), subgraph.access_count))
        .collect();
    priorities.sort_unstable_by(|a, b| a.1.cmp(&b.1));
    priorities
}

#[cfg(test)]
mod tests {
    use super::*;

    fn access(item_id: &str, access_count: usize, last_accessed: &str) -> AccessLog {
        AccessLog {
            item_id: item_id.to_string(),
            access_count,
            last_accessed: last_accessed.to_string(),
        }
    }

    #[test]
    fn identify_subgraphs_finds_connected_components() {
        let edges = vec![
            ("a".to_string(), "b".to_string()),
            ("b".to_string(), "c".to_string()),
            ("d".to_string(), "e".to_string()),
        ];
        let logs = vec![
            access("a", 10, "2026-01-01T00:00:00Z"),
            access("b", 5, "2026-01-02T00:00:00Z"),
            access("e", 3, "2026-01-03T00:00:00Z"),
        ];
        let subgraphs = identify_subgraphs(&edges, &logs);

        assert_eq!(subgraphs.len(), 2);
        assert_eq!(subgraphs[0].nodes.len(), 3);
        assert_eq!(subgraphs[0].edges.len(), 2);
        assert_eq!(subgraphs[0].root, "a".to_string());
        assert_eq!(subgraphs[0].access_count, 15);
        assert_eq!(subgraphs[0].last_accessed, "2026-01-02T00:00:00Z");
        assert_eq!(subgraphs[1].nodes.len(), 2);
        assert_eq!(subgraphs[1].edges.len(), 1);
        assert_eq!(subgraphs[1].root, "d".to_string());
        assert_eq!(subgraphs[1].access_count, 3);
        assert_eq!(subgraphs[1].last_accessed, "2026-01-03T00:00:00Z");
    }

    #[test]
    fn identify_subgraphs_excludes_single_node_components() {
        let edges = vec![
            ("a".to_string(), "b".to_string()),
            ("c".to_string(), "c".to_string()),
            ("d".to_string(), "e".to_string()),
        ];
        let subgraphs = identify_subgraphs(&edges, &[]);
        assert_eq!(subgraphs.len(), 2);
        assert!(!subgraphs.iter().any(|s| s.nodes.contains(&"c".to_string())));
    }

    #[test]
    fn prune_subgraph_records_pruned_nodes() {
        let subgraph = ReasoningSubgraph {
            root: "a".to_string(),
            nodes: vec!["a".to_string(), "b".to_string()],
            edges: vec![("a".to_string(), "b".to_string())],
            access_count: 42,
            last_accessed: "2026-01-01T00:00:00Z".to_string(),
        };
        let contradictions = vec![
            ("a".to_string(), "x".to_string()),
            ("x".to_string(), "y".to_string()),
        ];
        let receipt = prune_subgraph(&subgraph, &contradictions);
        assert_eq!(receipt.pruned_nodes, vec!["a".to_string(), "b".to_string()]);
    }

    #[test]
    fn prune_subgraph_preserves_contradictions() {
        let subgraph = ReasoningSubgraph {
            root: "a".to_string(),
            nodes: vec!["a".to_string(), "b".to_string(), "c".to_string()],
            edges: vec![("a".to_string(), "b".to_string()), ("b".to_string(), "c".to_string())],
            access_count: 11,
            last_accessed: "2026-01-01T00:00:00Z".to_string(),
        };
        let contradictions = vec![
            ("a".to_string(), "x".to_string()),
            ("y".to_string(), "c".to_string()),
            ("m".to_string(), "n".to_string()),
            ("b".to_string(), "c".to_string()),
        ];
        let receipt = prune_subgraph(&subgraph, &contradictions);
        assert_eq!(
            receipt.preserved_contradictions,
            vec![
                ("a".to_string(), "x".to_string()),
                ("y".to_string(), "c".to_string()),
                ("b".to_string(), "c".to_string())
            ]
        );
    }

    #[test]
    fn reconstruct_subgraph_rebuilds_original_from_receipt() {
        let subgraph = ReasoningSubgraph {
            root: "x".to_string(),
            nodes: vec!["x".to_string(), "y".to_string()],
            edges: vec![("x".to_string(), "y".to_string())],
            access_count: 7,
            last_accessed: "2026-01-01T00:00:00Z".to_string(),
        };
        let receipt = prune_subgraph(&subgraph, &[]);
        let rebuilt = reconstruct_subgraph(&receipt);
        assert_eq!(rebuilt, subgraph);
    }

    #[test]
    fn pruning_priority_returns_least_accessed_first() {
        let subgraphs = vec![
            ReasoningSubgraph {
                root: "a".to_string(),
                nodes: vec!["a".to_string()],
                edges: vec![],
                access_count: 15,
                last_accessed: "2026-01-01T00:00:00Z".to_string(),
            },
            ReasoningSubgraph {
                root: "b".to_string(),
                nodes: vec!["b".to_string()],
                edges: vec![],
                access_count: 2,
                last_accessed: "2026-01-01T00:00:00Z".to_string(),
            },
            ReasoningSubgraph {
                root: "c".to_string(),
                nodes: vec!["c".to_string()],
                edges: vec![],
                access_count: 9,
                last_accessed: "2026-01-01T00:00:00Z".to_string(),
            },
        ];
        let ordered = pruning_priority(&subgraphs);
        assert_eq!(
            ordered,
            vec![
                ("b".to_string(), 2),
                ("c".to_string(), 9),
                ("a".to_string(), 15),
            ]
        );
    }
}