Vapora/crates/vapora-agents/src/scoring.rs

use vapora_swarm::messages::AgentProfile;

use crate::learning_profile::LearningProfile;

/// Unified agent score combining SwarmCoordinator metrics and learning
/// expertise.
#[derive(Debug, Clone)]
pub struct AgentScore {
    /// Agent identifier
    pub agent_id: String,
    /// Base score from SwarmCoordinator: success_rate / (1 + current_load)
    pub base_score: f64,
    /// Expertise score from LearningProfile for specific task type (0.0-1.0)
    pub expertise_score: f64,
    /// Confidence in expertise score based on execution count (0.0-1.0)
    pub confidence: f64,
    /// Final combined score: 0.3*base + 0.5*expertise + 0.2*confidence
    pub final_score: f64,
    /// Human-readable explanation of scoring breakdown
    pub reasoning: String,
}

/// Service for ranking agents based on learning profiles and swarm metrics.
pub struct AgentScoringService;

impl AgentScoringService {
    /// Rank candidate agents for task assignment using combined scoring.
    ///
    /// Scoring formula:
    /// - base_score = success_rate / (1 + current_load) [from SwarmCoordinator]
    /// - expertise_score = learned success rate for task_type
    /// - confidence = min(1.0, total_executions / 20) [prevents overfitting]
    /// - final_score = 0.3*base + 0.5*expertise + 0.2*confidence
    ///
    /// Returns agents ranked by final_score (highest first).
    pub fn rank_agents(
        candidates: Vec<AgentProfile>,
        task_type: &str,
        learning_profiles: &[(String, LearningProfile)],
    ) -> Vec<AgentScore> {
        let mut scores: Vec<AgentScore> = candidates
            .into_iter()
            .map(|agent| {
                let base_score = agent.success_rate / (1.0 + agent.current_load);

                let (expertise_score, confidence) = learning_profiles
                    .iter()
                    .find(|(id, _)| id == &agent.id)
                    .map(|(_, profile)| {
                        (
                            profile.get_task_type_score(task_type),
                            profile.get_confidence(task_type),
                        )
                    })
                    .unwrap_or((agent.success_rate, 0.0));

                let final_score = 0.3 * base_score + 0.5 * expertise_score + 0.2 * confidence;

                let reasoning = format!(
                    "{}(base={:.2}, load={:.2}, expertise={:.2}, confidence={:.2})",
                    agent.id, base_score, agent.current_load, expertise_score, confidence
                );

                AgentScore {
                    agent_id: agent.id.clone(),
                    base_score,
                    expertise_score,
                    confidence,
                    final_score,
                    reasoning,
                }
            })
            .collect();

        scores.sort_by(|a, b| {
            b.final_score
                .partial_cmp(&a.final_score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        scores
    }

    /// Select best agent from candidates for task assignment.
    /// Returns the top-ranked agent or None if no candidates available.
    pub fn select_best(
        candidates: Vec<AgentProfile>,
        task_type: &str,
        learning_profiles: &[(String, LearningProfile)],
    ) -> Option<AgentScore> {
        Self::rank_agents(candidates, task_type, learning_profiles)
            .into_iter()
            .next()
    }

    /// Calculate blended score prioritizing task-type expertise.
    /// Uses recent_success_rate if available (recency bias from learning
    /// profile).
    pub fn rank_agents_with_recency(
        candidates: Vec<AgentProfile>,
        task_type: &str,
        learning_profiles: &[(String, LearningProfile)],
    ) -> Vec<AgentScore> {
        let mut scores: Vec<AgentScore> = candidates
            .into_iter()
            .map(|agent| {
                let base_score = agent.success_rate / (1.0 + agent.current_load);

                let (expertise_score, confidence) = learning_profiles
                    .iter()
                    .find(|(id, _)| id == &agent.id)
                    .map(|(_, profile)| {
                        // Use recent_success_rate if available (weighted 3x for last 7 days)
                        let recent = profile.get_recent_score(task_type);
                        let conf = profile.get_confidence(task_type);
                        (recent, conf)
                    })
                    .unwrap_or((agent.success_rate, 0.0));

                let final_score = 0.3 * base_score + 0.5 * expertise_score + 0.2 * confidence;

                let reasoning = format!(
                    "{}(recent={:.2}, confidence={:.2})",
                    agent.id, expertise_score, confidence
                );

                AgentScore {
                    agent_id: agent.id.clone(),
                    base_score,
                    expertise_score,
                    confidence,
                    final_score,
                    reasoning,
                }
            })
            .collect();

        scores.sort_by(|a, b| {
            b.final_score
                .partial_cmp(&a.final_score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        scores
    }
}

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

    fn create_mock_agent(id: &str, success_rate: f64, load: f64) -> AgentProfile {
        AgentProfile {
            id: id.to_string(),
            roles: vec![],
            capabilities: vec![],
            current_load: load,
            success_rate,
            availability: true,
        }
    }

    fn create_mock_learning(agent_id: &str, expertise: f64, confidence: f64) -> LearningProfile {
        use crate::learning_profile::TaskTypeExpertise;
        let mut profile = LearningProfile::new(agent_id.to_string());
        let exp = TaskTypeExpertise {
            success_rate: expertise,
            total_executions: (confidence * 20.0) as u32,
            recent_success_rate: expertise,
            avg_duration_ms: 100.0,
            learning_curve: Vec::new(),
            confidence,
        };
        profile.set_task_type_expertise("coding".to_string(), exp);
        profile
    }

    #[test]
    fn test_rank_agents_basic() {
        let candidates = vec![
            create_mock_agent("agent-a", 0.8, 0.4),
            create_mock_agent("agent-b", 0.6, 0.2),
            create_mock_agent("agent-c", 0.9, 0.5),
        ];

        let learning = vec![
            (
                "agent-a".to_string(),
                create_mock_learning("agent-a", 0.85, 0.8),
            ),
            (
                "agent-b".to_string(),
                create_mock_learning("agent-b", 0.70, 0.6),
            ),
            (
                "agent-c".to_string(),
                create_mock_learning("agent-c", 0.75, 0.5),
            ),
        ];

        let ranked = AgentScoringService::rank_agents(candidates, "coding", &learning);
        assert_eq!(ranked.len(), 3);
        // Verify sorted by final_score descending
        for i in 1..ranked.len() {
            assert!(ranked[i - 1].final_score >= ranked[i].final_score);
        }
    }

    #[test]
    fn test_select_best() {
        let candidates = vec![
            create_mock_agent("agent-a", 0.8, 0.4),
            create_mock_agent("agent-b", 0.6, 0.2),
        ];

        let learning = vec![
            (
                "agent-a".to_string(),
                create_mock_learning("agent-a", 0.85, 0.8),
            ),
            (
                "agent-b".to_string(),
                create_mock_learning("agent-b", 0.70, 0.6),
            ),
        ];

        let best = AgentScoringService::select_best(candidates, "coding", &learning);
        assert!(best.is_some());
        assert_eq!(best.unwrap().agent_id, "agent-a");
    }

    #[test]
    fn test_rank_agents_no_learning_data() {
        let candidates = vec![
            create_mock_agent("agent-a", 0.8, 0.4),
            create_mock_agent("agent-b", 0.6, 0.2),
        ];

        let ranked = AgentScoringService::rank_agents(candidates, "coding", &[]);
        assert_eq!(ranked.len(), 2);
        // Should still rank by base score when no learning data
        assert!(ranked[0].final_score > 0.0);
    }

    #[test]
    fn test_recency_bias_scoring() {
        let candidates = vec![
            create_mock_agent("agent-a", 0.5, 0.3),
            create_mock_agent("agent-b", 0.5, 0.3),
        ];

        let mut learning_a = LearningProfile::new("agent-a".to_string());
        use crate::learning_profile::TaskTypeExpertise;
        learning_a.set_task_type_expertise(
            "coding".to_string(),
            TaskTypeExpertise {
                success_rate: 0.7,
                total_executions: 20,
                recent_success_rate: 0.95, // Recent success much higher
                avg_duration_ms: 100.0,
                learning_curve: Vec::new(),
                confidence: 1.0,
            },
        );

        let learning = vec![("agent-a".to_string(), learning_a)];

        let ranked = AgentScoringService::rank_agents_with_recency(candidates, "coding", &learning);
        assert_eq!(ranked.len(), 2);
        // agent-a should rank higher due to recent success
        assert_eq!(ranked[0].agent_id, "agent-a");
    }

    #[test]
    fn test_confidence_weights_low_sample_count() {
        let candidates = vec![
            create_mock_agent("agent-a", 0.9, 0.0), // High success but...
            create_mock_agent("agent-b", 0.8, 0.0), // Moderate success
        ];

        let learning = vec![
            (
                "agent-a".to_string(),
                create_mock_learning("agent-a", 0.9, 0.05),
            ), // Low confidence
            (
                "agent-b".to_string(),
                create_mock_learning("agent-b", 0.8, 0.95),
            ), // High confidence
        ];

        let ranked = AgentScoringService::rank_agents(candidates, "coding", &learning);
        // agent-b should rank higher due to higher confidence despite lower expertise
        assert_eq!(ranked[0].agent_id, "agent-b");
    }

    #[test]
    fn test_empty_candidates() {
        let ranked = AgentScoringService::rank_agents(Vec::new(), "coding", &[]);
        assert_eq!(ranked.len(), 0);
    }
}