//! Continuous batching scheduler with SLA lanes
//!
//! This module implements a sophisticated scheduler that forms dynamic batches
//! based on SLA requirements, memory constraints, and performance optimization.

use serde::{Deserialize, Serialize};
use std::cmp::Ordering;
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use tracing::{debug, trace, warn};
use uuid::Uuid;

use crate::error::{InferenceError, InferenceResult};
use crate::request::{InferenceRequest, RequestId, RequestPriority};

/// Unique identifier for batches
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct BatchId(Uuid);

impl BatchId {
    /// Generate a new unique batch ID
    #[must_use]
    pub fn new() -> Self {
        Self(Uuid::new_v4())
    }
}

impl std::fmt::Display for BatchId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.fmt(f)
    }
}

/// SLA lane configuration for request segregation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SlaLane {
    /// Lane name for identification
    pub name: String,

    /// Priority level for this lane
    pub priority: RequestPriority,

    /// Maximum acceptable latency for requests in this lane
    pub max_latency: Duration,

    /// Maximum batch size for this lane
    pub max_batch_size: usize,

    /// Memory limit per batch (bytes)
    pub memory_limit: Option<usize>,

    /// Minimum batch wait time before processing
    pub min_wait_time: Option<Duration>,
}

/// Preemption decision result
#[derive(Debug, Clone)]
pub struct PreemptionDecision {
    /// Whether preemption should occur
    pub should_preempt: bool,

    /// Target batch to preempt (if any)
    pub target_batch_id: Option<BatchId>,

    /// Reason for preemption decision
    pub reason: String,

    /// Priority difference that triggered decision
    pub priority_delta: i32,
}

/// SLA violation information
#[derive(Debug, Clone)]
pub struct SlaViolation {
    /// Request that violated SLA
    pub request_id: RequestId,

    /// Lane the request was assigned to
    pub lane_name: String,

    /// Type of violation
    pub violation_type: String,

    /// Time of violation
    pub violation_time: Instant,

    /// Expected vs actual latency
    pub expected_latency: Duration,
    pub actual_latency: Duration,
}

/// Batch of requests ready for processing
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct InferenceBatch {
    /// Unique batch identifier
    pub id: BatchId,

    /// Requests in this batch
    pub requests: Vec<InferenceRequest>,

    /// Lane this batch belongs to
    pub lane_name: String,

    /// Batch creation timestamp
    pub created_at: Instant,

    /// Estimated memory usage
    pub estimated_memory_usage: usize,

    /// Batch priority (highest among constituent requests)
    pub priority: RequestPriority,

    /// Expected processing duration
    pub estimated_duration: Duration,
}

/// Lane statistics for monitoring
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LaneStats {
    /// Number of requests pending in this lane
    pub pending_count: usize,

    /// Number of requests currently processing
    pub processing_count: usize,

    /// Average latency for completed requests
    pub average_latency: Duration,

    /// SLA violation rate
    pub violation_rate: f64,

    /// Throughput (requests per second)
    pub throughput: f64,

    /// Memory utilization
    pub memory_utilization: f64,
}

/// Queue statistics across all lanes
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueueStats {
    /// Total pending requests across all lanes
    pub total_pending: usize,

    /// Total processing requests
    pub total_processing: usize,

    /// Statistics per lane
    pub lanes: HashMap<String, LaneStats>,

    /// Overall queue utilization
    pub queue_utilization: f64,

    /// Memory pressure level (0.0-1.0)
    pub memory_pressure: f64,
}

/// Configuration for batch scheduler
#[derive(Debug, Clone)]
pub struct BatchSchedulerConfig {
    /// Maximum batch size across all lanes
    pub max_batch_size: usize,

    /// Maximum wait time before forming batch
    pub max_wait_time: Duration,

    /// SLA lanes configuration
    pub sla_lanes: Vec<SlaLane>,

    /// Enable preemption of lower priority batches
    pub enable_preemption: bool,

    /// Memory pressure threshold for degradation
    pub memory_pressure_threshold: f64,

    /// Maximum memory available for batching
    pub memory_limit: usize,

    /// Degradation mode activation threshold
    pub degradation_mode_threshold: usize,

    /// Maximum queue size before rejection
    pub max_queue_size: usize,
}

impl Default for BatchSchedulerConfig {
    fn default() -> Self {
        Self {
            max_batch_size: 32,
            max_wait_time: Duration::from_millis(100),
            sla_lanes: vec![SlaLane {
                name: "standard".to_string(),
                priority: RequestPriority::Normal,
                max_latency: Duration::from_millis(200),
                max_batch_size: 32,
                memory_limit: None,
                min_wait_time: None,
            }],
            enable_preemption: false,
            memory_pressure_threshold: 0.8,
            memory_limit: 8 * 1024 * 1024 * 1024, // 8GB
            degradation_mode_threshold: 1000,
            max_queue_size: 2000,
        }
    }
}

/// Priority-ordered batch request for scheduling
#[derive(Debug)]
struct SchedulerRequest {
    request: InferenceRequest,
    assigned_lane: String,
    queued_at: Instant,
    priority_score: i32,
}

impl PartialEq for SchedulerRequest {
    fn eq(&self, other: &Self) -> bool {
        self.priority_score == other.priority_score
    }
}

impl Eq for SchedulerRequest {}

impl PartialOrd for SchedulerRequest {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for SchedulerRequest {
    fn cmp(&self, other: &Self) -> Ordering {
        // Higher priority score comes first
        other.priority_score.cmp(&self.priority_score)
    }
}

/// Lane state tracking
#[derive(Debug)]
struct LaneState {
    config: SlaLane,
    pending_requests: VecDeque<SchedulerRequest>,
    processing_batches: HashMap<BatchId, InferenceBatch>,
    completed_requests: VecDeque<(Instant, Duration)>, // (completion_time, latency)
    violation_count: usize,
    total_processed: usize,
}

impl LaneState {
    fn new(config: SlaLane) -> Self {
        Self {
            config,
            pending_requests: VecDeque::new(),
            processing_batches: HashMap::new(),
            completed_requests: VecDeque::new(),
            violation_count: 0,
            total_processed: 0,
        }
    }

    fn add_request(&mut self, request: SchedulerRequest) {
        self.pending_requests.push_back(request);
    }

    fn has_pending_requests(&self) -> bool {
        !self.pending_requests.is_empty()
    }

    fn pending_count(&self) -> usize {
        self.pending_requests.len()
    }

    fn processing_count(&self) -> usize {
        self.processing_batches.len()
    }

    fn should_form_batch(&self, current_memory_usage: usize, memory_limit: usize) -> bool {
        if self.pending_requests.is_empty() {
            return false;
        }

        // Always allow at least one request to form a batch
        let min_batch_size = 1;
        if self.pending_requests.len() < min_batch_size {
            return false;
        }

        // Check memory constraints
        let estimated_batch_memory = self.estimate_batch_memory();
        if current_memory_usage + estimated_batch_memory > memory_limit {
            return false;
        }

        // Check if we should wait for more requests (but be lenient for batch formation)
        if let Some(min_wait) = self.config.min_wait_time
            && let Some(oldest) = self.pending_requests.front()
        {
            let wait_time = oldest.queued_at.elapsed();

            // Check for urgent deadlines that should bypass min_wait_time
            let has_urgent_deadline = self.pending_requests.iter().any(|req| {
                if let Some(deadline) = req.request.deadline {
                    deadline.saturating_duration_since(Instant::now()) < Duration::from_millis(100)
                } else {
                    false
                }
            });

            // Be more aggressive about batch formation:
            // - Always form batch if we only have 1 request (avoid starvation)
            // - Form batch if we have any requests and have been waiting at least half the min_wait
            // - Or if we have urgent deadlines
            // - Or if we have reached desired batch size
            let half_min_wait = min_wait / 2;
            let desired_batch_size = std::cmp::max(1, self.config.max_batch_size / 2);

            let should_wait = wait_time < half_min_wait
                && self.pending_requests.len() < desired_batch_size
                && !has_urgent_deadline
                && self.pending_requests.len() > 1; // Don't wait for single requests

            if should_wait {
                return false;
            }
        }

        true
    }

    fn form_batch(&mut self) -> Option<InferenceBatch> {
        if self.pending_requests.is_empty() {
            return None;
        }

        let mut batch_requests = Vec::new();
        let mut batch_memory = 0;
        let max_batch_size = self.config.max_batch_size;
        let memory_limit = self.config.memory_limit.unwrap_or(usize::MAX);

        // Collect requests for batch - ensure we don't exceed max_batch_size
        while !self.pending_requests.is_empty() && batch_requests.len() < max_batch_size {
            let request = self
                .pending_requests
                .pop_front()
                .expect("checked is_empty before pop_front");
            let request_memory = request.request.estimated_memory_bytes();

            if batch_memory + request_memory <= memory_limit {
                batch_memory += request_memory;
                batch_requests.push(request.request);
            } else {
                // Put request back if it doesn't fit
                self.pending_requests.push_front(request);
                break;
            }
        }

        if batch_requests.is_empty() {
            return None;
        }

        // Calculate batch priority as highest priority among requests
        let batch_priority = batch_requests
            .iter()
            .map(|r| r.priority)
            .max()
            .unwrap_or(RequestPriority::Normal);

        // Estimate processing duration based on batch size and complexity
        let estimated_duration = Duration::from_millis(50 + (batch_requests.len() * 10) as u64);

        let batch = InferenceBatch {
            id: BatchId::new(),
            requests: batch_requests,
            lane_name: self.config.name.clone(),
            created_at: Instant::now(),
            estimated_memory_usage: batch_memory,
            priority: batch_priority,
            estimated_duration,
        };

        // Track batch as processing
        self.processing_batches.insert(batch.id, batch.clone());

        Some(batch)
    }

    fn form_batch_with_limit(&mut self, global_max_batch_size: usize) -> Option<InferenceBatch> {
        if self.pending_requests.is_empty() {
            return None;
        }

        let mut batch_requests = Vec::new();
        let mut batch_memory = 0;
        // Use the minimum of lane max batch size and global max batch size
        let max_batch_size = std::cmp::min(self.config.max_batch_size, global_max_batch_size);
        let memory_limit = self.config.memory_limit.unwrap_or(usize::MAX);

        // Collect requests for batch - ensure we don't exceed effective max_batch_size
        while !self.pending_requests.is_empty() && batch_requests.len() < max_batch_size {
            let request = self
                .pending_requests
                .pop_front()
                .expect("checked is_empty before pop_front");
            let request_memory = request.request.estimated_memory_bytes();

            if batch_memory + request_memory <= memory_limit {
                batch_memory += request_memory;
                batch_requests.push(request.request);
            } else {
                // Put request back if it doesn't fit
                self.pending_requests.push_front(request);
                break;
            }
        }

        if batch_requests.is_empty() {
            return None;
        }

        // Calculate batch priority as highest priority among requests
        let batch_priority = batch_requests
            .iter()
            .map(|r| r.priority)
            .max()
            .unwrap_or(RequestPriority::Normal);

        // Estimate processing duration based on batch size and complexity
        let estimated_duration = Duration::from_millis(50 + (batch_requests.len() * 10) as u64);

        let batch = InferenceBatch {
            id: BatchId::new(),
            requests: batch_requests,
            lane_name: self.config.name.clone(),
            created_at: Instant::now(),
            estimated_memory_usage: batch_memory,
            priority: batch_priority,
            estimated_duration,
        };

        // Track batch as processing
        self.processing_batches.insert(batch.id, batch.clone());

        Some(batch)
    }

    fn estimate_batch_memory(&self) -> usize {
        self.pending_requests
            .iter()
            .take(self.config.max_batch_size)
            .map(|req| req.request.estimated_memory_bytes())
            .sum()
    }

    fn calculate_memory_utilization(&self) -> f64 {
        // Calculate memory used by pending requests
        let pending_memory = self
            .pending_requests
            .iter()
            .map(|req| req.request.estimated_memory_bytes())
            .sum::<usize>();

        // Calculate memory used by processing batches
        let processing_memory = self
            .processing_batches
            .values()
            .map(|batch| {
                batch
                    .requests
                    .iter()
                    .map(super::request::InferenceRequest::estimated_memory_bytes)
                    .sum::<usize>()
            })
            .sum::<usize>();

        let total_memory_used = pending_memory + processing_memory;

        // Use lane-specific memory limit if available, otherwise use a reasonable default
        let memory_limit = self.config.memory_limit.unwrap_or(8_000_000_000); // Use configured limit or 8GB default

        if memory_limit > 0 {
            (total_memory_used as f64 / memory_limit as f64).min(1.0)
        } else {
            0.0
        }
    }

    fn get_stats(&self) -> LaneStats {
        let now = Instant::now();
        // If subtraction fails (process younger than 60s), use now as cutoff
        let recent_cutoff = now
            .checked_sub(Duration::from_secs(60))
            .unwrap_or_else(Instant::now);

        // Calculate recent metrics
        let recent_completions: Vec<_> = self
            .completed_requests
            .iter()
            .filter(|(completion_time, _)| *completion_time > recent_cutoff)
            .collect();

        let average_latency = if recent_completions.is_empty() {
            Duration::ZERO
        } else {
            let total_latency: Duration =
                recent_completions.iter().map(|(_, latency)| *latency).sum();
            total_latency / recent_completions.len() as u32
        };

        let violation_rate = if self.total_processed > 0 {
            self.violation_count as f64 / self.total_processed as f64
        } else {
            0.0
        };

        let throughput = recent_completions.len() as f64 / 60.0; // per second

        LaneStats {
            pending_count: self.pending_requests.len(),
            processing_count: self.processing_batches.len(),
            average_latency,
            violation_rate,
            throughput,
            memory_utilization: self.calculate_memory_utilization(),
        }
    }
}

/// Continuous batching scheduler
pub struct BatchScheduler {
    config: BatchSchedulerConfig,

    // Lane management
    lanes: Arc<RwLock<HashMap<String, LaneState>>>,

    // Global state
    current_memory_usage: Arc<RwLock<usize>>,
    degradation_mode: Arc<RwLock<bool>>,
    violation_history: Arc<RwLock<Vec<SlaViolation>>>,

    // Statistics
    total_requests_processed: Arc<RwLock<usize>>,
    last_stats_reset: Arc<RwLock<Instant>>,
}

impl BatchScheduler {
    /// Create a new batch scheduler
    pub async fn new(config: BatchSchedulerConfig) -> InferenceResult<Self> {
        if config.sla_lanes.is_empty() {
            return Err(InferenceError::invalid_request(
                "At least one SLA lane must be configured",
            ));
        }

        let mut lanes = HashMap::new();
        for lane_config in &config.sla_lanes {
            lanes.insert(
                lane_config.name.clone(),
                LaneState::new(lane_config.clone()),
            );
        }

        Ok(Self {
            config,
            lanes: Arc::new(RwLock::new(lanes)),
            current_memory_usage: Arc::new(RwLock::new(0)),
            degradation_mode: Arc::new(RwLock::new(false)),
            violation_history: Arc::new(RwLock::new(Vec::new())),
            total_requests_processed: Arc::new(RwLock::new(0)),
            last_stats_reset: Arc::new(RwLock::new(Instant::now())),
        })
    }

    /// Submit a request for scheduling
    pub async fn submit_request(&mut self, request: InferenceRequest) -> InferenceResult<()> {
        // Check queue size limits first
        let total_pending = {
            let lanes = self.lanes.read().await;
            lanes.values().map(LaneState::pending_count).sum::<usize>()
        };

        if total_pending >= self.config.max_queue_size {
            return Err(InferenceError::queue_full(
                self.config.max_queue_size,
                total_pending,
            ));
        }

        // Assign request to appropriate lane
        let lane_name = self.assign_to_lane(&request).await?;

        // Calculate priority score
        let priority_score = self.calculate_priority_score(&request);
        let request_id = request.id;

        let scheduler_request = SchedulerRequest {
            request,
            assigned_lane: lane_name.clone(),
            queued_at: Instant::now(),
            priority_score,
        };

        // Add to lane
        {
            let mut lanes = self.lanes.write().await;
            if let Some(lane) = lanes.get_mut(&lane_name) {
                lane.add_request(scheduler_request);
                debug!("Request {} assigned to lane {}", request_id, lane_name);
            } else {
                return Err(InferenceError::internal_error(
                    "Lane assignment",
                    "Lane not found",
                ));
            }
        }

        // Check if we should enter degradation mode (account for the request we just added)
        let new_total_pending = total_pending + 1;
        if new_total_pending >= self.config.degradation_mode_threshold {
            let mut degradation = self.degradation_mode.write().await;
            *degradation = true;
            warn!(
                "Entering degradation mode: {} pending requests",
                new_total_pending
            );
        }

        Ok(())
    }

    /// Get next batch for processing
    pub async fn get_next_batch(&mut self) -> InferenceResult<InferenceBatch> {
        let current_memory = *self.current_memory_usage.read().await;
        let global_max_batch_size = self.config.max_batch_size;
        let mut best_batch = None;
        let mut best_priority = RequestPriority::Low;

        {
            let mut lanes = self.lanes.write().await;

            for (_lane_name, lane_state) in lanes.iter_mut() {
                if lane_state.should_form_batch(current_memory, self.config.memory_limit)
                    && let Some(batch) = lane_state.form_batch_with_limit(global_max_batch_size)
                {
                    // Choose batch with highest priority
                    if batch.priority >= best_priority {
                        best_priority = batch.priority;
                        best_batch = Some(batch);
                    }
                }
            }
        }

        best_batch.ok_or_else(|| InferenceError::BatchFormationFailed {
            reason: "No eligible batches available".to_string(),
        })
    }

    /// Mark batch as executing
    pub async fn mark_batch_executing(&mut self, batch_id: BatchId) -> InferenceResult<()> {
        let lanes = self.lanes.read().await;
        for lane in lanes.values() {
            if lane.processing_batches.contains_key(&batch_id) {
                trace!("Batch {} marked as executing", batch_id);
                return Ok(());
            }
        }

        Err(InferenceError::internal_error(
            "Mark batch executing",
            "Batch not found",
        ))
    }

    /// Evaluate preemption decisions
    pub async fn evaluate_preemption(&self) -> InferenceResult<PreemptionDecision> {
        if !self.config.enable_preemption {
            return Ok(PreemptionDecision {
                should_preempt: false,
                target_batch_id: None,
                reason: "Preemption disabled".to_string(),
                priority_delta: 0,
            });
        }

        let lanes = self.lanes.read().await;

        // Find highest priority pending request
        let mut highest_pending_priority = RequestPriority::Low;
        let mut pending_deadline = None;

        for lane in lanes.values() {
            if let Some(request) = lane.pending_requests.front()
                && request.request.priority > highest_pending_priority
            {
                highest_pending_priority = request.request.priority;
                pending_deadline = request.request.deadline;
            }
        }

        // Find lowest priority executing batch
        let mut lowest_executing_priority = RequestPriority::Critical;
        let mut target_batch_id = None;

        for lane in lanes.values() {
            for batch in lane.processing_batches.values() {
                if batch.priority < lowest_executing_priority {
                    lowest_executing_priority = batch.priority;
                    target_batch_id = Some(batch.id);
                }
            }
        }

        // Check if preemption is justified
        let priority_delta = highest_pending_priority as i32 - lowest_executing_priority as i32;
        let should_preempt = priority_delta >= 2; // Preempt if 2+ priority levels difference

        // Also check deadline urgency
        let urgent_deadline = pending_deadline
            .is_some_and(|deadline| deadline - Instant::now() < Duration::from_millis(50));

        let final_decision = should_preempt || urgent_deadline;

        Ok(PreemptionDecision {
            should_preempt: final_decision,
            target_batch_id,
            reason: if urgent_deadline {
                "Urgent deadline approaching".to_string()
            } else if should_preempt {
                format!("Priority difference: {priority_delta}")
            } else {
                "No preemption needed".to_string()
            },
            priority_delta,
        })
    }

    /// Check for SLA violations
    pub async fn check_sla_violations(&self) -> InferenceResult<Vec<SlaViolation>> {
        let mut violations = Vec::new();
        let now = Instant::now();

        let lanes = self.lanes.read().await;
        for lane in lanes.values() {
            for request in &lane.pending_requests {
                let age = now - request.queued_at;

                // Check against lane max latency
                if age > lane.config.max_latency {
                    violations.push(SlaViolation {
                        request_id: request.request.id,
                        lane_name: lane.config.name.clone(),
                        violation_type: "Queue time exceeded".to_string(),
                        violation_time: now,
                        expected_latency: lane.config.max_latency,
                        actual_latency: age,
                    });
                }

                // Check against request-specific deadline
                if let Some(deadline) = request.request.deadline
                    && now > deadline
                {
                    let expected_duration = deadline - request.queued_at;
                    violations.push(SlaViolation {
                        request_id: request.request.id,
                        lane_name: lane.config.name.clone(),
                        violation_type: "Deadline exceeded".to_string(),
                        violation_time: now,
                        expected_latency: expected_duration,
                        actual_latency: age,
                    });
                }
            }
        }

        // Store violations in history
        if !violations.is_empty() {
            let mut history = self.violation_history.write().await;
            history.extend(violations.clone());

            // Keep only recent violations (last hour)
            // If subtraction fails (process younger than 1 hour), use now as cutoff
            let cutoff = now
                .checked_sub(Duration::from_secs(3600))
                .unwrap_or_else(Instant::now);
            history.retain(|v| v.violation_time > cutoff);
        }

        Ok(violations)
    }

    /// Set memory pressure level
    pub async fn set_memory_pressure(&mut self, pressure: f64) -> InferenceResult<()> {
        let memory_bytes = (pressure * self.config.memory_limit as f64) as usize;
        let mut current_memory = self.current_memory_usage.write().await;
        *current_memory = memory_bytes;

        if pressure > self.config.memory_pressure_threshold {
            warn!("High memory pressure: {:.1}%", pressure * 100.0);
        }

        Ok(())
    }

    /// Check if scheduler is in degradation mode
    pub async fn is_degradation_mode(&self) -> bool {
        *self.degradation_mode.read().await
    }

    /// Get queue statistics
    pub async fn queue_stats(&self) -> QueueStats {
        let lanes_guard = self.lanes.read().await;
        let mut lane_stats = HashMap::new();
        let mut total_pending = 0;
        let mut total_processing = 0;

        for (name, lane) in lanes_guard.iter() {
            let stats = lane.get_stats();
            total_pending += stats.pending_count;
            total_processing += stats.processing_count;
            lane_stats.insert(name.clone(), stats);
        }

        let memory_usage = *self.current_memory_usage.read().await;
        let memory_pressure = memory_usage as f64 / self.config.memory_limit as f64;

        QueueStats {
            total_pending,
            total_processing,
            lanes: lane_stats,
            queue_utilization: total_pending as f64 / self.config.max_queue_size as f64,
            memory_pressure,
        }
    }

    /// Get scheduler configuration
    #[must_use]
    pub fn config(&self) -> &BatchSchedulerConfig {
        &self.config
    }

    /// Assign request to appropriate SLA lane
    async fn assign_to_lane(&self, request: &InferenceRequest) -> InferenceResult<String> {
        // Find matching lane based on priority
        for lane_config in &self.config.sla_lanes {
            if lane_config.priority == request.priority {
                return Ok(lane_config.name.clone());
            }
        }

        // Fallback to lane with closest priority
        let mut best_lane = &self.config.sla_lanes[0];
        let mut best_diff = i32::MAX;

        for lane_config in &self.config.sla_lanes {
            let diff = (lane_config.priority as i32 - request.priority as i32).abs();
            if diff < best_diff {
                best_diff = diff;
                best_lane = lane_config;
            }
        }

        Ok(best_lane.name.clone())
    }

    /// Calculate priority score for request ordering
    fn calculate_priority_score(&self, request: &InferenceRequest) -> i32 {
        let mut score = request.priority as i32 * 1000;

        // Boost score for urgent deadlines
        if let Some(deadline) = request.deadline {
            let time_remaining = deadline.saturating_duration_since(Instant::now());
            if time_remaining < Duration::from_millis(100) {
                score += 500; // High urgency boost
            } else if time_remaining < Duration::from_secs(1) {
                score += 200; // Medium urgency boost
            }
        }

        // Penalize very large requests slightly
        if request.input_tokens.len() > 1024 {
            score -= 50;
        }

        score
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_scheduler_creation() {
        let config = BatchSchedulerConfig::default();
        let scheduler = BatchScheduler::new(config).await;
        assert!(scheduler.is_ok());
    }

    #[tokio::test]
    async fn test_lane_assignment() {
        let config = BatchSchedulerConfig {
            sla_lanes: vec![
                SlaLane {
                    name: "high".to_string(),
                    priority: RequestPriority::High,
                    max_latency: Duration::from_millis(50),
                    max_batch_size: 8,
                    memory_limit: None,
                    min_wait_time: None,
                },
                SlaLane {
                    name: "normal".to_string(),
                    priority: RequestPriority::Normal,
                    max_latency: Duration::from_millis(100),
                    max_batch_size: 16,
                    memory_limit: None,
                    min_wait_time: None,
                },
            ],
            ..Default::default()
        };

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

        let high_request = InferenceRequest {
            priority: RequestPriority::High,
            input_tokens: vec![1, 2, 3],
            ..Default::default()
        };

        let lane = scheduler.assign_to_lane(&high_request).await.unwrap();
        assert_eq!(lane, "high");
    }

    #[test]
    fn test_priority_score_calculation() {
        let scheduler_config = BatchSchedulerConfig::default();
        let scheduler = BatchScheduler {
            config: scheduler_config,
            lanes: Arc::new(RwLock::new(HashMap::new())),
            current_memory_usage: Arc::new(RwLock::new(0)),
            degradation_mode: Arc::new(RwLock::new(false)),
            violation_history: Arc::new(RwLock::new(Vec::new())),
            total_requests_processed: Arc::new(RwLock::new(0)),
            last_stats_reset: Arc::new(RwLock::new(Instant::now())),
        };

        let high_priority_request = InferenceRequest {
            priority: RequestPriority::High,
            ..Default::default()
        };

        let normal_priority_request = InferenceRequest {
            priority: RequestPriority::Normal,
            ..Default::default()
        };

        let high_score = scheduler.calculate_priority_score(&high_priority_request);
        let normal_score = scheduler.calculate_priority_score(&normal_priority_request);

        assert!(high_score > normal_score);
    }
}