remem-ai 0.6.70

use anyhow::Result;
use rusqlite::{params, Connection};

use super::constants::{
    DREAM_MAX_CLUSTERS, DREAM_MIN_CLUSTER_SIZE, DREAM_RECENCY_GUARD_SECS, TOPIC_KEY_PREFIX_LEN,
};

#[derive(Debug, Clone)]
pub(crate) struct MemoryCandidate {
    pub id: i64,
    pub version: i64,
    pub topic_key: Option<String>,
    pub title: String,
    pub content: String,
    pub memory_type: String,
    #[allow(dead_code)]
    pub updated_at_epoch: i64,
}

/// A group of memories that are candidates for merging.
#[derive(Debug)]
pub(crate) struct Cluster {
    pub members: Vec<MemoryCandidate>,
}

/// Load active memories for a project and group them into merge candidates.
pub(super) fn load_clusters(conn: &Connection, project: &str) -> Result<Vec<Cluster>> {
    let cutoff = chrono::Utc::now().timestamp() - DREAM_RECENCY_GUARD_SECS;
    let current_filter =
        crate::memory::memory_current_filter_sql("m.status", "m.expires_at_epoch", false);
    let state_filter = crate::memory::memory_state_key_current_filter_sql("m");
    let policy_filter = crate::memory::suppression::memory_policy_filter_sql("m");

    let mut stmt = conn.prepare(&format!(
        "SELECT id, version, topic_key, title, content, memory_type, updated_at_epoch
         FROM memories m
         WHERE m.project = ?1
           AND {current_filter}
           AND {state_filter}
           AND {policy_filter}
           AND m.updated_at_epoch < ?2
           AND COALESCE(
                m.owner_scope,
                CASE WHEN COALESCE(m.scope, 'project') = 'global' THEN 'user' ELSE 'repo' END
           ) = 'repo'
           AND COALESCE(
                m.owner_key,
                CASE WHEN COALESCE(m.scope, 'project') = 'global' THEN 'user:default' ELSE m.project END
           ) = ?1
         ORDER BY m.memory_type, m.topic_key, m.updated_at_epoch DESC"
    ))?;

    let candidates: Vec<MemoryCandidate> = stmt
        .query_map(params![project, cutoff], |row| {
            Ok(MemoryCandidate {
                id: row.get(0)?,
                version: row.get(1)?,
                topic_key: row.get(2)?,
                title: row.get::<_, String>(3)?,
                content: row.get::<_, String>(4)?,
                memory_type: row.get::<_, String>(5)?,
                updated_at_epoch: row.get(6)?,
            })
        })?
        .collect::<rusqlite::Result<Vec<MemoryCandidate>>>()?;

    Ok(cluster_candidates(candidates))
}

fn cluster_candidates(candidates: Vec<MemoryCandidate>) -> Vec<Cluster> {
    use std::collections::HashMap;

    // Group keyed memories by (memory_type, topic_key prefix) and unkeyed ones by
    // (memory_type, None) so clustering stays within each memory-type bucket.
    let mut groups: HashMap<(String, Option<String>), Vec<MemoryCandidate>> = HashMap::new();

    for c in candidates {
        let memory_type = c.memory_type.clone();
        let group_key = match &c.topic_key {
            Some(key) if !key.is_empty() => {
                // Truncate to prefix length for grouping.
                (
                    memory_type,
                    Some(key.chars().take(TOPIC_KEY_PREFIX_LEN).collect::<String>()),
                )
            }
            // NULL or empty topic_key: group by memory_type only.
            _ => (memory_type, None),
        };
        groups.entry(group_key).or_default().push(c);
    }

    // Keep only groups with ≥ MIN_CLUSTER_SIZE, limit total clusters
    let mut clusters: Vec<Cluster> = groups
        .into_values()
        .filter(|g| g.len() >= DREAM_MIN_CLUSTER_SIZE)
        .map(|members| Cluster { members })
        .collect();

    // Sort by cluster size descending (biggest benefit first)
    clusters.sort_by_key(|b| std::cmp::Reverse(b.members.len()));
    clusters.truncate(DREAM_MAX_CLUSTERS);
    clusters
}

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

    fn make(id: i64, topic_key: Option<&str>, memory_type: &str) -> MemoryCandidate {
        MemoryCandidate {
            id,
            version: 1,
            topic_key: topic_key.map(str::to_owned),
            title: format!("title-{}", id),
            content: format!("content-{}", id),
            memory_type: memory_type.to_owned(),
            updated_at_epoch: 1000 + id,
        }
    }

    #[test]
    fn test_cluster_by_topic_key_prefix() {
        let candidates = vec![
            make(1, Some("auth-middleware-design-v1"), "decision"),
            make(2, Some("auth-middleware-design-v2"), "decision"),
            make(3, Some("totally-different-topic"), "decision"),
        ];
        let clusters = cluster_candidates(candidates);
        // auth-middleware-design-v1 and v2 share a 20-char prefix
        assert_eq!(clusters.len(), 1);
        assert_eq!(clusters[0].members.len(), 2);
    }

    #[test]
    fn test_cluster_null_topic_key_by_type() {
        let candidates = vec![
            make(1, None, "preference"),
            make(2, None, "preference"),
            make(3, None, "decision"),
        ];
        let clusters = cluster_candidates(candidates);
        // 2 preference + 1 decision; only preference group has ≥ 2
        assert_eq!(clusters.len(), 1);
        assert_eq!(clusters[0].members.len(), 2);
    }

    #[test]
    fn test_cluster_same_prefix_keeps_memory_types_separate() {
        let candidates = vec![
            make(1, Some("auth-middleware-design-v1"), "decision"),
            make(2, Some("auth-middleware-design-v2"), "decision"),
            make(3, Some("auth-middleware-design-v3"), "preference"),
            make(4, Some("auth-middleware-design-v4"), "preference"),
        ];
        let mut cluster_shapes: Vec<(String, usize)> = cluster_candidates(candidates)
            .into_iter()
            .map(|cluster| {
                let memory_type = cluster.members[0].memory_type.clone();
                assert!(cluster
                    .members
                    .iter()
                    .all(|member| member.memory_type == memory_type));
                (memory_type, cluster.members.len())
            })
            .collect();

        cluster_shapes.sort();
        assert_eq!(
            cluster_shapes,
            vec![("decision".to_string(), 2), ("preference".to_string(), 2)]
        );
    }

    #[test]
    fn test_single_member_cluster_excluded() {
        let candidates = vec![
            make(1, Some("unique-key-aaa"), "decision"),
            make(2, Some("unique-key-bbb"), "decision"),
        ];
        let clusters = cluster_candidates(candidates);
        assert!(clusters.is_empty());
    }

    #[test]
    fn test_max_clusters_respected() {
        let candidates: Vec<MemoryCandidate> = (0..200)
            .map(|i| make(i, Some(&format!("topic-{:04}-suffix", i / 2)), "decision"))
            .collect();
        let clusters = cluster_candidates(candidates);
        assert!(clusters.len() <= DREAM_MAX_CLUSTERS);
    }

    fn setup_memories_table(conn: &Connection) {
        conn.execute_batch(
            "CREATE TABLE memories (
                 id,
                 version      INTEGER NOT NULL DEFAULT 1,
                 project      TEXT,
                 status       TEXT,
                 scope        TEXT DEFAULT 'project',
                 owner_scope  TEXT,
                 owner_key    TEXT,
                 topic_key    TEXT,
                 title        TEXT,
                 content      TEXT,
                 memory_type  TEXT,
                 updated_at_epoch INTEGER,
                 expires_at_epoch INTEGER,
                 state_key_id INTEGER
             );
             CREATE TABLE memory_state_keys (
                 id INTEGER PRIMARY KEY,
                 current_memory_id INTEGER
             );
             CREATE TABLE memory_suppressions (
                 id INTEGER PRIMARY KEY,
                 target_kind TEXT NOT NULL,
                 target_id INTEGER,
                 target_value TEXT,
                 status TEXT NOT NULL
             );
             CREATE TABLE memory_entities (memory_id INTEGER, entity_id INTEGER);
             CREATE TABLE entities (id INTEGER PRIMARY KEY, canonical_name TEXT)",
        )
        .unwrap();
    }

    #[test]
    fn test_load_clusters_propagates_row_error() {
        // Row with `id` stored as TEXT passes the recency-guard WHERE clause
        // (updated_at_epoch = 0 is always older than now-3600) but causes
        // rusqlite `row.get::<_, i64>(0)?` to return InvalidType.
        // The old filter_map(|r| r.ok()) silently dropped such rows; the fix
        // propagates the error as Err.
        let conn = Connection::open_in_memory().unwrap();
        setup_memories_table(&conn);
        // Insert a row with TEXT in `id` (not coercible to i64) and
        // updated_at_epoch=0 so it passes the recency-guard filter.
        conn.execute(
            "INSERT INTO memories
             (id, project, status, topic_key, title, content, memory_type, updated_at_epoch)
             VALUES (?1, ?2, 'active', NULL, 'title', 'content', 'preference', 0)",
            rusqlite::params!["not-an-integer", "test-project"],
        )
        .unwrap();

        let result = load_clusters(&conn, "test-project");
        assert!(
            result.is_err(),
            "load_clusters must propagate row deserialization errors, not silently drop them"
        );
    }

    #[test]
    fn test_load_clusters_propagates_null_text_column_error() {
        // Rows with NULL title, content, or memory_type are data-integrity
        // violations: the fix must surface them as errors, not silently
        // substitute empty strings that corrupt downstream merge generation.
        let conn = Connection::open_in_memory().unwrap();
        setup_memories_table(&conn);
        // NULL title — updated_at_epoch=0 passes the recency guard.
        conn.execute(
            "INSERT INTO memories
             (id, project, status, topic_key, title, content, memory_type, updated_at_epoch)
             VALUES (1, 'test-project', 'active', NULL, NULL, 'content', 'preference', 0)",
            [],
        )
        .unwrap();

        let result = load_clusters(&conn, "test-project");
        assert!(
            result.is_err(),
            "load_clusters must propagate NULL title as an error, not silently replace it with \"\""
        );
    }

    #[test]
    fn test_load_clusters_excludes_global_user_owned_memories() -> anyhow::Result<()> {
        let conn = Connection::open_in_memory()?;
        setup_memories_table(&conn);
        for id in 1..=2 {
            conn.execute(
                "INSERT INTO memories
                 (id, project, status, scope, owner_scope, owner_key, topic_key, title, content,
                  memory_type, updated_at_epoch)
                 VALUES (?1, 'test-project', 'active', 'global', 'user', 'user:default',
                         'global-topic-shared', ?2, ?3, 'preference', 0)",
                rusqlite::params![
                    id,
                    format!("global title {id}"),
                    format!("global content {id}")
                ],
            )?;
        }

        let clusters = load_clusters(&conn, "test-project")?;

        assert!(
            clusters.is_empty(),
            "dream should not create repo merge clusters from global/user-owned memories"
        );
        Ok(())
    }
}