consciousness/src/query.rs

// query.rs — peg-based query language for the memory graph
//
// Grammar-driven: the peg definition IS the language spec.
// Evaluates against node properties, graph metrics, and edge attributes.
// Designed for ad-hoc exploration without memorizing 35+ subcommands.
//
// Syntax:
//   expr | stage | stage ...
//
// Stages (piped):
//   sort FIELD          sort descending (default for exploration)
//   sort FIELD asc      sort ascending
//   limit N             cap results
//   select F,F,...      output specific fields as TSV
//   count               just show count
//
// Examples:
//   degree > 15 | sort degree | limit 10
//   category = core | select degree,weight
//   neighbors('identity.md') WHERE strength > 0.5 | sort strength
//   key ~ 'journal.*' AND degree > 10 | count
//   * | sort weight asc | limit 20

use crate::capnp_store::{NodeType, Provenance, RelationType, Store};
use crate::graph::Graph;
use regex::Regex;

// -- AST types --

#[derive(Debug, Clone)]
pub enum Expr {
    All,
    Comparison { field: String, op: CmpOp, value: Value },
    And(Box<Expr>, Box<Expr>),
    Or(Box<Expr>, Box<Expr>),
    Not(Box<Expr>),
    Neighbors { key: String, filter: Option<Box<Expr>> },
}

#[derive(Debug, Clone)]
pub enum Value {
    Num(f64),
    Str(String),
    Ident(String),
    FnCall(FnCall),
}

#[derive(Debug, Clone)]
pub enum FnCall {
    Community(String),
    Degree(String),
}

#[derive(Debug, Clone, Copy)]
pub enum CmpOp {
    Gt, Lt, Ge, Le, Eq, Ne, Match,
}

#[derive(Debug, Clone)]
pub enum Stage {
    Sort { field: String, ascending: bool },
    Limit(usize),
    Select(Vec<String>),
    Count,
}

#[derive(Debug, Clone)]
pub struct Query {
    pub expr: Expr,
    pub stages: Vec<Stage>,
}

// -- PEG grammar --

peg::parser! {
    pub grammar query_parser() for str {
        rule _() = [' ' | '\t']*

        pub rule query() -> Query
            = e:expr() s:stages() { Query { expr: e, stages: s } }

        rule stages() -> Vec<Stage>
            = s:(_ "|" _ s:stage() { s })* { s }

        rule stage() -> Stage
            = "sort" _ f:field() _ a:asc_desc() { Stage::Sort { field: f, ascending: a } }
            / "limit" _ n:integer() { Stage::Limit(n) }
            / "select" _ f:field_list() { Stage::Select(f) }
            / "count" { Stage::Count }

        rule asc_desc() -> bool
            = "asc" { true }
            / "desc" { false }
            / { false }  // default: descending

        rule field_list() -> Vec<String>
            = f:field() fs:(_ "," _ f:field() { f })* {
                let mut v = vec![f];
                v.extend(fs);
                v
            }

        rule integer() -> usize
            = n:$(['0'..='9']+) { n.parse().unwrap() }

        pub rule expr() -> Expr = precedence! {
            a:(@) _ "OR" _ b:@ { Expr::Or(Box::new(a), Box::new(b)) }
            --
            a:(@) _ "AND" _ b:@ { Expr::And(Box::new(a), Box::new(b)) }
            --
            "NOT" _ e:@ { Expr::Not(Box::new(e)) }
            --
            "neighbors" _ "(" _ k:string() _ ")" _ w:where_clause()? {
                Expr::Neighbors { key: k, filter: w.map(Box::new) }
            }
            f:field() _ op:cmp_op() _ v:value() {
                Expr::Comparison { field: f, op, value: v }
            }
            "*" { Expr::All }
            "(" _ e:expr() _ ")" { e }
        }

        rule where_clause() -> Expr
            = "WHERE" _ e:expr() { e }

        rule field() -> String
            = s:$(['a'..='z' | 'A'..='Z' | '_']['a'..='z' | 'A'..='Z' | '0'..='9' | '_']*) {
                s.to_string()
            }

        rule cmp_op() -> CmpOp
            = ">=" { CmpOp::Ge }
            / "<=" { CmpOp::Le }
            / "!=" { CmpOp::Ne }
            / ">"  { CmpOp::Gt }
            / "<"  { CmpOp::Lt }
            / "="  { CmpOp::Eq }
            / "~"  { CmpOp::Match }

        rule value() -> Value
            = f:fn_call() { Value::FnCall(f) }
            / n:number() { Value::Num(n) }
            / s:string() { Value::Str(s) }
            / i:ident()  { Value::Ident(i) }

        rule fn_call() -> FnCall
            = "community" _ "(" _ k:string() _ ")" { FnCall::Community(k) }
            / "degree" _ "(" _ k:string() _ ")" { FnCall::Degree(k) }

        rule number() -> f64
            = n:$(['0'..='9']+ ("." ['0'..='9']+)?) {
                n.parse().unwrap()
            }

        rule string() -> String
            = "'" s:$([^ '\'']*) "'" { s.to_string() }

        rule ident() -> String
            = s:$(['a'..='z' | 'A'..='Z' | '_']['a'..='z' | 'A'..='Z' | '0'..='9' | '_' | '-' | '.']*) {
                s.to_string()
            }
    }
}

// -- Field resolution --

/// Resolve a field value from a node + graph context, returning a comparable Value.
fn resolve_field(field: &str, key: &str, store: &Store, graph: &Graph) -> Option<Value> {
    let node = store.nodes.get(key)?;
    match field {
        "key"           => Some(Value::Str(key.to_string())),
        "weight"        => Some(Value::Num(node.weight as f64)),
        "category"      => Some(Value::Str(node.category.label().to_string())),
        "node_type"     => Some(Value::Str(node_type_label(node.node_type).to_string())),
        "provenance"    => Some(Value::Str(provenance_label(node.provenance).to_string())),
        "emotion"       => Some(Value::Num(node.emotion as f64)),
        "retrievals"    => Some(Value::Num(node.retrievals as f64)),
        "uses"          => Some(Value::Num(node.uses as f64)),
        "wrongs"        => Some(Value::Num(node.wrongs as f64)),
        "created"       => Some(Value::Str(node.created.clone())),
        "content"       => Some(Value::Str(node.content.clone())),
        "degree"        => Some(Value::Num(graph.degree(key) as f64)),
        "community_id"  => {
            graph.communities().get(key).map(|&c| Value::Num(c as f64))
        }
        "clustering_coefficient" => {
            Some(Value::Num(graph.clustering_coefficient(key) as f64))
        }
        "schema_fit"    => {
            node.schema_fit.map(|f| Value::Num(f as f64))
        }
        _ => None,
    }
}

fn node_type_label(nt: NodeType) -> &'static str {
    match nt {
        NodeType::EpisodicSession => "episodic_session",
        NodeType::EpisodicDaily   => "episodic_daily",
        NodeType::EpisodicWeekly  => "episodic_weekly",
        NodeType::Semantic        => "semantic",
    }
}

fn provenance_label(p: Provenance) -> &'static str {
    match p {
        Provenance::Manual  => "manual",
        Provenance::Journal => "journal",
        Provenance::Agent   => "agent",
        Provenance::Dream   => "dream",
        Provenance::Derived => "derived",
    }
}

fn rel_type_label(r: RelationType) -> &'static str {
    match r {
        RelationType::Link   => "link",
        RelationType::Causal => "causal",
        RelationType::Auto   => "auto",
    }
}

// -- Comparison logic --

fn as_num(v: &Value) -> Option<f64> {
    match v {
        Value::Num(n) => Some(*n),
        Value::Str(s) => s.parse().ok(),
        Value::Ident(s) => s.parse().ok(),
        Value::FnCall(_) => None,
    }
}

fn as_str(v: &Value) -> String {
    match v {
        Value::Str(s) | Value::Ident(s) => s.clone(),
        Value::Num(n) => format!("{}", n),
        Value::FnCall(_) => String::new(),
    }
}

fn compare(lhs: &Value, op: CmpOp, rhs: &Value) -> bool {
    match op {
        CmpOp::Match => {
            let text = as_str(lhs);
            let pat = as_str(rhs);
            match Regex::new(&pat) {
                Ok(re) => re.is_match(&text),
                Err(_) => false,
            }
        }
        CmpOp::Eq => {
            if let (Some(a), Some(b)) = (as_num(lhs), as_num(rhs)) {
                a == b
            } else {
                as_str(lhs) == as_str(rhs)
            }
        }
        CmpOp::Ne => {
            if let (Some(a), Some(b)) = (as_num(lhs), as_num(rhs)) {
                a != b
            } else {
                as_str(lhs) != as_str(rhs)
            }
        }
        CmpOp::Gt | CmpOp::Lt | CmpOp::Ge | CmpOp::Le => {
            let a = as_num(lhs).unwrap_or(f64::NAN);
            let b = as_num(rhs).unwrap_or(f64::NAN);
            match op {
                CmpOp::Gt => a > b,
                CmpOp::Lt => a < b,
                CmpOp::Ge => a >= b,
                CmpOp::Le => a <= b,
                _ => unreachable!(),
            }
        }
    }
}

// -- Evaluator --

fn resolve_fn(f: &FnCall, store: &Store, graph: &Graph) -> Value {
    match f {
        FnCall::Community(key) => {
            let resolved = store.resolve_key(key).unwrap_or_else(|_| key.clone());
            graph.communities().get(&resolved)
                .map(|&c| Value::Num(c as f64))
                .unwrap_or(Value::Num(f64::NAN))
        }
        FnCall::Degree(key) => {
            let resolved = store.resolve_key(key).unwrap_or_else(|_| key.clone());
            Value::Num(graph.degree(&resolved) as f64)
        }
    }
}

fn resolve_value(v: &Value, store: &Store, graph: &Graph) -> Value {
    match v {
        Value::FnCall(f) => resolve_fn(f, store, graph),
        other => other.clone(),
    }
}

fn eval_node(expr: &Expr, key: &str, store: &Store, graph: &Graph) -> bool {
    match expr {
        Expr::All => true,
        Expr::Comparison { field, op, value } => {
            let lhs = match resolve_field(field, key, store, graph) {
                Some(v) => v,
                None => return false,
            };
            let rhs = resolve_value(value, store, graph);
            compare(&lhs, *op, &rhs)
        }
        Expr::And(a, b) => {
            eval_node(a, key, store, graph) && eval_node(b, key, store, graph)
        }
        Expr::Or(a, b) => {
            eval_node(a, key, store, graph) || eval_node(b, key, store, graph)
        }
        Expr::Not(e) => !eval_node(e, key, store, graph),
        Expr::Neighbors { .. } => false,
    }
}

fn eval_edge(
    expr: &Expr,
    _source: &str,
    target: &str,
    strength: f32,
    rel_type: RelationType,
    store: &Store,
    graph: &Graph,
) -> bool {
    match expr {
        Expr::All => true,
        Expr::Comparison { field, op, value } => {
            let lhs = match field.as_str() {
                "strength" => Value::Num(strength as f64),
                "rel_type" => Value::Str(rel_type_label(rel_type).to_string()),
                _ => match resolve_field(field, target, store, graph) {
                    Some(v) => v,
                    None => return false,
                },
            };
            let rhs = resolve_value(value, store, graph);
            compare(&lhs, *op, &rhs)
        }
        Expr::And(a, b) => {
            eval_edge(a, _source, target, strength, rel_type, store, graph)
                && eval_edge(b, _source, target, strength, rel_type, store, graph)
        }
        Expr::Or(a, b) => {
            eval_edge(a, _source, target, strength, rel_type, store, graph)
                || eval_edge(b, _source, target, strength, rel_type, store, graph)
        }
        Expr::Not(e) => !eval_edge(e, _source, target, strength, rel_type, store, graph),
        Expr::Neighbors { .. } => false,
    }
}

// -- Query result --

pub struct QueryResult {
    pub key: String,
}

// -- Query executor --

pub fn execute_query(
    store: &Store,
    graph: &Graph,
    query_str: &str,
) -> Result<Vec<QueryResult>, String> {
    let q = query_parser::query(query_str)
        .map_err(|e| format!("Parse error: {}", e))?;

    let mut results = match &q.expr {
        Expr::Neighbors { key, filter } => {
            let resolved = store.resolve_key(key).unwrap_or_else(|_| key.clone());
            let edges = graph.edges_of(&resolved);
            let mut out = Vec::new();
            for edge in edges {
                let include = match filter {
                    Some(f) => eval_edge(
                        f, &resolved, &edge.target,
                        edge.strength, edge.rel_type, store, graph,
                    ),
                    None => true,
                };
                if include {
                    out.push(QueryResult { key: edge.target.clone() });
                }
            }
            out
        }
        _ => {
            let mut out = Vec::new();
            for key in store.nodes.keys() {
                if store.nodes[key].deleted { continue; }
                if eval_node(&q.expr, key, store, graph) {
                    out.push(QueryResult { key: key.clone() });
                }
            }
            out
        }
    };

    // Apply pipeline stages
    let mut has_sort = false;
    for stage in &q.stages {
        match stage {
            Stage::Sort { field, ascending } => {
                has_sort = true;
                let asc = *ascending;
                results.sort_by(|a, b| {
                    let va = resolve_field(field, &a.key, store, graph)
                        .and_then(|v| as_num(&v));
                    let vb = resolve_field(field, &b.key, store, graph)
                        .and_then(|v| as_num(&v));
                    match (va, vb) {
                        (Some(a), Some(b)) => if asc {
                            a.partial_cmp(&b).unwrap_or(std::cmp::Ordering::Equal)
                        } else {
                            b.partial_cmp(&a).unwrap_or(std::cmp::Ordering::Equal)
                        },
                        // String fallback for non-numeric fields
                        _ => {
                            let sa = resolve_field(field, &a.key, store, graph)
                                .map(|v| as_str(&v)).unwrap_or_default();
                            let sb = resolve_field(field, &b.key, store, graph)
                                .map(|v| as_str(&v)).unwrap_or_default();
                            if asc { sa.cmp(&sb) } else { sb.cmp(&sa) }
                        }
                    }
                });
            }
            Stage::Limit(n) => {
                results.truncate(*n);
            }
            Stage::Select(_) | Stage::Count => {} // handled in output
        }
    }

    // Default sort by degree desc if no explicit sort
    if !has_sort {
        results.sort_by(|a, b| {
            let da = graph.degree(&a.key);
            let db = graph.degree(&b.key);
            db.cmp(&da)
        });
    }

    Ok(results)
}

/// Extract the output stages from a parsed query (for cmd_query formatting)
pub fn output_stages(query_str: &str) -> Result<Vec<Stage>, String> {
    let q = query_parser::query(query_str)
        .map_err(|e| format!("Parse error: {}", e))?;
    Ok(q.stages)
}

/// Format a field value for display
pub fn format_field(field: &str, key: &str, store: &Store, graph: &Graph) -> String {
    match resolve_field(field, key, store, graph) {
        Some(Value::Num(n)) => {
            if n == n.floor() && n.abs() < 1e15 {
                format!("{}", n as i64)
            } else {
                format!("{:.3}", n)
            }
        }
        Some(Value::Str(s)) => {
            if field == "content" {
                let truncated: String = s.chars().take(80).collect();
                if s.len() > 80 { format!("{}...", truncated) } else { truncated }
            } else {
                s
            }
        }
        Some(Value::Ident(s)) => s,
        Some(Value::FnCall(_)) => "?".to_string(),
        None => "-".to_string(),
    }
}