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()],
    );

    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());

    // Build a role-level learning profile keyed by stable_id ("developer").
    // All developer agents share one profile since stable_id = role.
    let now = Utc::now();
    let dev_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();

    let dev_expertise = TaskTypeExpertise::from_executions(dev_executions, "coding");

    assert!(dev_expertise.success_rate > 0.9);

    let mut role_profile = ProfileAdapter::create_learning_profile("developer".to_string());
    role_profile =
        ProfileAdapter::add_task_type_expertise(role_profile, "coding".to_string(), dev_expertise);

    // Insert under stable_id = "developer" so assign_task can find it.
    coordinator
        .update_learning_profile("developer", role_profile)
        .ok();

    // Verify profile is stored under stable_id key
    let stored = coordinator.get_learning_profile("developer");
    assert!(stored.is_some(), "Role-level profile must be stored");

    // 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 — learning profile found → profile-based routing.
    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 one developer was assigned the task
    let all_agents = coordinator.registry().list_all();
    let total_assigned: u32 = all_agents.iter().map(|a| a.current_tasks).sum();
    assert_eq!(total_assigned, 1, "Exactly one agent must have the task");

    // Complete the task using the agent that received it
    if let Some(assigned) = all_agents.iter().find(|a| a.current_tasks > 0) {
        coordinator.complete_task(&task_id, &assigned.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 {
        if let Ok(tid) = coordinator
            .assign_task(
                "developer",
                format!("Task {}", i),
                "Test description".to_string(),
                "Context".to_string(),
                1,
            )
            .await
        {
            let agents = coordinator.registry().list_all();
            if let Some(assigned) = agents.iter().find(|a| a.current_tasks > 0) {
                coordinator.complete_task(&tid, &assigned.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 role-level profile is intact
    let all_profiles = coordinator.get_all_learning_profiles();
    assert!(
        all_profiles.contains_key("developer"),
        "Role-level profile must survive task processing"
    );
}

/// 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()],
    );

    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());

    // Build a role-level learning profile for "developer" (stable_id = role).
    let now = Utc::now();
    let role_execs: Vec<ExecutionData> = (0..20)
        .map(|i| ExecutionData {
            timestamp: now - Duration::days(i),
            duration_ms: 100,
            success: i < 19, // 95% success
        })
        .collect();

    let role_expertise = TaskTypeExpertise::from_executions(role_execs, "coding");

    let mut role_profile = ProfileAdapter::create_learning_profile("developer".to_string());
    role_profile =
        ProfileAdapter::add_task_type_expertise(role_profile, "coding".to_string(), role_expertise);

    coordinator
        .update_learning_profile("developer", role_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 — profile-based scoring should pick some developer.
    let mut assigned_count = 0;
    for i in 0..3 {
        if let Ok(tid) = coordinator
            .assign_task(
                "developer",
                format!("Coding Task {}", i),
                "Implement feature".to_string(),
                "Production".to_string(),
                1,
            )
            .await
        {
            assigned_count += 1;
            let agents = coordinator.registry().list_all();
            if let Some(assigned) = agents.iter().find(|a| a.current_tasks > 0) {
                coordinator.complete_task(&tid, &assigned.id).await.ok();
            }
        }
    }

    assert!(
        assigned_count > 0,
        "Tasks should be assigned 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()],
    );

    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 — keyed by stable_id = "developer".
    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("developer".to_string());
    profile = ProfileAdapter::add_task_type_expertise(
        profile,
        "coding".to_string(),
        initial_expertise.clone(),
    );

    coordinator
        .update_learning_profile("developer", profile)
        .ok();

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

    // Simulate improvement: add successful recent execution
    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("developer".to_string());
    updated_profile = ProfileAdapter::add_task_type_expertise(
        updated_profile,
        "coding".to_string(),
        initial_expertise,
    );

    coordinator
        .update_learning_profile("developer", updated_profile)
        .ok();

    // Verify improvement is reflected
    let final_profile = coordinator.get_learning_profile("developer").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);
}
