use std::collections::HashMap;
use std::sync::Arc;

use chrono::{Duration, Utc};
use vapora_agents::{
    AgentCoordinator, AgentMetadata, AgentRegistry, ExecutionData, ProfileAdapter,
    TaskTypeExpertise,
};
use vapora_llm_router::{BudgetManager, RoleBudget};

/// End-to-end integration test: Learning + Budget interaction
///
/// Verifies that:
/// 1. Agents with better learning profiles are selected for tasks
/// 2. Budget enforcement doesn't interfere with learning-based assignment
/// 3. Learning profiles improve over time with new executions
/// 4. Budget manager tracks spending correctly during task assignment
#[tokio::test]
async fn test_end_to_end_learning_with_budget_enforcement() {
    // Create agent registry and coordinator
    let registry = Arc::new(AgentRegistry::new(10));

    // Register two developers with different expertise
    let developer_a = AgentMetadata::new(
        "developer".to_string(),
        "Developer A - Coding Specialist".to_string(),
        "claude".to_string(),
        "claude-opus-4-5".to_string(),
        vec!["coding".to_string(), "testing".to_string()],
    );

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

    let dev_a_id = developer_a.id.clone();
    let dev_b_id = developer_b.id.clone();

    registry.register_agent(developer_a).ok();
    registry.register_agent(developer_b).ok();

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

    // Create budget manager with limits
    let mut budgets = HashMap::new();
    budgets.insert(
        "developer".to_string(),
        RoleBudget {
            role: "developer".to_string(),
            monthly_limit_cents: 100000, // $1000 for testing
            weekly_limit_cents: 25000,   // $250 for testing
            fallback_provider: "ollama".to_string(),
            alert_threshold: 0.8,
        },
    );

    let budget_manager = Arc::new(BudgetManager::new(budgets));
    let coordinator = coordinator.with_budget_manager(budget_manager.clone());

    // Simulate historical executions for developer_a (excellent at coding)
    let now = Utc::now();
    let dev_a_executions: Vec<ExecutionData> = (0..30)
        .map(|i| ExecutionData {
            timestamp: now - Duration::days(i),
            duration_ms: 200 + (i as u64 * 5),
            success: i < 28, // 93% success rate
        })
        .collect();

    // Simulate historical executions for developer_b (mediocre at coding)
    let dev_b_executions: Vec<ExecutionData> = (0..30)
        .map(|i| ExecutionData {
            timestamp: now - Duration::days(i),
            duration_ms: 300 + (i as u64 * 10),
            success: i < 20, // 67% success rate
        })
        .collect();

    // Calculate expertise from executions
    let dev_a_expertise = TaskTypeExpertise::from_executions(dev_a_executions, "coding");
    let dev_b_expertise = TaskTypeExpertise::from_executions(dev_b_executions, "coding");

    // Verify expertise calculations
    assert!(dev_a_expertise.success_rate > 0.9);
    assert!(dev_b_expertise.success_rate > 0.6 && dev_b_expertise.success_rate < 0.7);
    assert!(dev_a_expertise.success_rate > dev_b_expertise.success_rate);

    // Create learning profiles
    let mut profile_a = ProfileAdapter::create_learning_profile(dev_a_id.clone());
    profile_a =
        ProfileAdapter::add_task_type_expertise(profile_a, "coding".to_string(), dev_a_expertise);

    let mut profile_b = ProfileAdapter::create_learning_profile(dev_b_id.clone());
    profile_b =
        ProfileAdapter::add_task_type_expertise(profile_b, "coding".to_string(), dev_b_expertise);

    // Update coordinator with learning profiles
    coordinator
        .update_learning_profile(&dev_a_id, profile_a.clone())
        .ok();
    coordinator
        .update_learning_profile(&dev_b_id, profile_b.clone())
        .ok();

    // Verify profiles are stored
    let stored_a = coordinator.get_learning_profile(&dev_a_id);
    let stored_b = coordinator.get_learning_profile(&dev_b_id);
    assert!(stored_a.is_some());
    assert!(stored_b.is_some());

    // Check budget status before task assignment
    let budget_status = budget_manager.check_budget("developer").await.unwrap();
    assert!(!budget_status.exceeded);
    assert!(!budget_status.near_threshold);
    assert_eq!(budget_status.monthly_remaining_cents, 100000);

    // Assign a coding task (should go to developer_a based on learning)
    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");

    // Verify task was assigned (we can check via registry)
    let all_agents = coordinator.registry().list_all();
    let dev_a_tasks = all_agents
        .iter()
        .find(|a| a.id == dev_a_id)
        .map(|a| a.current_tasks)
        .unwrap_or(0);

    let _dev_b_tasks = all_agents
        .iter()
        .find(|a| a.id == dev_b_id)
        .map(|a| a.current_tasks)
        .unwrap_or(0);

    // Developer A (high expertise) should be selected
    assert!(
        dev_a_tasks > 0,
        "Developer A (high expertise) should have been assigned the task"
    );

    // Simulate task completion
    coordinator.complete_task(&task_id, &dev_a_id).await.ok();

    // Verify budget status is still within limits
    let budget_status = budget_manager.check_budget("developer").await.unwrap();
    assert!(!budget_status.exceeded, "Budget should not be exceeded");

    // Simulate multiple tasks to test cumulative budget tracking
    for i in 0..5 {
        let task = coordinator
            .assign_task(
                "developer",
                format!("Task {}", i),
                "Test description".to_string(),
                "Context".to_string(),
                1,
            )
            .await;

        if task.is_ok() {
            let agents = coordinator.registry().list_all();
            if let Some(dev_a) = agents.iter().find(|a| a.id == dev_a_id) {
                coordinator
                    .complete_task(&format!("task-{}", i), &dev_a.id)
                    .await
                    .ok();
            }
        }
    }

    // Final budget status check
    let final_budget = budget_manager.check_budget("developer").await.unwrap();
    assert!(
        final_budget.monthly_utilization < 1.0,
        "Should not exceed monthly budget"
    );

    // Verify learning profiles are still intact
    let all_profiles = coordinator.get_all_learning_profiles();
    assert_eq!(all_profiles.len(), 2, "Both profiles should be stored");
}

/// Test that budget enforcement doesn't break learning-based selection
#[tokio::test]
async fn test_learning_selection_with_budget_constraints() {
    let registry = Arc::new(AgentRegistry::new(10));

    let agent_expert = AgentMetadata::new(
        "developer".to_string(),
        "Expert Developer".to_string(),
        "claude".to_string(),
        "claude-opus-4-5".to_string(),
        vec!["coding".to_string()],
    );

    let agent_novice = AgentMetadata::new(
        "developer".to_string(),
        "Novice Developer".to_string(),
        "claude".to_string(),
        "claude-sonnet-4".to_string(),
        vec!["coding".to_string()],
    );

    let expert_id = agent_expert.id.clone();
    let novice_id = agent_novice.id.clone();

    registry.register_agent(agent_expert).ok();
    registry.register_agent(agent_novice).ok();

    let coordinator = AgentCoordinator::with_registry(registry);

    // Create tight budget to test constraint handling
    let mut budgets = HashMap::new();
    budgets.insert(
        "developer".to_string(),
        RoleBudget {
            role: "developer".to_string(),
            monthly_limit_cents: 10000, // $100 (tight)
            weekly_limit_cents: 2500,   // $25 (tight)
            fallback_provider: "ollama".to_string(),
            alert_threshold: 0.9, // Alert at 90%
        },
    );

    let budget_manager = Arc::new(BudgetManager::new(budgets));
    let coordinator = coordinator.with_budget_manager(budget_manager.clone());

    // Create expertise profiles
    let now = Utc::now();
    let expert_execs: Vec<ExecutionData> = (0..20)
        .map(|i| ExecutionData {
            timestamp: now - Duration::days(i),
            duration_ms: 100,
            success: i < 19, // 95% success
        })
        .collect();

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

    let expert_expertise = TaskTypeExpertise::from_executions(expert_execs, "coding");
    let novice_expertise = TaskTypeExpertise::from_executions(novice_execs, "coding");

    let mut expert_profile = ProfileAdapter::create_learning_profile(expert_id.clone());
    expert_profile = ProfileAdapter::add_task_type_expertise(
        expert_profile,
        "coding".to_string(),
        expert_expertise,
    );

    let mut novice_profile = ProfileAdapter::create_learning_profile(novice_id.clone());
    novice_profile = ProfileAdapter::add_task_type_expertise(
        novice_profile,
        "coding".to_string(),
        novice_expertise,
    );

    coordinator
        .update_learning_profile(&expert_id, expert_profile)
        .ok();
    coordinator
        .update_learning_profile(&novice_id, novice_profile)
        .ok();

    // Verify budget status
    let status = budget_manager.check_budget("developer").await.unwrap();
    assert!(
        !status.exceeded && !status.near_threshold,
        "Initial budget should be healthy"
    );

    // Assign multiple tasks - expert should be consistently selected
    let mut expert_count = 0;
    for i in 0..3 {
        if let Ok(_task_id) = coordinator
            .assign_task(
                "developer",
                format!("Coding Task {}", i),
                "Implement feature".to_string(),
                "Production".to_string(),
                1,
            )
            .await
        {
            let agents = coordinator.registry().list_all();
            if let Some(expert) = agents.iter().find(|a| a.id == expert_id) {
                if expert.current_tasks > 0 {
                    expert_count += 1;
                }
            }
        }
    }

    // Expert should have been selected more often
    assert!(
        expert_count > 0,
        "Expert should have been selected despite budget constraints"
    );
}

/// Test that learning profile improvements are tracked correctly
#[tokio::test]
async fn test_learning_profile_improvement_with_budget_tracking() {
    let registry = Arc::new(AgentRegistry::new(10));

    let agent = AgentMetadata::new(
        "developer".to_string(),
        "Improving Developer".to_string(),
        "claude".to_string(),
        "claude-sonnet-4".to_string(),
        vec!["coding".to_string()],
    );

    let agent_id = agent.id.clone();
    registry.register_agent(agent).ok();

    let coordinator = AgentCoordinator::with_registry(registry);

    // Create budget manager
    let mut budgets = HashMap::new();
    budgets.insert(
        "developer".to_string(),
        RoleBudget {
            role: "developer".to_string(),
            monthly_limit_cents: 50000,
            weekly_limit_cents: 12500,
            fallback_provider: "ollama".to_string(),
            alert_threshold: 0.8,
        },
    );

    let budget_manager = Arc::new(BudgetManager::new(budgets));
    let coordinator = coordinator.with_budget_manager(budget_manager.clone());

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

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

    let mut profile = ProfileAdapter::create_learning_profile(agent_id.clone());
    profile = ProfileAdapter::add_task_type_expertise(
        profile,
        "coding".to_string(),
        initial_expertise.clone(),
    );

    coordinator
        .update_learning_profile(&agent_id, profile.clone())
        .ok();

    // Check initial profile
    let stored_profile = coordinator.get_learning_profile(&agent_id).unwrap();
    assert_eq!(
        stored_profile.get_task_type_score("coding"),
        initial_expertise.success_rate
    );

    // Simulate improvement: add successful recent executions
    let new_exec = ExecutionData {
        timestamp: now,
        duration_ms: 120,
        success: true,
    };

    initial_expertise.update_with_execution(&new_exec);
    assert!(
        initial_expertise.success_rate > 0.5,
        "Success rate should improve with new successful execution"
    );

    // Update profile with improved expertise
    let mut updated_profile = ProfileAdapter::create_learning_profile(agent_id.clone());
    updated_profile = ProfileAdapter::add_task_type_expertise(
        updated_profile,
        "coding".to_string(),
        initial_expertise,
    );

    coordinator
        .update_learning_profile(&agent_id, updated_profile)
        .ok();

    // Verify improvement is reflected
    let final_profile = coordinator.get_learning_profile(&agent_id).unwrap();
    let final_score = final_profile.get_task_type_score("coding");
    assert!(final_score > 0.5, "Final score should reflect improvement");

    // Verify budget tracking is unaffected
    let status = budget_manager.check_budget("developer").await.unwrap();
    assert!(!status.exceeded);
}
