//! Tests for continuous batching scheduler with SLA lanes
//!
//! This module contains comprehensive tests for the continuous batching
//! system that are designed to FAIL FIRST, then drive implementation.

use rtx_inference::{
    error::InferenceError,
    request::{InferenceRequest, RequestId, RequestPriority},
    scheduler::{BatchScheduler, BatchSchedulerConfig, SlaLane},
};
use std::time::{Duration, Instant};
use tokio::time::sleep;

/// Test that scheduler can be created with valid configuration
#[tokio::test]
async fn test_scheduler_creation() {
    let config = BatchSchedulerConfig {
        max_batch_size: 32,
        max_wait_time: Duration::from_millis(100),
        sla_lanes: vec![
            SlaLane {
                name: "premium".to_string(),
                priority: RequestPriority::High,
                max_latency: Duration::from_millis(50),
                max_batch_size: 16,
                memory_limit: Some(1024 * 1024 * 100), // 100MB
                min_wait_time: Some(Duration::from_millis(10)),
            },
            SlaLane {
                name: "standard".to_string(),
                priority: RequestPriority::Normal,
                max_latency: Duration::from_millis(150),
                max_batch_size: 32,
                memory_limit: Some(1024 * 1024 * 200), // 200MB
                min_wait_time: Some(Duration::from_millis(20)),
            },
        ],
        ..Default::default()
    };

    let scheduler = BatchScheduler::new(config).await;
    assert!(scheduler.is_ok());

    let scheduler = scheduler.unwrap();
    assert_eq!(scheduler.config().max_batch_size, 32);
    assert_eq!(scheduler.config().sla_lanes.len(), 2);
}

/// Test request submission and queue management
#[tokio::test]
async fn test_request_submission() {
    let config = BatchSchedulerConfig::default();
    let mut scheduler = BatchScheduler::new(config).await.unwrap();

    let request = InferenceRequest {
        id: RequestId::new(),
        model_name: "llama-7b".to_string(),
        input_tokens: vec![1, 2, 3, 4, 5],
        max_new_tokens: 100,
        temperature: 0.8,
        priority: RequestPriority::Normal,
        deadline: Some(Instant::now() + Duration::from_millis(1000)),
        ..Default::default()
    };

    let result = scheduler.submit_request(request.clone()).await;
    assert!(result.is_ok());

    let queue_stats = scheduler.queue_stats().await;
    assert_eq!(queue_stats.total_pending, 1);
    assert_eq!(queue_stats.lanes.get("standard").unwrap().pending_count, 1);
}

/// Test SLA lane assignment based on request priority
#[tokio::test]
async fn test_sla_lane_assignment() {
    let config = BatchSchedulerConfig {
        sla_lanes: vec![
            SlaLane {
                name: "premium".to_string(),
                priority: RequestPriority::High,
                max_latency: Duration::from_millis(50),
                max_batch_size: 8,
                memory_limit: Some(1024 * 1024 * 50), // 50MB
                min_wait_time: Some(Duration::from_millis(5)),
            },
            SlaLane {
                name: "standard".to_string(),
                priority: RequestPriority::Normal,
                max_latency: Duration::from_millis(150),
                max_batch_size: 32,
                memory_limit: Some(1024 * 1024 * 200), // 200MB
                min_wait_time: Some(Duration::from_millis(20)),
            },
        ],
        ..Default::default()
    };

    let mut scheduler = BatchScheduler::new(config).await.unwrap();

    // Submit high priority request
    let high_priority_request = InferenceRequest {
        id: RequestId::new(),
        priority: RequestPriority::High,
        input_tokens: vec![1, 2, 3],
        ..Default::default()
    };

    scheduler
        .submit_request(high_priority_request)
        .await
        .unwrap();

    // Submit normal priority request
    let normal_priority_request = InferenceRequest {
        id: RequestId::new(),
        priority: RequestPriority::Normal,
        input_tokens: vec![4, 5, 6],
        ..Default::default()
    };

    scheduler
        .submit_request(normal_priority_request)
        .await
        .unwrap();

    let queue_stats = scheduler.queue_stats().await;
    assert_eq!(queue_stats.lanes.get("premium").unwrap().pending_count, 1);
    assert_eq!(queue_stats.lanes.get("standard").unwrap().pending_count, 1);
}

/// Test dynamic batch formation with different batch sizes
#[tokio::test]
async fn test_dynamic_batch_formation() {
    let config = BatchSchedulerConfig {
        max_batch_size: 4,
        max_wait_time: Duration::from_millis(50),
        ..Default::default()
    };

    let mut scheduler = BatchScheduler::new(config).await.unwrap();

    // Submit multiple requests rapidly
    let mut request_ids = Vec::new();
    for i in 0..6 {
        let request = InferenceRequest {
            id: RequestId::new(),
            input_tokens: vec![i; 10],
            max_new_tokens: 50,
            ..Default::default()
        };
        request_ids.push(request.id);
        scheduler.submit_request(request).await.unwrap();
    }

    // Should form batches of size 4, 2
    let batch1 = scheduler.get_next_batch().await.unwrap();
    assert_eq!(batch1.requests.len(), 4);

    let batch2 = scheduler.get_next_batch().await.unwrap();
    assert_eq!(batch2.requests.len(), 2);
}

/// Test batch scheduling with SLA constraints
#[tokio::test]
async fn test_sla_constrained_batching() {
    let config = BatchSchedulerConfig {
        sla_lanes: vec![
            SlaLane {
                name: "premium".to_string(),
                priority: RequestPriority::High,
                max_latency: Duration::from_millis(30),
                max_batch_size: 2, // Small batch size for low latency
                memory_limit: Some(1024 * 1024 * 10), // 10MB
                min_wait_time: Some(Duration::from_millis(1)),
            },
            SlaLane {
                name: "standard".to_string(),
                priority: RequestPriority::Normal,
                max_latency: Duration::from_millis(100),
                max_batch_size: 8,
                memory_limit: Some(1024 * 1024 * 50), // 50MB
                min_wait_time: Some(Duration::from_millis(5)),
            },
        ],
        ..Default::default()
    };

    let mut scheduler = BatchScheduler::new(config).await.unwrap();

    // Submit premium requests - should batch quickly due to SLA
    for i in 0..3 {
        let request = InferenceRequest {
            id: RequestId::new(),
            priority: RequestPriority::High,
            input_tokens: vec![i; 5],
            ..Default::default()
        };
        scheduler.submit_request(request).await.unwrap();
    }

    // Premium batch should form with max 2 requests
    let premium_batch = scheduler.get_next_batch().await.unwrap();
    assert_eq!(premium_batch.requests.len(), 2);
    assert_eq!(premium_batch.lane_name, "premium");

    // Remaining premium request should form another batch
    let remaining_batch = scheduler.get_next_batch().await.unwrap();
    assert_eq!(remaining_batch.requests.len(), 1);
    assert_eq!(remaining_batch.lane_name, "premium");
}

/// Test preemption of lower priority batches
#[tokio::test]
async fn test_batch_preemption() {
    let config = BatchSchedulerConfig {
        enable_preemption: true,
        sla_lanes: vec![
            SlaLane {
                name: "premium".to_string(),
                priority: RequestPriority::High,
                max_latency: Duration::from_millis(50),
                max_batch_size: 4,
                memory_limit: Some(1024 * 1024 * 25), // 25MB
                min_wait_time: Some(Duration::from_millis(2)),
            },
            SlaLane {
                name: "standard".to_string(),
                priority: RequestPriority::Normal,
                max_latency: Duration::from_millis(200),
                max_batch_size: 8,
                memory_limit: Some(1024 * 1024 * 50), // 50MB
                min_wait_time: Some(Duration::from_millis(5)),
            },
        ],
        ..Default::default()
    };

    let mut scheduler = BatchScheduler::new(config).await.unwrap();

    // Submit normal priority requests
    for i in 0..4 {
        let request = InferenceRequest {
            id: RequestId::new(),
            priority: RequestPriority::Normal,
            input_tokens: vec![i; 20],
            ..Default::default()
        };
        scheduler.submit_request(request).await.unwrap();
    }

    // Get standard batch for processing
    let standard_batch = scheduler.get_next_batch().await.unwrap();
    assert_eq!(standard_batch.lane_name, "standard");

    // Mark batch as executing
    scheduler
        .mark_batch_executing(standard_batch.id)
        .await
        .unwrap();

    // Submit high priority request - should trigger preemption consideration
    let high_priority_request = InferenceRequest {
        id: RequestId::new(),
        priority: RequestPriority::High,
        input_tokens: vec![99; 10],
        deadline: Some(Instant::now() + Duration::from_millis(30)),
        ..Default::default()
    };

    scheduler
        .submit_request(high_priority_request)
        .await
        .unwrap();

    // Should suggest preemption
    let preemption_decision = scheduler.evaluate_preemption().await.unwrap();
    assert!(preemption_decision.should_preempt);
    assert_eq!(preemption_decision.target_batch_id, Some(standard_batch.id));
}

/// Test SLA violation detection and handling
#[tokio::test]
async fn test_sla_violation_detection() {
    let config = BatchSchedulerConfig {
        sla_lanes: vec![SlaLane {
            name: "premium".to_string(),
            priority: RequestPriority::High,
            max_latency: Duration::from_millis(100),
            max_batch_size: 4,
            memory_limit: Some(1024 * 1024 * 25), // 25MB
            min_wait_time: Some(Duration::from_millis(2)),
        }],
        ..Default::default()
    };

    let mut scheduler = BatchScheduler::new(config).await.unwrap();

    let request = InferenceRequest {
        id: RequestId::new(),
        priority: RequestPriority::High,
        input_tokens: vec![1, 2, 3],
        deadline: Some(Instant::now() + Duration::from_millis(50)),
        ..Default::default()
    };

    let request_id = request.id;
    scheduler.submit_request(request).await.unwrap();

    // Wait longer than deadline
    sleep(Duration::from_millis(60)).await;

    // Check for SLA violations
    let violations = scheduler.check_sla_violations().await.unwrap();
    assert!(!violations.is_empty());
    assert!(violations.iter().any(|v| v.request_id == request_id));
}

/// Test memory pressure handling in batch formation
#[tokio::test]
async fn test_memory_pressure_batching() {
    let config = BatchSchedulerConfig {
        max_batch_size: 16,
        memory_pressure_threshold: 0.8, // 80% memory usage
        ..Default::default()
    };

    let mut scheduler = BatchScheduler::new(config).await.unwrap();

    // Simulate high memory pressure
    scheduler.set_memory_pressure(0.85).await.unwrap();

    // Submit requests with varying memory requirements
    let large_request = InferenceRequest {
        id: RequestId::new(),
        input_tokens: vec![1; 2048], // Large context
        max_new_tokens: 1024,
        ..Default::default()
    };

    scheduler.submit_request(large_request).await.unwrap();

    let small_request = InferenceRequest {
        id: RequestId::new(),
        input_tokens: vec![1; 128], // Small context
        max_new_tokens: 64,
        ..Default::default()
    };

    scheduler.submit_request(small_request).await.unwrap();

    // Under memory pressure, should create smaller batches
    let batch = scheduler.get_next_batch().await.unwrap();
    assert!(batch.requests.len() <= 2); // Should limit batch size due to memory pressure
    assert!(batch.estimated_memory_usage <= scheduler.config().memory_limit);
}

/// Test graceful degradation under extreme load
#[tokio::test]
async fn test_graceful_degradation() {
    let config = BatchSchedulerConfig {
        max_batch_size: 4,
        max_queue_size: 10, // Limited queue size
        degradation_mode_threshold: 8,
        ..Default::default()
    };

    let mut scheduler = BatchScheduler::new(config).await.unwrap();

    // Submit more requests than queue can handle
    let mut rejected_count = 0;
    for i in 0..15 {
        let request = InferenceRequest {
            id: RequestId::new(),
            input_tokens: vec![i; 10],
            ..Default::default()
        };

        match scheduler.submit_request(request).await {
            Ok(_) => {}
            Err(InferenceError::QueueFull { .. }) => rejected_count += 1,
            Err(e) => panic!("Unexpected error: {}", e),
        }
    }

    // Should reject some requests to maintain stability
    assert!(rejected_count > 0);
    assert!(rejected_count >= 5); // Should reject excess requests

    // Queue should be in degradation mode
    assert!(scheduler.is_degradation_mode().await);
}