use rtx_auto::{
    agents::DataEngineeringAgent,
    error::AutoError,
    proposal::{Proposal, ProposalStatus, ProposalType},
};
use rtx_runtime::Runtime;
use rtx_tensor::{DType, Device, Tensor};
use std::sync::Arc;

#[tokio::test]
async fn test_data_engineering_agent_creation() {
    let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
    let agent = DataEngineeringAgent::new(runtime.clone());
    assert!(agent.is_ok(), "Failed to create DataEngineeringAgent");
}

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

    let device = Device::cuda(0).unwrap_or(Device::Cpu);
    let tensors = vec![
        Tensor::zeros_typed(&[1000, 512], DType::F32, &device).unwrap(),
        Tensor::zeros_typed(&[512, 256], DType::F32, &device).unwrap(),
    ];

    let access_patterns = agent.analyze_access_patterns(&tensors).await;
    assert!(access_patterns.is_ok());

    let patterns = access_patterns.unwrap();
    assert!(!patterns.is_empty(), "Should detect access patterns");
}

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

    let device = Device::cuda(0).unwrap_or(Device::Cpu);
    let tensor = Tensor::zeros_typed(&[1024, 768], DType::F32, &device).unwrap();
    let proposals = agent.generate_layout_proposals(&tensor).await;

    assert!(proposals.is_ok());
    let proposals = proposals.unwrap();
    assert!(
        !proposals.is_empty(),
        "Should generate at least one layout proposal"
    );

    for proposal in &proposals {
        assert_eq!(proposal.proposal_type(), ProposalType::DataLayout);
        assert_eq!(proposal.status(), ProposalStatus::Pending);
        assert!(proposal.expected_performance_gain() > 0.0);
    }
}

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

    let device = Device::cuda(0).unwrap_or(Device::Cpu);
    let tensors = vec![
        Tensor::zeros_typed(&[100, 100], DType::F32, &device).unwrap(),
        Tensor::zeros_typed(&[100, 100], DType::F32, &device).unwrap(),
        Tensor::zeros_typed(&[100, 100], DType::F32, &device).unwrap(),
    ];

    let proposals = agent.generate_coalescing_proposals(&tensors).await;
    assert!(proposals.is_ok());

    let proposals = proposals.unwrap();
    assert!(
        !proposals.is_empty(),
        "Should generate coalescing proposals"
    );

    for proposal in &proposals {
        assert_eq!(proposal.proposal_type(), ProposalType::MemoryCoalescing);
        assert!(proposal.description().to_lowercase().contains("coalesce"));
    }
}

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

    let device = Device::cuda(0).unwrap_or(Device::Cpu);
    let tensor = Tensor::zeros_typed(&[2048, 2048], DType::F32, &device).unwrap();
    let proposals = agent.generate_cache_optimization_proposals(&tensor).await;

    assert!(proposals.is_ok());
    let proposals = proposals.unwrap();

    for proposal in &proposals {
        assert_eq!(proposal.proposal_type(), ProposalType::CacheOptimization);
        assert!(proposal.expected_performance_gain() > 0.0);
        assert!(proposal.description().len() > 10);
    }
}