use rtx_auto::{
    agents::{
        DataEngineeringAgent, KernelSynthesizerAgent, ParallelPlannerAgent, QuantGuardianAgent,
    },
    error::AutoError,
    proposal::{Proposal, ProposalType, ProposalValidator},
    rollback::RollbackManager,
};
use rtx_runtime::Runtime;
use rtx_tensor::{DType, Device, Tensor};
use std::collections::HashMap;
use std::sync::Arc;

#[tokio::test]
async fn test_agent_collaboration() {
    let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));

    // Create all agents
    let data_agent = DataEngineeringAgent::new(runtime.clone()).unwrap();
    let parallel_agent = ParallelPlannerAgent::new(runtime.clone()).unwrap();
    let quant_agent = QuantGuardianAgent::new(runtime.clone()).unwrap();
    let kernel_agent = KernelSynthesizerAgent::new(runtime.clone()).unwrap();

    // Create a tensor for optimization
    let device = Device::cuda(0).unwrap_or(Device::Cpu);
    let tensor = Tensor::randn(&[1024, 1024], &device).unwrap();

    // Get proposals from each agent
    let data_proposals = data_agent.generate_layout_proposals(&tensor).await.unwrap();
    let quant_proposals = quant_agent
        .generate_quantization_proposals(&tensor, 0.95)
        .await
        .unwrap();

    assert!(!data_proposals.is_empty());
    assert!(!quant_proposals.is_empty());

    // Verify that proposals can be combined
    let mut all_proposals = data_proposals;
    all_proposals.extend(quant_proposals);

    assert!(
        all_proposals.len() >= 2,
        "Should have proposals from multiple agents"
    );
}

#[tokio::test]
async fn test_end_to_end_optimization_pipeline() {
    let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));

    // Set up the pipeline components
    let data_agent = DataEngineeringAgent::new(runtime.clone()).unwrap();
    let validator = ProposalValidator::new(runtime.clone()).unwrap();
    let mut rollback_manager = RollbackManager::new(runtime.clone()).unwrap();

    // Create initial state
    let mut initial_state = HashMap::new();
    initial_state.insert("performance_baseline".to_string(), vec![1.0]);
    initial_state.insert("accuracy".to_string(), vec![0.95]);

    // Create checkpoint
    let checkpoint = rollback_manager
        .create_checkpoint(
            rtx_auto::rollback::CheckpointType::BeforeOptimization,
            initial_state.clone(),
            "Baseline before optimization".to_string(),
        )
        .await
        .unwrap();

    rollback_manager.save_checkpoint(&checkpoint).await.unwrap();

    // Generate proposals
    let device = Device::cuda(0).unwrap_or(Device::Cpu);
    let tensor = Tensor::randn(&[512, 512], &device).unwrap();
    let proposals = data_agent.generate_layout_proposals(&tensor).await.unwrap();

    // Validate and rank proposals
    let ranked_proposals = validator.rank_proposals(&proposals).await.unwrap();
    assert!(!ranked_proposals.is_empty(), "Should have ranked proposals");

    // Select best proposal
    let (best_proposal, _score) = &ranked_proposals[0];
    assert_eq!(best_proposal.proposal_type(), ProposalType::DataLayout);
}

#[tokio::test]
async fn test_multi_agent_consensus() {
    let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));

    let data_agent = DataEngineeringAgent::new(runtime.clone()).unwrap();
    let kernel_agent = KernelSynthesizerAgent::new(runtime.clone()).unwrap();
    let validator = ProposalValidator::new(runtime.clone()).unwrap();

    // Get proposals from multiple agents for the same workload
    let device = Device::cuda(0).unwrap_or(Device::Cpu);
    let tensor = Tensor::randn(&[256, 256], &device).unwrap();
    let data_proposals = data_agent.generate_layout_proposals(&tensor).await.unwrap();

    // Also get quantization proposals to have multiple proposal types
    let quant_agent = QuantGuardianAgent::new(runtime.clone()).unwrap();
    let quant_proposals = quant_agent
        .generate_quantization_proposals(&tensor, 0.95)
        .await
        .unwrap();

    // Combine and rank all proposals
    let mut all_proposals = data_proposals;
    all_proposals.extend(quant_proposals);

    let ranked = validator.rank_proposals(&all_proposals).await.unwrap();
    assert!(!ranked.is_empty(), "Should rank combined proposals");

    // Verify diversity of proposal types
    let proposal_types: std::collections::HashSet<_> = ranked
        .iter()
        .map(|(proposal, _)| proposal.proposal_type())
        .collect();
    assert!(
        proposal_types.len() >= 2,
        "Should have proposals from different agents"
    );
}

#[tokio::test]
async fn test_adaptive_optimization_cycle() {
    let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));

    let quant_agent = QuantGuardianAgent::new(runtime.clone()).unwrap();
    let mut rollback_manager = RollbackManager::new(runtime.clone()).unwrap();

    // Set up adaptive thresholds
    rollback_manager
        .set_rollback_threshold("accuracy", 0.90)
        .await
        .unwrap();

    // Initial state with good performance
    let mut good_state = HashMap::new();
    good_state.insert("accuracy".to_string(), vec![0.95]);
    good_state.insert("latency".to_string(), vec![100.0]);

    let checkpoint = rollback_manager
        .create_checkpoint(
            rtx_auto::rollback::CheckpointType::Automatic,
            good_state,
            "Good performance state".to_string(),
        )
        .await
        .unwrap();

    rollback_manager.save_checkpoint(&checkpoint).await.unwrap();

    // Generate aggressive quantization proposal
    let device = Device::cuda(0).unwrap_or(Device::Cpu);
    let tensor = Tensor::randn(&[128, 128], &device).unwrap();
    let proposals = quant_agent
        .generate_quantization_proposals(&tensor, 0.85)
        .await
        .unwrap();

    assert!(
        !proposals.is_empty(),
        "Should generate quantization proposals"
    );

    // Simulate applying proposal with mixed results
    let mut result_state = HashMap::new();
    result_state.insert("accuracy".to_string(), vec![0.88]); // Below threshold
    result_state.insert("latency".to_string(), vec![50.0]); // Better latency

    let should_rollback = rollback_manager
        .should_trigger_automatic_rollback(&result_state)
        .await
        .unwrap();
    assert!(
        should_rollback,
        "Should trigger rollback due to accuracy degradation"
    );
}

#[tokio::test]
async fn test_proposal_conflict_resolution() {
    let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
    let validator = ProposalValidator::new(runtime.clone()).unwrap();

    // Create conflicting proposals
    let proposal1 = Proposal::new(
        ProposalType::DataLayout,
        "Optimize for sequential access".to_string(),
        1.5,
    );

    let proposal2 = Proposal::new(
        ProposalType::DataLayout,
        "Optimize for random access".to_string(),
        1.3,
    );

    let proposals = vec![proposal1, proposal2];

    // Detect conflicts
    let conflicts = validator.detect_conflicts(&proposals).await;
    assert!(conflicts.is_ok());

    let conflicts = conflicts.unwrap();
    assert!(!conflicts.is_empty(), "Should detect conflicting proposals");
}

#[tokio::test]
async fn test_performance_regression_detection() {
    let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
    let quant_agent = QuantGuardianAgent::new(runtime.clone()).unwrap();

    // Create reference tensor with known values
    let device = Device::cuda(0).unwrap_or(Device::Cpu);
    let reference = Tensor::ones(&[64, 64], &device).unwrap();

    // Create degraded version
    let degraded_data = vec![0.5f32; 64 * 64];
    let degraded = Tensor::from_vec(degraded_data, &[64, 64], &device).unwrap();

    // Monitor accuracy should detect significant degradation
    let metrics = quant_agent
        .monitor_accuracy(&reference, &degraded)
        .await
        .unwrap();
    assert!(metrics.mse > 0.0, "Should detect mean squared error");

    // Use threshold 0.8 - degradation detected when mse > (1.0 - 0.8) = 0.2
    // MSE between [1,1,1,...] and [0.5,0.5,0.5,...] is 0.25, which is > 0.2
    let is_degraded = quant_agent
        .detect_accuracy_degradation(&reference, &degraded, 0.8)
        .await
        .unwrap();
    assert!(is_degraded, "Should detect accuracy degradation");
}