poimen-memory/crates/mem-core/src/query.rs

use crate::domain::{ProjectId, QueryId};
use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;

/// A single standing query.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Query {
    pub id: String,
    pub question: String,
    #[serde(default)]
    pub exit_gate: bool,
}

/// Synthesis query (optional, for L2).
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct SynthesisQuery {
    pub question: String,
    #[serde(default)]
    pub exit_gate: bool,
}

/// Defaults applied to queries.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Defaults {
    #[serde(default = "default_memory_budget")]
    pub memory_budget: u32,
    #[serde(default = "default_chunk_tokens")]
    pub chunk_tokens: u32,
    #[serde(default)]
    pub exit_gate: bool,
}

fn default_memory_budget() -> u32 {
    1024
}

fn default_chunk_tokens() -> u32 {
    5000
}

impl Default for Defaults {
    fn default() -> Self {
        Self {
            memory_budget: default_memory_budget(),
            chunk_tokens: default_chunk_tokens(),
            exit_gate: false,
        }
    }
}

/// Complete set of queries for a project.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QuerySet {
    pub project: String,
    pub roots: Vec<String>,
    pub sources: Vec<String>,
    pub queries: Vec<Query>,
    #[serde(default)]
    pub synthesis: Option<SynthesisQuery>,
    #[serde(default)]
    pub defaults: Defaults,
}

/// Load error with context.
#[derive(Debug, Clone)]
pub struct QueryLoadError {
    pub file: String,
    pub query_id: Option<String>,
    pub field: Option<String>,
    pub message: String,
}

impl std::fmt::Display for QueryLoadError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match (&self.query_id, &self.field) {
            (Some(id), Some(field)) => {
                write!(f, "{}: query '{}', field '{}': {}", self.file, id, field, self.message)
            }
            (Some(id), None) => {
                write!(f, "{}: query '{}': {}", self.file, id, self.message)
            }
            (None, Some(field)) => {
                write!(f, "{}: field '{}': {}", self.file, field, self.message)
            }
            (None, None) => {
                write!(f, "{}: {}", self.file, self.message)
            }
        }
    }
}

impl std::error::Error for QueryLoadError {}

/// Valid charset for query ids: lowercase, digits, hyphens only.
fn is_valid_query_id(id: &str) -> bool {
    !id.is_empty() && id.chars().all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-')
}

impl QuerySet {
    /// Load and validate a query set from a YAML file.
    pub fn load<P: AsRef<Path>>(path: P) -> Result<Self> {
        let path = path.as_ref();
        let filename = path.to_string_lossy().to_string();
        let contents = std::fs::read_to_string(path)?;
        
        // Parse YAML
        let mut set: QuerySet = serde_yaml::from_str(&contents)
            .map_err(|e| anyhow!("Failed to parse {}: {}", filename, e))?;

        // Validate project
        if set.project.trim().is_empty() {
            return Err(anyhow!(QueryLoadError {
                file: filename,
                query_id: None,
                field: Some("project".to_string()),
                message: "project field is required and cannot be empty".to_string(),
            }));
        }

        // Validate at least one query
        if set.queries.is_empty() {
            return Err(anyhow!(QueryLoadError {
                file: filename,
                query_id: None,
                field: Some("queries".to_string()),
                message: "at least one query is required".to_string(),
            }));
        }

        // Validate each query
        let mut seen_ids = std::collections::HashSet::new();
        for query in &mut set.queries {
            // Check ID is not empty
            if query.id.trim().is_empty() {
                return Err(anyhow!(QueryLoadError {
                    file: filename,
                    query_id: None,
                    field: Some("id".to_string()),
                    message: "query id cannot be empty".to_string(),
                }));
            }

            // Check ID charset
            if !is_valid_query_id(&query.id) {
                return Err(anyhow!(QueryLoadError {
                    file: filename.clone(),
                    query_id: Some(query.id.clone()),
                    field: Some("id".to_string()),
                    message: format!(
                        "query id '{}' must match [a-z0-9-]+ (it becomes a filename)",
                        query.id
                    ),
                }));
            }

            // Check for duplicate IDs
            if seen_ids.contains(&query.id) {
                return Err(anyhow!(QueryLoadError {
                    file: filename,
                    query_id: Some(query.id.clone()),
                    field: Some("id".to_string()),
                    message: format!("duplicate query id '{}'", query.id),
                }));
            }
            seen_ids.insert(query.id.clone());

            // Check question is not empty
            if query.question.trim().is_empty() {
                return Err(anyhow!(QueryLoadError {
                    file: filename,
                    query_id: Some(query.id.clone()),
                    field: Some("question".to_string()),
                    message: "question cannot be empty".to_string(),
                }));
            }

            // Apply defaults if exit_gate not set
            // (defaults already applied via serde default)
        }

        // Validate synthesis if present
        if let Some(ref synthesis) = set.synthesis {
            if synthesis.question.trim().is_empty() {
                return Err(anyhow!(QueryLoadError {
                    file: filename,
                    query_id: None,
                    field: Some("synthesis.question".to_string()),
                    message: "synthesis question cannot be empty".to_string(),
                }));
            }
        }

        // Validate defaults
        if set.defaults.memory_budget == 0 {
            return Err(anyhow!(QueryLoadError {
                file: filename,
                query_id: None,
                field: Some("defaults.memory_budget".to_string()),
                message: "memory_budget must be greater than 0".to_string(),
            }));
        }

        if set.defaults.chunk_tokens == 0 {
            return Err(anyhow!(QueryLoadError {
                file: filename,
                query_id: None,
                field: Some("defaults.chunk_tokens".to_string()),
                message: "chunk_tokens must be greater than 0".to_string(),
            }));
        }

        Ok(set)
    }

    /// Get a query by ID.
    pub fn query(&self, id: &str) -> Option<&Query> {
        self.queries.iter().find(|q| q.id == id)
    }
}

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

    #[test]
    fn test_valid_query_id() {
        assert!(is_valid_query_id("architecture-decisions"));
        assert!(is_valid_query_id("infra-root-causes"));
        assert!(is_valid_query_id("id123"));
        assert!(is_valid_query_id("a"));
        assert!(is_valid_query_id("a-b-c-123"));

        assert!(!is_valid_query_id(""));
        assert!(!is_valid_query_id("infra/root-causes"));
        assert!(!is_valid_query_id("UPPERCASE"));
        assert!(!is_valid_query_id("with space"));
        assert!(!is_valid_query_id("with_underscore"));
    }

    #[test]
    fn test_defaults() {
        let defaults = Defaults::default();
        assert_eq!(defaults.memory_budget, 1024);
        assert_eq!(defaults.chunk_tokens, 5000);
        assert!(!defaults.exit_gate);
    }
}