poimen-memory/crates/mem-cli/src/query/query_reasoner.rs

//! Query Reasoning (Phase 5.3)
//!
//! Complex question decomposition, multi-hop reasoning, constraint satisfaction,
//! and answer validation.

use std::collections::HashMap;
use sqlx::PgPool;
use serde::{Deserialize, Serialize};
use tracing::{debug, warn};

/// Question type/intent
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum QuestionType {
    /// "What is X?" - Simple fact lookup
    Factual,
    /// "How does A relate to B?" - Relationship query
    Relationship,
    /// "Find all X that satisfy Y" - Set query with constraints
    SetQuery,
    /// "Why is X true?" - Multi-hop reasoning
    Causal,
    /// "Compare A vs B" - Comparative reasoning
    Comparative,
    /// "What are consequences of X?" - Forward chaining
    Consequence,
}

/// Decomposed sub-query
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SubQuery {
    /// Sub-query ID
    pub id: String,
    /// The actual question (natural language)
    pub question: String,
    /// Question type
    pub question_type: QuestionType,
    /// Entity IDs to query
    pub entity_ids: Vec<String>,
    /// Relation types to follow
    pub relation_types: Vec<String>,
    /// Constraints to apply
    pub constraints: Vec<Constraint>,
    /// Expected result type
    pub result_type: ResultType,
}

/// Constraint on results
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Constraint {
    /// Constraint type (e.g., "confidence", "relation_type", "distance")
    pub constraint_type: String,
    /// Operator (e.g., ">=", "==", "in", "not_in")
    pub operator: String,
    /// Value to compare against
    pub value: String,
}

/// Result type
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ResultType {
    /// Single entity
    Entity,
    /// Multiple entities
    Entities,
    /// Relationship/edge
    Edge,
    /// Multiple relationships
    Edges,
    /// Boolean (yes/no)
    Boolean,
    /// Count
    Count,
}

/// Reasoning step result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ReasoningStep {
    /// Step index
    pub step_id: usize,
    /// Sub-query executed
    pub sub_query: SubQuery,
    /// Results from this step
    pub results: Vec<String>,
    /// Confidence in results
    pub confidence: f32,
    /// Constraints satisfied
    pub constraints_satisfied: usize,
    /// Constraints total
    pub constraints_total: usize,
}

/// Final answer with reasoning
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ReasonedAnswer {
    /// Original question
    pub question: String,
    /// Final answer(s)
    pub answers: Vec<String>,
    /// Answer confidence
    pub confidence: f32,
    /// Reasoning steps
    pub reasoning_steps: Vec<ReasoningStep>,
    /// Evidence supporting answer
    pub evidence: Vec<String>,
    /// Explanation
    pub explanation: String,
}

/// Query Reasoner
pub struct QueryReasoner {
    pool: PgPool,
}

impl QueryReasoner {
    pub fn new(pool: PgPool) -> Self {
        QueryReasoner { pool }
    }

    /// Decompose complex question into sub-queries
    pub fn decompose_question(&self, question: &str) -> Result<Vec<SubQuery>, String> {
        if question.is_empty() {
            return Ok(vec![]);
        }

        let question_lower = question.to_lowercase();
        let question_type = self.classify_question(question);

        let mut sub_queries = Vec::new();

        // Detect entities in question (simple heuristic: capitalized words)
        let entities = self.extract_entities_from_question(question);

        // Detect relation keywords
        let relations = self.extract_relations_from_question(question);

        // Create base sub-query
        let base_query = SubQuery {
            id: "sq_1".to_string(),
            question: question.to_string(),
            question_type: question_type.clone(),
            entity_ids: entities.clone(),
            relation_types: relations.clone(),
            constraints: self.extract_constraints_from_question(question),
            result_type: self.infer_result_type(&question_type),
        };

        sub_queries.push(base_query);

        // For complex questions, generate follow-up sub-queries
        if matches!(question_type, QuestionType::Causal | QuestionType::Comparative) {
            // Add explanation sub-query
            sub_queries.push(SubQuery {
                id: "sq_2".to_string(),
                question: format!("Explain the reasoning for: {}", question),
                question_type: QuestionType::Causal,
                entity_ids: entities,
                relation_types: relations,
                constraints: vec![],
                result_type: ResultType::Entities,
            });
        }

        Ok(sub_queries)
    }

    /// Execute reasoning over sub-queries
    pub async fn reason_over_subqueries(
        &self,
        sub_queries: Vec<SubQuery>,
        project_id: &str,
    ) -> Result<ReasonedAnswer, String> {
        let original_question = sub_queries
            .first()
            .map(|q| q.question.clone())
            .unwrap_or_default();

        let mut reasoning_steps = Vec::new();
        let mut all_results = Vec::new();
        let mut total_confidence = 0.0;

        for (idx, sub_query) in sub_queries.iter().enumerate() {
            // Execute sub-query
            let results = self.execute_subquery(sub_query, project_id).await?;

            // Apply constraints
            let filtered_results = self.apply_constraints(&results, &sub_query.constraints);

            let constraint_satisfaction = if sub_query.constraints.is_empty() {
                1.0
            } else {
                (filtered_results.len() as f32 / results.len().max(1) as f32).min(1.0)
            };

            let confidence = 0.9 * constraint_satisfaction;

            reasoning_steps.push(ReasoningStep {
                step_id: idx + 1,
                sub_query: sub_query.clone(),
                results: filtered_results.clone(),
                confidence,
                constraints_satisfied: filtered_results.len(),
                constraints_total: sub_query.constraints.len(),
            });

            all_results.extend(filtered_results);
            total_confidence += confidence;
        }

        let avg_confidence = if reasoning_steps.is_empty() {
            0.0
        } else {
            total_confidence / reasoning_steps.len() as f32
        };

        // Deduplicate results
        let unique_results: Vec<String> = all_results.into_iter().collect::<std::collections::HashSet<_>>().into_iter().collect();

        // Generate explanation
        let explanation = self.generate_explanation(&reasoning_steps, &unique_results);

        Ok(ReasonedAnswer {
            question: original_question,
            answers: unique_results.clone(),
            confidence: avg_confidence,
            reasoning_steps,
            evidence: unique_results.clone(),
            explanation,
        })
    }

    /// Validate answer against constraints
    pub fn validate_answer(
        &self,
        answer: &str,
        constraints: &[Constraint],
    ) -> Result<bool, String> {
        if constraints.is_empty() {
            return Ok(true);
        }

        for constraint in constraints {
            if !self.check_constraint(answer, constraint) {
                return Ok(false);
            }
        }

        Ok(true)
    }

    /// Check if answer satisfies single constraint
    pub fn check_constraint(&self, value: &str, constraint: &Constraint) -> bool {
        match constraint.operator.as_str() {
            "==" | "eq" => value == constraint.value,
            "!=" | "ne" => value != constraint.value,
            "contains" => value.contains(&constraint.value),
            "not_contains" => !value.contains(&constraint.value),
            "in" => {
                let values: Vec<&str> = constraint.value.split(',').map(|s| s.trim()).collect();
                values.contains(&value)
            }
            "not_in" => {
                let values: Vec<&str> = constraint.value.split(',').map(|s| s.trim()).collect();
                !values.contains(&value)
            }
            _ => true,
        }
    }

    // ========== Private Helper Methods ==========

    /// Classify question intent
    fn classify_question(&self, question: &str) -> QuestionType {
        let lower = question.to_lowercase();

        if lower.contains("how does") || lower.contains("how is") {
            QuestionType::Relationship
        } else if lower.contains("why") {
            QuestionType::Causal
        } else if lower.contains("compare") || lower.contains("versus") || lower.contains(" vs ") {
            QuestionType::Comparative
        } else if lower.contains("consequences") || lower.contains("results in") {
            QuestionType::Consequence
        } else if lower.contains("find all") || lower.contains("list all") {
            QuestionType::SetQuery
        } else {
            QuestionType::Factual
        }
    }

    /// Extract entity names from question
    fn extract_entities_from_question(&self, question: &str) -> Vec<String> {
        let words: Vec<&str> = question.split_whitespace().collect();
        let mut entities = Vec::new();

        for word in words {
            if word.chars().next().map_or(false, |c| c.is_uppercase()) && word.len() > 2 {
                entities.push(word.to_string());
            }
        }

        entities.into_iter().collect::<std::collections::HashSet<_>>().into_iter().collect()
    }

    /// Extract relation keywords from question
    fn extract_relations_from_question(&self, question: &str) -> Vec<String> {
        let lower = question.to_lowercase();
        let mut relations = Vec::new();

        if lower.contains("depend") {
            relations.push("depends_on".to_string());
        }
        if lower.contains("relate") {
            relations.push("related_to".to_string());
        }
        if lower.contains("use") {
            relations.push("uses".to_string());
        }
        if lower.contains("contain") {
            relations.push("contains".to_string());
        }
        if lower.contains("require") {
            relations.push("requires".to_string());
        }

        relations
    }

    /// Extract constraints from question
    fn extract_constraints_from_question(&self, question: &str) -> Vec<Constraint> {
        let mut constraints = Vec::new();

        let lower = question.to_lowercase();

        if lower.contains("high confidence") || lower.contains("high reliability") {
            constraints.push(Constraint {
                constraint_type: "confidence".to_string(),
                operator: ">=".to_string(),
                value: "0.8".to_string(),
            });
        }

        if lower.contains("low confidence") {
            constraints.push(Constraint {
                constraint_type: "confidence".to_string(),
                operator: "<".to_string(),
                value: "0.5".to_string(),
            });
        }

        constraints
    }

    /// Infer expected result type
    fn infer_result_type(&self, question_type: &QuestionType) -> ResultType {
        match question_type {
            QuestionType::Factual => ResultType::Entity,
            QuestionType::Relationship => ResultType::Edge,
            QuestionType::SetQuery => ResultType::Entities,
            QuestionType::Causal => ResultType::Entities,
            QuestionType::Comparative => ResultType::Edges,
            QuestionType::Consequence => ResultType::Entities,
        }
    }

    /// Execute single sub-query
    async fn execute_subquery(
        &self,
        _sub_query: &SubQuery,
        _project_id: &str,
    ) -> Result<Vec<String>, String> {
        // Stub: would query database based on sub_query
        Ok(vec![])
    }

    /// Apply constraints to results
    fn apply_constraints(&self, results: &[String], constraints: &[Constraint]) -> Vec<String> {
        if constraints.is_empty() {
            return results.to_vec();
        }

        results
            .iter()
            .filter(|result| {
                constraints.iter().all(|c| self.check_constraint(result, c))
            })
            .cloned()
            .collect()
    }

    /// Generate human-readable explanation
    fn generate_explanation(
        &self,
        steps: &[ReasoningStep],
        answers: &[String],
    ) -> String {
        if steps.is_empty() {
            return "No reasoning steps available".to_string();
        }

        let mut explanation = format!("Found {} answer(s) through {} reasoning step(s): ", answers.len(), steps.len());

        for (idx, step) in steps.iter().enumerate() {
            explanation.push_str(&format!(
                "Step {}: {} (confidence: {:.2}, {} constraints satisfied). ",
                step.step_id,
                step.sub_query.question,
                step.confidence,
                step.constraints_satisfied
            ));
        }

        explanation
    }
}