engram-retired/crates/engram-sync/src/engine.rs

/// SyncEngine — orchestrates peer sync and swarm activation.
///
/// The engine holds a list of known peers, a handle to the local database,
/// and owns our identity (UUID + API key). It drives two workflows:
///
/// 1. **Delta sync**: periodic pull-then-push with each peer, filtering by
///    tier allowlist and skipping nodes we already have.
///
/// 2. **Swarm activation**: run spreading activation locally, then fan out
///    to all trusted peers in parallel, and merge results by strength.
use crate::client::SyncClient;
use crate::types::{
    MergedActivatedNode, Peer, PeerActivationResult, PeerStatus, PeerSyncResult,
    SerializableActivatedNode, SyncConfig, SyncDelta, SyncReport, SwarmActivateRequest,
    SwarmActivateResponse,
};
use engram_core::types::{now_ms, MemoryTier};
use engram_core::EngramDb;
use std::sync::{Arc, Mutex};
use uuid::Uuid;

pub struct SyncEngine {
    db: Arc<Mutex<EngramDb>>,
    peers: Vec<Peer>,
    our_id: Uuid,
    our_name: String,
    api_key: String,
    default_sync_tiers: Vec<MemoryTier>,
}

impl SyncEngine {
    pub fn new(db: Arc<Mutex<EngramDb>>, config: SyncConfig) -> Self {
        Self {
            db,
            peers: Vec::new(),
            our_id: config.our_id,
            our_name: config.our_name,
            api_key: config.api_key,
            default_sync_tiers: config.default_sync_tiers,
        }
    }

    // ── Peer registry ─────────────────────────────────────────────────────────

    pub fn add_peer(&mut self, peer: Peer) {
        // Replace if already registered
        self.peers.retain(|p| p.id != peer.id);
        self.peers.push(peer);
    }

    pub fn remove_peer(&mut self, peer_id: Uuid) {
        self.peers.retain(|p| p.id != peer_id);
    }

    pub fn list_peers(&self) -> &[Peer] {
        &self.peers
    }

    pub fn get_peer(&self, peer_id: Uuid) -> Option<&Peer> {
        self.peers.iter().find(|p| p.id == peer_id)
    }

    pub fn our_id(&self) -> Uuid {
        self.our_id
    }

    pub fn our_name(&self) -> &str {
        &self.our_name
    }

    pub fn api_key(&self) -> &str {
        &self.api_key
    }

    // ── Full sync cycle ───────────────────────────────────────────────────────

    /// Sync with every registered peer. Returns a report summarising what moved.
    pub async fn sync_all(&mut self) -> SyncReport {
        let mut report = SyncReport {
            peers_synced: 0,
            nodes_received: 0,
            nodes_sent: 0,
            errors: Vec::new(),
        };

        // Clone the peer list so we can mutate self afterwards
        let peers: Vec<Peer> = self.peers.clone();
        for peer in peers {
            match self.sync_peer(&peer).await {
                Ok(result) => {
                    report.peers_synced += 1;
                    report.nodes_received += result.nodes_received;
                    report.nodes_sent += result.nodes_sent;
                    // Update last_sync_at
                    if let Some(p) = self.peers.iter_mut().find(|p| p.id == peer.id) {
                        p.last_sync_at = now_ms();
                    }
                }
                Err(e) => {
                    report.errors.push(format!("peer {}: {}", peer.name, e));
                }
            }
        }

        report
    }

    /// Sync a single peer: pull delta, apply it, then push our delta.
    pub async fn sync_peer(&self, peer: &Peer) -> anyhow::Result<PeerSyncResult> {
        let client = SyncClient::new(peer.clone(), self.our_id);
        let tiers = effective_tiers(peer, &self.default_sync_tiers);

        // Pull: get everything from the peer since their last known sync time
        let remote_delta = client.pull_delta(peer.last_sync_at).await?;
        let nodes_received = self.apply_delta(remote_delta, &tiers).await?;

        // Push: send everything we have that the peer hasn't seen
        let our_delta = self.generate_delta(peer.last_sync_at, &tiers)?;
        let nodes_sent = our_delta.nodes.len();
        client.push_delta(&our_delta).await?;

        Ok(PeerSyncResult {
            peer_id: peer.id,
            nodes_received,
            nodes_sent,
        })
    }

    // ── Delta generation ──────────────────────────────────────────────────────

    /// Build a delta containing all nodes/edges modified since `since`
    /// that belong to one of the allowed `tiers`.
    pub fn generate_delta(
        &self,
        since: i64,
        tiers: &[MemoryTier],
    ) -> anyhow::Result<SyncDelta> {
        let db = self
            .db
            .lock()
            .map_err(|_| anyhow::anyhow!("db lock poisoned"))?;

        // Scan all nodes; filter by tier and modification time
        let all_nodes = db.scan_nodes()?;
        let nodes: Vec<_> = all_nodes
            .into_iter()
            .filter(|n| tiers.contains(&n.tier) && n.last_activated >= since)
            .collect();

        // Edges: include all edges between nodes in our set
        let node_ids: std::collections::HashSet<Uuid> = nodes.iter().map(|n| n.id).collect();
        let all_edges = db.scan_edges()?;
        let edges: Vec<_> = all_edges
            .into_iter()
            .filter(|e| node_ids.contains(&e.from_id) && node_ids.contains(&e.to_id))
            .collect();

        Ok(SyncDelta {
            peer_id: self.our_id,
            since,
            nodes,
            edges,
            tombstones: Vec::new(), // tombstone tracking requires a separate log; not yet implemented
            generated_at: now_ms(),
        })
    }

    // ── Delta application ─────────────────────────────────────────────────────

    /// Merge an incoming delta into the local database.
    ///
    /// - Nodes/edges we already have (same UUID) are skipped — local wins.
    /// - Tombstones cause deletion.
    /// - Only nodes in the allowed tiers are accepted.
    ///
    /// Returns the number of nodes actually written.
    pub async fn apply_delta(
        &self,
        delta: SyncDelta,
        allowed_tiers: &[MemoryTier],
    ) -> anyhow::Result<usize> {
        let db = self
            .db
            .lock()
            .map_err(|_| anyhow::anyhow!("db lock poisoned"))?;

        let mut written = 0usize;

        // Apply tombstones first
        for id in &delta.tombstones {
            let _ = db.delete_node(*id);
        }

        // Merge nodes
        for node in delta.nodes {
            if !allowed_tiers.contains(&node.tier) {
                continue;
            }
            // Skip if we already have this UUID
            if db.get_node(node.id)?.is_some() {
                continue;
            }
            db.put_node(node)?;
            written += 1;
        }

        // Merge edges (if both endpoints exist)
        for edge in delta.edges {
            let from_exists = db.get_node(edge.from_id)?.is_some();
            let to_exists = db.get_node(edge.to_id)?.is_some();
            if from_exists && to_exists {
                db.put_edge(edge)?;
            }
        }

        Ok(written)
    }

    // ── Swarm activation ──────────────────────────────────────────────────────

    /// Run spreading activation locally, then fan out to all trusted peers,
    /// and merge all results into a unified ranked list.
    pub async fn swarm_activate(
        &self,
        req: SwarmActivateRequest,
    ) -> anyhow::Result<SwarmActivateResponse> {
        // Local activation
        let local_results: Vec<SerializableActivatedNode> = {
            let db = self
                .db
                .lock()
                .map_err(|_| anyhow::anyhow!("db lock poisoned"))?;
            let activated = db.activate(
                &req.seeds,
                &req.query_embedding,
                req.max_depth,
                req.limit,
            )?;
            activated.into_iter().map(Into::into).collect()
        };

        let mut peer_results: Vec<PeerActivationResult> = Vec::new();

        if req.include_peers {
            // Fan out to all trusted peers in parallel
            let trusted_peers: Vec<Peer> = self
                .peers
                .iter()
                .filter(|p| p.trusted)
                .cloned()
                .collect();

            let mut handles = Vec::new();
            for peer in trusted_peers {
                let seeds = req.seeds.clone();
                let embedding = req.query_embedding.clone();
                let max_depth = req.max_depth;
                let limit = req.limit;
                let our_id = self.our_id;

                handles.push(tokio::spawn(async move {
                    let client = SyncClient::new(peer.clone(), our_id);
                    match client
                        .remote_activate(&seeds, &embedding, max_depth, limit)
                        .await
                    {
                        Ok(results) => PeerActivationResult {
                            peer_id: peer.id,
                            peer_name: peer.name,
                            results,
                            error: None,
                        },
                        Err(e) => PeerActivationResult {
                            peer_id: peer.id,
                            peer_name: peer.name,
                            results: Vec::new(),
                            error: Some(e.to_string()),
                        },
                    }
                }));
            }

            for handle in handles {
                match handle.await {
                    Ok(result) => peer_results.push(result),
                    Err(e) => peer_results.push(PeerActivationResult {
                        peer_id: Uuid::nil(),
                        peer_name: "unknown".into(),
                        results: Vec::new(),
                        error: Some(format!("task error: {}", e)),
                    }),
                }
            }
        }

        let merged = merge_activation_results(&local_results, &peer_results, req.limit);

        Ok(SwarmActivateResponse {
            local_results,
            peer_results,
            merged,
        })
    }

    // ── Peer health check ─────────────────────────────────────────────────────

    /// Check which peers are reachable and return their status.
    pub async fn peer_statuses(&self) -> Vec<PeerStatus> {
        let mut statuses = Vec::new();
        for peer in &self.peers {
            let client = SyncClient::new(peer.clone(), self.our_id);
            // We attempt a health check by pulling an empty delta (since=now)
            let reachable = client.pull_delta(now_ms()).await.is_ok();
            statuses.push(PeerStatus {
                peer_id: peer.id,
                peer_name: peer.name.clone(),
                address: peer.address.clone(),
                last_sync_at: peer.last_sync_at,
                reachable,
                sync_tiers: peer.sync_tiers.clone(),
                trusted: peer.trusted,
            });
        }
        statuses
    }
}

// ── Merge logic ───────────────────────────────────────────────────────────────

/// Merge local and peer activation results into a unified ranked list.
///
/// Deduplication: two nodes are considered duplicates if they have the same UUID.
/// When duplicates occur, the one with the highest activation strength is kept.
/// The merged list is sorted by activation_strength descending and truncated to `limit`.
pub fn merge_activation_results(
    local: &[SerializableActivatedNode],
    peer_results: &[PeerActivationResult],
    limit: usize,
) -> Vec<MergedActivatedNode> {
    use std::collections::HashMap;

    // (uuid -> MergedActivatedNode) — we keep the strongest instance of each node
    let mut map: HashMap<Uuid, MergedActivatedNode> = HashMap::new();

    // Process local results first (source_peer = None)
    for item in local {
        let id = item.node.id;
        let candidate = MergedActivatedNode {
            content: String::from_utf8_lossy(&item.node.content).to_string(),
            node_type: item.node.node_type.clone(),
            tier: item.node.tier.clone(),
            activation_strength: item.activation_strength,
            source_peer: None,
            hops: item.hops,
            node: item.node.clone(),
        };
        map.entry(id)
            .and_modify(|existing| {
                if item.activation_strength > existing.activation_strength {
                    *existing = candidate.clone();
                }
            })
            .or_insert(candidate);
    }

    // Process peer results
    for peer_result in peer_results {
        if peer_result.error.is_some() {
            continue;
        }
        for item in &peer_result.results {
            let id = item.node.id;
            let candidate = MergedActivatedNode {
                content: String::from_utf8_lossy(&item.node.content).to_string(),
                node_type: item.node.node_type.clone(),
                tier: item.node.tier.clone(),
                activation_strength: item.activation_strength,
                source_peer: Some(peer_result.peer_id),
                hops: item.hops,
                node: item.node.clone(),
            };
            map.entry(id)
                .and_modify(|existing| {
                    if item.activation_strength > existing.activation_strength {
                        *existing = candidate.clone();
                    }
                })
                .or_insert(candidate);
        }
    }

    // Sort by strength descending, take top N
    let mut merged: Vec<MergedActivatedNode> = map.into_values().collect();
    merged.sort_by(|a, b| {
        b.activation_strength
            .partial_cmp(&a.activation_strength)
            .unwrap_or(std::cmp::Ordering::Equal)
    });
    merged.truncate(limit);
    merged
}

// ── Tier helpers ──────────────────────────────────────────────────────────────

/// Compute the effective sync tiers for a peer — intersection of the peer's
/// own allowlist and our defaults. Untrusted peers are further restricted
/// to Semantic only.
fn effective_tiers<'a>(peer: &Peer, defaults: &'a [MemoryTier]) -> Vec<MemoryTier> {
    if !peer.trusted {
        // Untrusted: Semantic only, regardless of configuration
        return vec![MemoryTier::Semantic];
    }
    if peer.sync_tiers.is_empty() {
        defaults.to_vec()
    } else {
        // Intersection: only tiers that both us and the peer agree on
        defaults
            .iter()
            .filter(|t| peer.sync_tiers.contains(t))
            .cloned()
            .collect()
    }
}