poimen-memory/crates/mem-store/src/pg_repo.rs

use anyhow::{anyhow, Result};
use pgvector::Vector;
use serde::{Deserialize, Serialize};
use sqlx::{PgPool, Row};

/// Memory level
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum Level {
    L0,
    L1,
    L2,
    R,
}

impl std::fmt::Display for Level {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Level::L0 => write!(f, "L0"),
            Level::L1 => write!(f, "L1"),
            Level::L2 => write!(f, "L2"),
            Level::R => write!(f, "R"),
        }
    }
}

impl std::str::FromStr for Level {
    type Err = anyhow::Error;
    fn from_str(s: &str) -> Result<Self> {
        match s {
            "L0" => Ok(Level::L0),
            "L1" => Ok(Level::L1),
            "L2" => Ok(Level::L2),
            "R" => Ok(Level::R),
            _ => Err(anyhow!("invalid level: {}", s)),
        }
    }
}

/// Vector kind (text embedding or symptom projection)
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum VectorKind {
    Text,
    Symptom,
}

impl std::fmt::Display for VectorKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            VectorKind::Text => write!(f, "text"),
            VectorKind::Symptom => write!(f, "symptom"),
        }
    }
}

/// Search scope (project-specific or federated)
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Scope {
    Project(String),  // project-specific search
    AllProjects,      // federated across all projects (for tool-failure lookups)
}

/// Memory node (content-addressable by sha256)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryNode {
    pub sha256: String,
    pub level: Level,
    pub project: String,
    pub query_id: Option<String>,  // NULL at L2, R
    pub run_id: String,
    pub t: i32,
    pub source: Option<String>,  // set at L0; URI at R
    pub text: String,
}

/// Scored search result with metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScoredNode {
    pub node: MemoryNode,
    pub distance: f32,  // raw cosine distance (not similarity)
    pub matched_kind: VectorKind,
}

/// Failure signature hit (exact-match tier)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SignatureHit {
    pub node_sha: String,
    pub tool: String,
    pub raw: String,
    pub seen_count: i32,
}

/// PostgreSQL repository for memory projection
pub struct PgRepo {
    pool: PgPool,
}

impl PgRepo {
    /// Connect to Postgres and run migrations
    pub async fn connect(url: &str) -> Result<Self> {
        let pool = PgPool::connect(url).await?;
        sqlx::migrate!("./migrations")
            .run(&pool)
            .await?;
        Ok(Self { pool })
    }

    /// Upsert a memory node (ON CONFLICT DO NOTHING — idempotent)
    pub async fn upsert_node(&self, node: &MemoryNode) -> Result<()> {
        let level_str = node.level.to_string();
        sqlx::query(
            r#"
            INSERT INTO memory_node
            (sha256, level, project, query_id, run_id, t, source, text)
            VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
            ON CONFLICT (sha256) DO NOTHING
            "#,
        )
        .bind(&node.sha256)
        .bind(&level_str)
        .bind(&node.project)
        .bind(&node.query_id)
        .bind(&node.run_id)
        .bind(node.t)
        .bind(&node.source)
        .bind(&node.text)
        .execute(&self.pool)
        .await?;
        Ok(())
    }

    /// Store an embedding for a node (text or symptom kind)
    pub async fn upsert_vector(
        &self,
        node_sha: &str,
        kind: VectorKind,
        embedding: &[f32],
    ) -> Result<()> {
        if embedding.len() != 768 {
            return Err(anyhow!(
                "invalid embedding dimension: expected 768, got {}",
                embedding.len()
            ));
        }
        let kind_str = kind.to_string();
        let vec = Vector::from(embedding.to_vec());
        sqlx::query(
            r#"
            INSERT INTO memory_vector (node_sha, kind, embedding)
            VALUES ($1, $2, $3)
            ON CONFLICT (node_sha, kind) DO UPDATE
            SET embedding = $3
            "#,
        )
        .bind(node_sha)
        .bind(&kind_str)
        .bind(&vec)
        .execute(&self.pool)
        .await?;
        Ok(())
    }

    /// Insert edges (child → parent pointers) — fails if endpoints don't exist
    pub async fn insert_edges(&self, child_sha: &str, parent_shas: &[String]) -> Result<()> {
        // Batch insert edges; foreign key constraints ensure endpoints exist
        for parent_sha in parent_shas {
            sqlx::query(
                r#"
                INSERT INTO memory_edge (child_sha, parent_sha)
                VALUES ($1, $2)
                ON CONFLICT (child_sha, parent_sha) DO NOTHING
                "#,
            )
            .bind(child_sha)
            .bind(parent_sha)
            .execute(&self.pool)
            .await?;
        }
        Ok(())
    }

    /// Search by embedding (cosine distance, kind-filtered)
    /// - For text kind: uses partial index WHERE kind = 'text'
    /// - For symptom kind: uses partial index WHERE kind = 'symptom'
    pub async fn search(
        &self,
        query_embedding: &[f32],
        kind: VectorKind,
        levels: &[Level],
        scope: &Scope,
        k: usize,
    ) -> Result<Vec<ScoredNode>> {
        if query_embedding.len() != 768 {
            return Err(anyhow!("query embedding must be 768-dim"));
        }

        let q_vec = Vector::from(query_embedding.to_vec());
        let kind_str = kind.to_string();
        let level_strs: Vec<String> = levels.iter().map(|l| l.to_string()).collect();

        // Build the WHERE clause based on scope
        let (where_clause, project_param) = match scope {
            Scope::Project(proj) => ("AND n.project = $4".to_string(), Some(proj.clone())),
            Scope::AllProjects => ("".to_string(), None),
        };

        let query_sql = format!(
            r#"
            SELECT n.sha256, n.level, n.project, n.query_id, n.run_id, n.t, n.source, n.text,
                   v.embedding <=> $1 AS distance
            FROM memory_vector v
            JOIN memory_node n ON v.node_sha = n.sha256
            WHERE v.kind = $2
              AND n.level = ANY($3)
              {}
              AND (SELECT COUNT(*) FROM memory_supersede WHERE old_sha = n.sha256) = 0
            ORDER BY v.embedding <=> $1 ASC
            LIMIT $5
            "#,
            if where_clause.is_empty() { "" } else { &where_clause }
        );

        let rows = if let Some(proj) = project_param {
            sqlx::query(&query_sql)
                .bind(&q_vec)
                .bind(&kind_str)
                .bind(&level_strs)
                .bind(&proj)
                .bind(k as i64)
                .fetch_all(&self.pool)
                .await?
        } else {
            sqlx::query(&query_sql)
                .bind(&q_vec)
                .bind(&kind_str)
                .bind(&level_strs)
                .bind(k as i64)
                .fetch_all(&self.pool)
                .await?
        };

        let mut results = Vec::new();
        for row in rows {
            let level_str: String = row.get("level");
            let node = MemoryNode {
                sha256: row.get("sha256"),
                level: level_str.parse()?,
                project: row.get("project"),
                query_id: row.get("query_id"),
                run_id: row.get("run_id"),
                t: row.get("t"),
                source: row.get("source"),
                text: row.get("text"),
            };
            let distance: f32 = row.get("distance");
            results.push(ScoredNode {
                node,
                distance,
                matched_kind: kind,
            });
        }

        Ok(results)
    }

    /// Exact-match lookup on failure signature
    pub async fn lookup_signature(&self, sig_sha: &str) -> Result<Option<SignatureHit>> {
        let row = sqlx::query(
            r#"
            SELECT node_sha, tool, raw, seen_count
            FROM failure_signature
            WHERE sig_sha = $1
            "#,
        )
        .bind(sig_sha)
        .fetch_optional(&self.pool)
        .await?;

        Ok(row.map(|r| SignatureHit {
            node_sha: r.get("node_sha"),
            tool: r.get("tool"),
            raw: r.get("raw"),
            seen_count: r.get("seen_count"),
        }))
    }

    /// Traverse parent nodes via edges
    pub async fn parents_of(&self, sha: &str) -> Result<Vec<MemoryNode>> {
        let rows = sqlx::query(
            r#"
            SELECT n.sha256, n.level, n.project, n.query_id, n.run_id, n.t, n.source, n.text
            FROM memory_node n
            JOIN memory_edge e ON e.parent_sha = n.sha256
            WHERE e.child_sha = $1
            "#,
        )
        .bind(sha)
        .fetch_all(&self.pool)
        .await?;

        let mut parents = Vec::new();
        for row in rows {
            let level_str: String = row.get("level");
            parents.push(MemoryNode {
                sha256: row.get("sha256"),
                level: level_str.parse()?,
                project: row.get("project"),
                query_id: row.get("query_id"),
                run_id: row.get("run_id"),
                t: row.get("t"),
                source: row.get("source"),
                text: row.get("text"),
            });
        }
        Ok(parents)
    }

    /// Clear all nodes for a project (edges cascade delete)
    pub async fn clear_project(&self, project: &str) -> Result<()> {
        sqlx::query("DELETE FROM memory_node WHERE project = $1")
            .bind(project)
            .execute(&self.pool)
            .await?;
        Ok(())
    }

    /// Get node count (for testing)
    pub async fn node_count(&self) -> Result<i64> {
        let row = sqlx::query("SELECT COUNT(*) as cnt FROM memory_node")
            .fetch_one(&self.pool)
            .await?;
        Ok(row.get("cnt"))
    }

    /// Get edge count (for testing)
    pub async fn edge_count(&self) -> Result<i64> {
        let row = sqlx::query("SELECT COUNT(*) as cnt FROM memory_edge")
            .fetch_one(&self.pool)
            .await?;
        Ok(row.get("cnt"))
    }
}

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

    #[test]
    fn test_level_display() {
        assert_eq!(Level::L0.to_string(), "L0");
        assert_eq!(Level::L1.to_string(), "L1");
        assert_eq!(Level::L2.to_string(), "L2");
        assert_eq!(Level::R.to_string(), "R");
    }

    #[test]
    fn test_vector_kind_display() {
        assert_eq!(VectorKind::Text.to_string(), "text");
        assert_eq!(VectorKind::Symptom.to_string(), "symptom");
    }

    #[test]
    fn test_scope_variants() {
        let proj_scope = Scope::Project("test".to_string());
        let all_scope = Scope::AllProjects;
        assert_ne!(proj_scope, all_scope);
    }
}