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
    }
}

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

    fn create_reasoner_mock() -> QueryReasoner {
        let pool = sqlx::postgres::PgPoolOptions::new()
            .max_connections(1)
            .build_lazy();
        QueryReasoner::new(pool)
    }

    #[test]
    fn test_question_type_factual() {
        let reasoner = create_reasoner_mock();
        let qt = reasoner.classify_question("What is Kubernetes?");
        assert_eq!(qt, QuestionType::Factual);
    }

    #[test]
    fn test_question_type_relationship() {
        let reasoner = create_reasoner_mock();
        let qt = reasoner.classify_question("How does Docker relate to Kubernetes?");
        assert_eq!(qt, QuestionType::Relationship);
    }

    #[test]
    fn test_question_type_causal() {
        let reasoner = create_reasoner_mock();
        let qt = reasoner.classify_question("Why is Kubernetes essential?");
        assert_eq!(qt, QuestionType::Causal);
    }

    #[test]
    fn test_question_type_comparative() {
        let reasoner = create_reasoner_mock();
        let qt = reasoner.classify_question("Compare Docker versus Kubernetes");
        assert_eq!(qt, QuestionType::Comparative);
    }

    #[test]
    fn test_question_type_set_query() {
        let reasoner = create_reasoner_mock();
        let qt = reasoner.classify_question("Find all containerization tools");
        assert_eq!(qt, QuestionType::SetQuery);
    }

    #[test]
    fn test_question_type_consequence() {
        let reasoner = create_reasoner_mock();
        let qt = reasoner.classify_question("What are the consequences of using Kubernetes?");
        assert_eq!(qt, QuestionType::Consequence);
    }

    #[test]
    fn test_extract_entities() {
        let reasoner = create_reasoner_mock();
        let entities = reasoner.extract_entities_from_question("How does Kubernetes work with Docker?");
        assert!(entities.contains(&"Kubernetes".to_string()));
        assert!(entities.contains(&"Docker".to_string()));
    }

    #[test]
    fn test_extract_relations_depends() {
        let reasoner = create_reasoner_mock();
        let relations = reasoner.extract_relations_from_question("What does Kubernetes depend on?");
        assert!(relations.contains(&"depends_on".to_string()));
    }

    #[test]
    fn test_extract_relations_uses() {
        let reasoner = create_reasoner_mock();
        let relations = reasoner.extract_relations_from_question("Kubernetes uses containers");
        assert!(relations.contains(&"uses".to_string()));
    }

    #[test]
    fn test_extract_constraints_high_confidence() {
        let reasoner = create_reasoner_mock();
        let constraints = reasoner.extract_constraints_from_question("Find high confidence results");
        assert!(constraints.iter().any(|c| c.constraint_type == "confidence"));
    }

    #[test]
    fn test_constraint_equals() {
        let reasoner = create_reasoner_mock();
        let constraint = Constraint {
            constraint_type: "type".to_string(),
            operator: "==".to_string(),
            value: "entity".to_string(),
        };
        assert!(reasoner.check_constraint("entity", &constraint));
        assert!(!reasoner.check_constraint("edge", &constraint));
    }

    #[test]
    fn test_constraint_in() {
        let reasoner = create_reasoner_mock();
        let constraint = Constraint {
            constraint_type: "type".to_string(),
            operator: "in".to_string(),
            value: "entity,edge,fact".to_string(),
        };
        assert!(reasoner.check_constraint("entity", &constraint));
        assert!(reasoner.check_constraint("edge", &constraint));
        assert!(!reasoner.check_constraint("other", &constraint));
    }

    #[test]
    fn test_constraint_contains() {
        let reasoner = create_reasoner_mock();
        let constraint = Constraint {
            constraint_type: "text".to_string(),
            operator: "contains".to_string(),
            value: "test".to_string(),
        };
        assert!(reasoner.check_constraint("this is a test", &constraint));
        assert!(!reasoner.check_constraint("this is not it", &constraint));
    }

    #[test]
    fn test_subquery_structure() {
        let sq = SubQuery {
            id: "sq1".to_string(),
            question: "What is X?".to_string(),
            question_type: QuestionType::Factual,
            entity_ids: vec!["e1".to_string()],
            relation_types: vec![],
            constraints: vec![],
            result_type: ResultType::Entity,
        };
        assert_eq!(sq.question_type, QuestionType::Factual);
    }

    #[test]
    fn test_reasoning_step_structure() {
        let step = ReasoningStep {
            step_id: 1,
            sub_query: SubQuery {
                id: "sq1".to_string(),
                question: "Test".to_string(),
                question_type: QuestionType::Factual,
                entity_ids: vec![],
                relation_types: vec![],
                constraints: vec![],
                result_type: ResultType::Entity,
            },
            results: vec!["answer1".to_string()],
            confidence: 0.9,
            constraints_satisfied: 1,
            constraints_total: 1,
        };
        assert_eq!(step.step_id, 1);
        assert_eq!(step.confidence, 0.9);
    }

    #[test]
    fn test_reasoned_answer_structure() {
        let answer = ReasonedAnswer {
            question: "Test question".to_string(),
            answers: vec!["answer1".to_string()],
            confidence: 0.9,
            reasoning_steps: vec![],
            evidence: vec![],
            explanation: "Explanation".to_string(),
        };
        assert_eq!(answer.answers.len(), 1);
    }

    #[test]
    fn test_decompose_empty_question() {
        let reasoner = create_reasoner_mock();
        let result = reasoner.decompose_question("").unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_decompose_simple_question() {
        let reasoner = create_reasoner_mock();
        let result = reasoner.decompose_question("What is Kubernetes?").unwrap();
        assert!(!result.is_empty());
        assert_eq!(result[0].question_type, QuestionType::Factual);
    }

    #[test]
    fn test_decompose_complex_question() {
        let reasoner = create_reasoner_mock();
        let result = reasoner.decompose_question("Why is Kubernetes important?").unwrap();
        assert!(result.len() >= 1);
    }

    #[test]
    fn test_infer_result_type_factual() {
        let reasoner = create_reasoner_mock();
        let rt = reasoner.infer_result_type(&QuestionType::Factual);
        assert_eq!(rt, ResultType::Entity);
    }

    #[test]
    fn test_infer_result_type_set_query() {
        let reasoner = create_reasoner_mock();
        let rt = reasoner.infer_result_type(&QuestionType::SetQuery);
        assert_eq!(rt, ResultType::Entities);
    }

    #[test]
    fn test_constraint_serialization() {
        let constraint = Constraint {
            constraint_type: "test".to_string(),
            operator: "==".to_string(),
            value: "val".to_string(),
        };
        let json = serde_json::to_string(&constraint).unwrap();
        assert!(json.contains("test"));
    }

    #[test]
    fn test_subquery_serialization() {
        let sq = SubQuery {
            id: "sq1".to_string(),
            question: "Test?".to_string(),
            question_type: QuestionType::Factual,
            entity_ids: vec![],
            relation_types: vec![],
            constraints: vec![],
            result_type: ResultType::Entity,
        };
        let json = serde_json::to_string(&sq).unwrap();
        assert!(json.contains("Test?"));
    }

    #[test]
    fn test_validate_answer_no_constraints() {
        let reasoner = create_reasoner_mock();
        let valid = reasoner.validate_answer("answer", &[]).unwrap();
        assert!(valid);
    }

    #[test]
    fn test_validate_answer_with_constraint() {
        let reasoner = create_reasoner_mock();
        let constraint = Constraint {
            constraint_type: "type".to_string(),
            operator: "==".to_string(),
            value: "entity".to_string(),
        };
        let valid = reasoner.validate_answer("entity", &[constraint]).unwrap();
        assert!(valid);
    }

    #[test]
    fn test_apply_constraints_empty() {
        let reasoner = create_reasoner_mock();
        let results = vec!["r1".to_string(), "r2".to_string()];
        let filtered = reasoner.apply_constraints(&results, &[]);
        assert_eq!(filtered.len(), 2);
    }

    #[test]
    fn test_apply_constraints_filter() {
        let reasoner = create_reasoner_mock();
        let results = vec!["entity".to_string(), "edge".to_string()];
        let constraint = Constraint {
            constraint_type: "type".to_string(),
            operator: "==".to_string(),
            value: "entity".to_string(),
        };
        let filtered = reasoner.apply_constraints(&results, &[constraint]);
        assert_eq!(filtered.len(), 1);
        assert_eq!(filtered[0], "entity");
    }

    #[test]
    fn test_generate_explanation() {
        let reasoner = create_reasoner_mock();
        let step = ReasoningStep {
            step_id: 1,
            sub_query: SubQuery {
                id: "sq1".to_string(),
                question: "Test".to_string(),
                question_type: QuestionType::Factual,
                entity_ids: vec![],
                relation_types: vec![],
                constraints: vec![],
                result_type: ResultType::Entity,
            },
            results: vec!["ans".to_string()],
            confidence: 0.9,
            constraints_satisfied: 0,
            constraints_total: 0,
        };
        let expl = reasoner.generate_explanation(&[step], &["ans".to_string()]);
        assert!(expl.contains("reasoning"));
    }
}