Vapora/crates/vapora-agents/tests/learning_integration_test.rs

use chrono::{Duration, Utc};
use vapora_agents::{
    ExecutionData, ProfileAdapter, TaskTypeExpertise,
    AgentScoringService,
};
use vapora_swarm::messages::AgentProfile;

#[test]
fn test_end_to_end_learning_flow() {
    // Simulate historical executions for agent
    let now = Utc::now();
    let executions: Vec<ExecutionData> = (0..20)
        .map(|i| ExecutionData {
            timestamp: now - Duration::days(i),
            duration_ms: 100 + (i as u64 * 10),
            success: i < 18, // 18 successes out of 20 = 90%
        })
        .collect();

    // Calculate expertise from executions
    let expertise = TaskTypeExpertise::from_executions(executions, "coding");
    assert!((expertise.success_rate - 0.9).abs() < 0.01);
    assert_eq!(expertise.total_executions, 20);

    // Create learning profile for agent
    let mut profile = ProfileAdapter::create_learning_profile("agent-1".to_string());

    // Add expertise to profile
    profile = ProfileAdapter::add_task_type_expertise(profile, "coding".to_string(), expertise);

    // Verify expertise is stored
    assert_eq!(profile.get_task_type_score("coding"), 0.9);
    assert!(profile.get_confidence("coding") > 0.9); // 20/20 is high confidence
}

#[test]
fn test_learning_profile_improves_over_time() {
    let now = Utc::now();

    // Initial executions: 50% success
    let initial_execs: Vec<ExecutionData> = (0..10)
        .map(|i| ExecutionData {
            timestamp: now - Duration::days(i * 2),
            duration_ms: 100,
            success: i < 5,
        })
        .collect();

    let mut initial_expertise = TaskTypeExpertise::from_executions(initial_execs, "coding");
    assert!((initial_expertise.success_rate - 0.5).abs() < 0.01);

    // New successful execution
    let new_exec = ExecutionData {
        timestamp: now,
        duration_ms: 100,
        success: true,
    };
    initial_expertise.update_with_execution(&new_exec);

    // Expertise should improve
    assert!(initial_expertise.success_rate > 0.5);
    assert_eq!(initial_expertise.total_executions, 11);
}

#[test]
fn test_agent_scoring_with_learning() {
    // Create candidate agents
    let candidates = vec![
        AgentProfile {
            id: "agent-a".to_string(),
            roles: vec!["developer".to_string()],
            capabilities: vec!["coding".to_string()],
            current_load: 0.3,
            success_rate: 0.8,
            availability: true,
        },
        AgentProfile {
            id: "agent-b".to_string(),
            roles: vec!["developer".to_string()],
            capabilities: vec!["coding".to_string()],
            current_load: 0.1,
            success_rate: 0.8,
            availability: true,
        },
    ];

    // Create learning profiles
    let mut profile_a = ProfileAdapter::create_learning_profile("agent-a".to_string());
    profile_a = ProfileAdapter::add_task_type_expertise(
        profile_a,
        "coding".to_string(),
        TaskTypeExpertise {
            success_rate: 0.95,
            total_executions: 50,
            recent_success_rate: 0.95,
            avg_duration_ms: 100.0,
            learning_curve: Vec::new(),
            confidence: 1.0,
        },
    );

    let mut profile_b = ProfileAdapter::create_learning_profile("agent-b".to_string());
    profile_b = ProfileAdapter::add_task_type_expertise(
        profile_b,
        "coding".to_string(),
        TaskTypeExpertise {
            success_rate: 0.70,
            total_executions: 30,
            recent_success_rate: 0.70,
            avg_duration_ms: 120.0,
            learning_curve: Vec::new(),
            confidence: 1.0,
        },
    );

    let learning_profiles = vec![
        ("agent-a".to_string(), profile_a),
        ("agent-b".to_string(), profile_b),
    ];

    // Score agents
    let ranked = AgentScoringService::rank_agents(candidates, "coding", &learning_profiles);
    assert_eq!(ranked.len(), 2);

    // agent-a should rank higher due to superior expertise despite higher load
    assert_eq!(ranked[0].agent_id, "agent-a");
    assert!(ranked[0].final_score > ranked[1].final_score);
}

#[test]
fn test_recency_bias_affects_ranking() {
    let candidates = vec![
        AgentProfile {
            id: "agent-x".to_string(),
            roles: vec!["developer".to_string()],
            capabilities: vec!["coding".to_string()],
            current_load: 0.3,
            success_rate: 0.8,
            availability: true,
        },
        AgentProfile {
            id: "agent-y".to_string(),
            roles: vec!["developer".to_string()],
            capabilities: vec!["coding".to_string()],
            current_load: 0.3,
            success_rate: 0.8,
            availability: true,
        },
    ];

    // agent-x has high overall success but recent failures
    let mut profile_x = ProfileAdapter::create_learning_profile("agent-x".to_string());
    profile_x = ProfileAdapter::add_task_type_expertise(
        profile_x,
        "coding".to_string(),
        TaskTypeExpertise {
            success_rate: 0.85,
            total_executions: 40,
            recent_success_rate: 0.60, // Recent poor performance
            avg_duration_ms: 100.0,
            learning_curve: Vec::new(),
            confidence: 1.0,
        },
    );

    // agent-y has consistent good recent performance
    let mut profile_y = ProfileAdapter::create_learning_profile("agent-y".to_string());
    profile_y = ProfileAdapter::add_task_type_expertise(
        profile_y,
        "coding".to_string(),
        TaskTypeExpertise {
            success_rate: 0.80,
            total_executions: 30,
            recent_success_rate: 0.90, // Recent strong performance
            avg_duration_ms: 110.0,
            learning_curve: Vec::new(),
            confidence: 1.0,
        },
    );

    let learning_profiles = vec![
        ("agent-x".to_string(), profile_x),
        ("agent-y".to_string(), profile_y),
    ];

    // Rank with recency bias
    let ranked = AgentScoringService::rank_agents_with_recency(candidates, "coding", &learning_profiles);
    assert_eq!(ranked.len(), 2);

    // agent-y should rank higher due to recent success despite lower overall rate
    assert_eq!(ranked[0].agent_id, "agent-y");
}

#[test]
fn test_confidence_prevents_overfitting() {
    let candidates = vec![
        AgentProfile {
            id: "agent-new".to_string(),
            roles: vec!["developer".to_string()],
            capabilities: vec!["coding".to_string()],
            current_load: 0.0,
            success_rate: 0.8,
            availability: true,
        },
        AgentProfile {
            id: "agent-exp".to_string(),
            roles: vec!["developer".to_string()],
            capabilities: vec!["coding".to_string()],
            current_load: 0.0,
            success_rate: 0.8,
            availability: true,
        },
    ];

    // agent-new: High expertise but low confidence (few samples)
    let mut profile_new = ProfileAdapter::create_learning_profile("agent-new".to_string());
    profile_new = ProfileAdapter::add_task_type_expertise(
        profile_new,
        "coding".to_string(),
        TaskTypeExpertise {
            success_rate: 1.0, // Perfect so far
            total_executions: 2,
            recent_success_rate: 1.0,
            avg_duration_ms: 100.0,
            learning_curve: Vec::new(),
            confidence: 0.1, // Low confidence - only 2/20 executions
        },
    );

    // agent-exp: Slightly lower expertise but high confidence
    let mut profile_exp = ProfileAdapter::create_learning_profile("agent-exp".to_string());
    profile_exp = ProfileAdapter::add_task_type_expertise(
        profile_exp,
        "coding".to_string(),
        TaskTypeExpertise {
            success_rate: 0.95,
            total_executions: 50,
            recent_success_rate: 0.95,
            avg_duration_ms: 100.0,
            learning_curve: Vec::new(),
            confidence: 1.0,
        },
    );

    let learning_profiles = vec![
        ("agent-new".to_string(), profile_new),
        ("agent-exp".to_string(), profile_exp),
    ];

    let ranked = AgentScoringService::rank_agents(candidates, "coding", &learning_profiles);

    // agent-exp should rank higher despite slightly lower expertise due to confidence weighting
    assert_eq!(ranked[0].agent_id, "agent-exp");
}

#[test]
fn test_multiple_task_types_independent() {
    let mut profile = ProfileAdapter::create_learning_profile("agent-1".to_string());

    // Agent excels at coding
    let coding_exp = TaskTypeExpertise {
        success_rate: 0.95,
        total_executions: 30,
        recent_success_rate: 0.95,
        avg_duration_ms: 100.0,
        learning_curve: Vec::new(),
        confidence: 1.0,
    };

    // Agent struggles with documentation
    let docs_exp = TaskTypeExpertise {
        success_rate: 0.60,
        total_executions: 20,
        recent_success_rate: 0.65,
        avg_duration_ms: 250.0,
        learning_curve: Vec::new(),
        confidence: 1.0,
    };

    profile = ProfileAdapter::add_task_type_expertise(profile, "coding".to_string(), coding_exp);
    profile = ProfileAdapter::add_task_type_expertise(profile, "documentation".to_string(), docs_exp);

    // Verify independence
    assert_eq!(profile.get_task_type_score("coding"), 0.95);
    assert_eq!(profile.get_task_type_score("documentation"), 0.60);
}

#[tokio::test]
async fn test_coordinator_assignment_with_learning_scores() {
    use vapora_agents::{
        AgentRegistry, AgentMetadata, AgentCoordinator,
    };
    use std::sync::Arc;

    // Create registry with test agents
    let registry = Arc::new(AgentRegistry::new(10));

    // Register two agents for developer role
    let agent_a = AgentMetadata::new(
        "developer".to_string(),
        "Agent A - Coding Specialist".to_string(),
        "claude".to_string(),
        "claude-opus-4-5".to_string(),
        vec!["coding".to_string(), "testing".to_string()],
    );

    let agent_b = AgentMetadata::new(
        "developer".to_string(),
        "Agent B - Generalist".to_string(),
        "claude".to_string(),
        "claude-sonnet-4".to_string(),
        vec!["coding".to_string(), "documentation".to_string()],
    );

    let agent_a_id = agent_a.id.clone();
    let agent_b_id = agent_b.id.clone();

    registry.register_agent(agent_a).ok();
    registry.register_agent(agent_b).ok();

    // Create coordinator
    let coordinator = AgentCoordinator::with_registry(registry);

    // Create learning profiles: Agent A excels at coding, Agent B is mediocre
    let now = Utc::now();
    let agent_a_executions: Vec<ExecutionData> = (0..20)
        .map(|i| ExecutionData {
            timestamp: now - Duration::days(i),
            duration_ms: 100,
            success: i < 19, // 95% success rate
        })
        .collect();

    let agent_b_executions: Vec<ExecutionData> = (0..20)
        .map(|i| ExecutionData {
            timestamp: now - Duration::days(i),
            duration_ms: 100,
            success: i < 14, // 70% success rate
        })
        .collect();

    let agent_a_expertise = TaskTypeExpertise::from_executions(agent_a_executions, "coding");
    let agent_b_expertise = TaskTypeExpertise::from_executions(agent_b_executions, "coding");

    let mut agent_a_profile = ProfileAdapter::create_learning_profile(agent_a_id.clone());
    agent_a_profile =
        ProfileAdapter::add_task_type_expertise(agent_a_profile, "coding".to_string(), agent_a_expertise);

    let mut agent_b_profile = ProfileAdapter::create_learning_profile(agent_b_id.clone());
    agent_b_profile =
        ProfileAdapter::add_task_type_expertise(agent_b_profile, "coding".to_string(), agent_b_expertise);

    // Update coordinator with learning profiles
    coordinator
        .update_learning_profile(&agent_a_id, agent_a_profile)
        .ok();
    coordinator
        .update_learning_profile(&agent_b_id, agent_b_profile)
        .ok();

    // Assign a coding task
    let _task_id = coordinator
        .assign_task(
            "developer",
            "Implement authentication module".to_string(),
            "Create secure login and token validation".to_string(),
            "Security critical".to_string(),
            2,
        )
        .await
        .expect("Should assign task");

    // Get the registry to verify which agent was selected
    let registry = coordinator.registry();
    let agent_a_tasks = registry.list_all()
        .iter()
        .find(|a| a.id == agent_a_id)
        .map(|a| a.current_tasks)
        .unwrap_or(0);

    let agent_b_tasks = registry.list_all()
        .iter()
        .find(|a| a.id == agent_b_id)
        .map(|a| a.current_tasks)
        .unwrap_or(0);

    // Agent A (higher expertise in coding) should have been selected
    assert!(agent_a_tasks > 0, "Agent A (coding specialist) should have 1+ tasks");
    assert_eq!(agent_b_tasks, 0, "Agent B (generalist) should have 0 tasks");

    // Verify learning profiles are stored
    let stored_profiles = coordinator.get_all_learning_profiles();
    assert!(stored_profiles.contains_key(&agent_a_id), "Agent A profile should be stored");
    assert!(stored_profiles.contains_key(&agent_b_id), "Agent B profile should be stored");
}