Vapora/crates/vapora-knowledge-graph/src/reasoning.rs

use crate::models::*;
use std::collections::HashMap;

/// Reasoning engine for inferring insights from execution history
pub struct ReasoningEngine;

impl ReasoningEngine {
    /// Analyze execution records for patterns
    pub fn find_patterns(records: &[ExecutionRecord]) -> Vec<String> {
        let mut patterns = Vec::new();

        // Pattern 1: Task type failure rate
        let mut task_failure_rates: HashMap<String, (u32, u32)> = HashMap::new();
        for record in records {
            task_failure_rates
                .entry(record.task_type.clone())
                .or_insert((0, 0));

            let entry = task_failure_rates.get_mut(&record.task_type).unwrap();
            entry.0 += 1;
            if !record.success {
                entry.1 += 1;
            }
        }

        for (task_type, (total, failures)) in task_failure_rates.iter() {
            let failure_rate = *failures as f64 / *total as f64;
            if failure_rate > 0.3 {
                patterns.push(format!(
                    "High failure rate detected for '{}': {:.1}%",
                    task_type,
                    failure_rate * 100.0
                ));
            }
        }

        // Pattern 2: Duration anomalies
        if records.len() > 3 {
            let mut durations: Vec<u64> = records.iter().map(|r| r.duration_ms).collect();
            durations.sort_unstable();
            let median = durations[durations.len() / 2];
            let avg = durations.iter().sum::<u64>() as f64 / durations.len() as f64;

            if avg > median as f64 * 2.0 {
                patterns.push(format!(
                    "High latency variance detected: avg={:.0}ms, median={}ms",
                    avg, median
                ));
            }
        }

        // Pattern 3: Agent expertise misalignment
        let mut agent_task_types: HashMap<String, HashMap<String, u32>> = HashMap::new();
        for record in records {
            agent_task_types
                .entry(record.agent_id.clone())
                .or_default()
                .entry(record.task_type.clone())
                .and_modify(|c| *c += 1)
                .or_insert(1);
        }

        for (agent_id, task_map) in agent_task_types {
            if task_map.len() > 5 {
                patterns.push(format!(
                    "Agent '{}' is assigned to {} different task types (possible overallocation)",
                    agent_id,
                    task_map.len()
                ));
            }
        }

        patterns
    }

    /// Predict success probability for a task
    pub fn predict_success(
        task_type: &str,
        similar_records: &[ExecutionRecord],
    ) -> (f64, String) {
        if similar_records.is_empty() {
            return (0.5, "No historical data available".to_string());
        }

        let success_count = similar_records.iter().filter(|r| r.success).count() as f64;
        let success_rate = success_count / similar_records.len() as f64;

        let reasoning = if success_rate > 0.9 {
            format!(
                "High success probability based on {} successful similar tasks",
                success_count as usize
            )
        } else if success_rate > 0.7 {
            "Moderate success probability with some historical challenges".to_string()
        } else if success_rate > 0.5 {
            "Below-average success rate - recommend expert agent".to_string()
        } else {
            format!(
                "Critical: {} only {} successful of {} similar tasks",
                task_type,
                success_count as usize,
                similar_records.len()
            )
        };

        (success_rate, reasoning)
    }

    /// Estimate task duration
    pub fn estimate_duration(similar_records: &[ExecutionRecord]) -> (u64, String) {
        if similar_records.is_empty() {
            return (300_000, "No historical data - using default 5 minutes".to_string());
        }

        let mut durations: Vec<u64> = similar_records.iter().map(|r| r.duration_ms).collect();
        durations.sort_unstable();

        let median = durations[durations.len() / 2];
        let avg = durations.iter().sum::<u64>() / durations.len() as u64;
        let max = durations[durations.len() - 1];

        let estimate = (avg + median) / 2;
        let reasoning = if max > estimate * 3 {
            format!(
                "High variance in execution time: avg={}ms, max={}ms",
                avg, max
            )
        } else {
            format!("Based on {} similar tasks: avg={}ms", durations.len(), avg)
        };

        (estimate, reasoning)
    }

    /// Recommend best agent for a task
    pub fn recommend_agent(
        task_type: &str,
        agent_profiles: &[AgentProfile],
    ) -> Option<(String, String)> {
        agent_profiles
            .iter()
            .filter(|profile| {
                profile.primary_task_types.contains(&task_type.to_string())
                    || profile.expertise_score > 75.0
            })
            .max_by(|a, b| {
                a.expertise_score
                    .partial_cmp(&b.expertise_score)
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .map(|profile| {
                let reasoning = format!(
                    "Agent '{}' has {:.1}% expertise in {} tasks",
                    profile.agent_id,
                    profile.expertise_score,
                    profile.primary_task_types.join(", ")
                );
                (profile.agent_id.clone(), reasoning)
            })
    }

    /// Chain reasoning - find root cause chains
    pub fn find_root_cause_chain(records: &[ExecutionRecord]) -> Vec<Vec<String>> {
        let mut chains: Vec<Vec<String>> = Vec::new();
        let mut visited = std::collections::HashSet::new();

        for record in records.iter().filter(|r| !r.success && r.root_cause.is_some()) {
            if visited.contains(&record.id) {
                continue;
            }

            let mut chain = vec![record.root_cause.clone().unwrap()];
            visited.insert(record.id.clone());

            // Find similar errors in history
            for other in records
                .iter()
                .filter(|r| !visited.contains(&r.id) && r.solution.is_some())
            {
                if let Some(ref root) = other.root_cause {
                    if root.contains(chain.last().unwrap()) {
                        chain.push(other.solution.clone().unwrap());
                        visited.insert(other.id.clone());
                        break;
                    }
                }
            }

            if chain.len() > 1 {
                chains.push(chain);
            }
        }

        chains
    }

    /// Generate actionable insights
    pub fn generate_insights(
        stats: &GraphStatistics,
        patterns: &[String],
        profiles: &[AgentProfile],
    ) -> Vec<String> {
        let mut insights = Vec::new();

        // Insight 1: Overall system health
        if stats.success_rate > 0.95 {
            insights.push("✓ System is very reliable (>95% success rate)".to_string());
        } else if stats.success_rate < 0.75 {
            insights.push(format!(
                "⚠ System reliability is concerning ({:.1}% success rate)",
                stats.success_rate * 100.0
            ));
        }

        // Insight 2: Performance metrics
        if stats.avg_duration_ms > 60_000.0 {
            insights.push(format!(
                "⚠ Average task duration is high ({:.0}ms)",
                stats.avg_duration_ms
            ));
        }

        // Insight 3: Coverage
        insights.push(format!(
            "ℹ Knowledge graph contains {} execution records across {} agents",
            stats.total_records, stats.distinct_agents
        ));

        // Insight 4: Agent expertise distribution
        let avg_expertise = profiles.iter().map(|p| p.expertise_score).sum::<f64>()
            / profiles.len().max(1) as f64;
        if avg_expertise < 70.0 {
            insights.push(format!(
                "⚠ Average agent expertise is below target ({:.0}%)",
                avg_expertise
            ));
        }

        // Insight 5: Patterns
        insights.extend(patterns.iter().cloned());

        insights
    }
}

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

    #[test]
    fn test_predict_success() {
        let records = vec![
            ExecutionRecord {
                id: "1".to_string(),
                task_id: "t1".to_string(),
                agent_id: "a1".to_string(),
                task_type: "dev".to_string(),
                description: "test".to_string(),
                root_cause: None,
                solution: None,
                duration_ms: 1000,
                input_tokens: 100,
                output_tokens: 50,
                success: true,
                error: None,
                timestamp: Utc::now(),
            },
            ExecutionRecord {
                id: "2".to_string(),
                task_id: "t2".to_string(),
                agent_id: "a1".to_string(),
                task_type: "dev".to_string(),
                description: "test".to_string(),
                root_cause: None,
                solution: None,
                duration_ms: 1000,
                input_tokens: 100,
                output_tokens: 50,
                success: true,
                error: None,
                timestamp: Utc::now(),
            },
        ];

        let (prob, _) = ReasoningEngine::predict_success("dev", &records);
        assert_eq!(prob, 1.0);
    }

    #[test]
    fn test_estimate_duration() {
        let records = vec![
            ExecutionRecord {
                id: "1".to_string(),
                task_id: "t1".to_string(),
                agent_id: "a1".to_string(),
                task_type: "dev".to_string(),
                description: "test".to_string(),
                root_cause: None,
                solution: None,
                duration_ms: 1000,
                input_tokens: 100,
                output_tokens: 50,
                success: true,
                error: None,
                timestamp: Utc::now(),
            },
        ];

        let (duration, _) = ReasoningEngine::estimate_duration(&records);
        assert_eq!(duration, 1000);
    }
}