
/// Statistics for dynamic priority scheduling
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct DynamicSchedulerStats {
    /// Total requests processed
    pub total_requests: u64,
    /// Requests completed
    pub completed_requests: u64,
    /// Requests that met SLA
    pub sla_met: u64,
    /// Requests that missed SLA
    pub sla_missed: u64,
    /// Requests dropped (hard deadline exceeded)
    pub dropped_requests: u64,
    /// Total priority promotions
    pub promotions: u64,
    /// Average time-to-first-token (ms)
    pub avg_ttft_ms: f64,
    /// p99 time-to-first-token (ms)
    pub p99_ttft_ms: f64,
    /// Tokens allocated per priority level
    pub tokens_by_priority: [u64; 4],
    /// Current queue depth per priority
    pub queue_depth_by_priority: [usize; 4],
}

/// Dynamic batch priority scheduler
///
/// Implements advanced priority scheduling with:
/// - Age-based priority promotion (MLFQ-style)
/// - Deadline-aware scheduling for SLA compliance
/// - Fair share token budget allocation
/// - Urgency-based boosting for time-sensitive requests
pub struct DynamicPriorityScheduler {
    /// Configuration
    config: DynamicPriorityConfig,
    /// All requests by ID
    requests: HashMap<u64, DynamicRequest>,
    /// Priority queues (one per level)
    priority_queues: [VecDeque<u64>; 4],
    /// Running requests
    running: Vec<u64>,
    /// Next request ID
    next_request_id: u64,
    /// Statistics
    stats: DynamicSchedulerStats,
    /// TTFT samples for percentile calculation
    ttft_samples: Vec<f64>,
    /// Total batch token budget
    batch_token_budget: usize,
}

impl DynamicPriorityScheduler {
    /// Create a new dynamic priority scheduler
    pub fn new(batch_token_budget: usize) -> Self {
        Self::with_config(batch_token_budget, DynamicPriorityConfig::default())
    }

    /// Create with custom configuration
    pub fn with_config(batch_token_budget: usize, config: DynamicPriorityConfig) -> Self {
        Self {
            config,
            requests: HashMap::new(),
            priority_queues: [
                VecDeque::new(),
                VecDeque::new(),
                VecDeque::new(),
                VecDeque::new(),
            ],
            running: Vec::new(),
            next_request_id: 0,
            stats: DynamicSchedulerStats::default(),
            ttft_samples: Vec::new(),
            batch_token_budget,
        }
    }

    /// Add a request with priority and optional deadline
    pub fn add_request(
        &mut self,
        input_ids: Vec<u32>,
        max_tokens: usize,
        priority: Priority,
        deadline: Option<Deadline>,
    ) -> u64 {
        let request_id = self.next_request_id;
        self.next_request_id += 1;

        let mut request =
            DynamicRequest::new(request_id, input_ids, max_tokens).with_priority(priority);
        if let Some(d) = deadline {
            request = request.with_deadline(d);
        }

        // Add to appropriate priority queue
        let queue_idx = priority as usize;
        self.priority_queues[queue_idx].push_back(request_id);
        self.requests.insert(request_id, request);

        self.stats.total_requests += 1;
        self.update_queue_depths();

        request_id
    }

    /// Add a simple request (Normal priority, no deadline)
    pub fn add_simple_request(&mut self, input_ids: Vec<u32>, max_tokens: usize) -> u64 {
        self.add_request(input_ids, max_tokens, Priority::Normal, None)
    }

    /// Perform age-based priority promotion
    pub fn promote_aged_requests(&mut self) {
        if !self.config.enable_age_promotion {
            return;
        }

        let promotion_threshold = self.config.promotion_interval_ms;
        let max_priority = self.config.max_promoted_priority;

        // Check each queue except Critical (can't promote beyond Critical)
        const QUEUE_PRIORITIES: [(usize, Priority); 3] = [
            (0, Priority::Low),
            (1, Priority::Normal),
            (2, Priority::High),
        ];
        for &(queue_idx, current_priority) in &QUEUE_PRIORITIES {
            if current_priority >= max_priority {
                continue;
            }
            let to_promote = self.find_promotable_in_queue(queue_idx, promotion_threshold);
            for request_id in to_promote {
                self.promote_request(request_id);
            }
        }
    }

    /// Scan a priority queue and return the ids of requests whose wait time
    /// has crossed the next promotion threshold.
    fn find_promotable_in_queue(&self, queue_idx: usize, promotion_threshold: u64) -> Vec<u64> {
        let mut to_promote = Vec::new();
        for &request_id in &self.priority_queues[queue_idx] {
            if let Some(request) = self.requests.get(&request_id) {
                let promotions_time = promotion_threshold * (request.promotions as u64 + 1);
                if request.wait_time_ms() >= promotions_time {
                    to_promote.push(request_id);
                }
            }
        }
        to_promote
    }

    /// Promote a single request to next priority level
    fn promote_request(&mut self, request_id: u64) {
        if let Some(request) = self.requests.get_mut(&request_id) {
            let current_idx = request.effective_priority as usize;
            let max_idx = self.config.max_promoted_priority as usize;

            if current_idx < max_idx {
                // Remove from current queue
                self.priority_queues[current_idx].retain(|&id| id != request_id);

                // Promote
                let new_priority = match current_idx + 1 {
                    1 => Priority::Normal,
                    2 => Priority::High,
                    3 => Priority::Critical,
                    _ => return,
                };
                request.effective_priority = new_priority;
                request.promotions += 1;

                // Add to new queue (at front since it's been waiting)
                self.priority_queues[current_idx + 1].push_front(request_id);
                self.stats.promotions += 1;
            }
        }
    }

    /// Drop expired requests (hard deadline exceeded)
    pub fn drop_expired(&mut self) -> Vec<u64> {
        let mut dropped = Vec::new();

        for queue in &mut self.priority_queues {
            let mut to_remove = Vec::new();
            for &request_id in queue.iter() {
                if let Some(request) = self.requests.get(&request_id) {
                    if request.is_expired() {
                        to_remove.push(request_id);
                    }
                }
            }

            for request_id in to_remove {
                queue.retain(|&id| id != request_id);
                if let Some(mut request) = self.requests.remove(&request_id) {
                    request.state = SequenceState::Failed;
                    dropped.push(request_id);
                    self.stats.dropped_requests += 1;
                }
            }
        }

        self.update_queue_depths();
        dropped
    }

    /// Schedule requests using dynamic priority and token budgets
    ///
    /// Returns (scheduled_request_ids, tokens_allocated_per_request)
    pub fn schedule(&mut self, available_slots: usize) -> Vec<(u64, usize)> {
        self.promote_aged_requests();
        self.drop_expired();

        let mut scheduled = Vec::new();
        let mut remaining_budget = self.batch_token_budget;
        let mut remaining_slots = available_slots;
        let budgets = self.priority_budgets();

        // Schedule from highest priority to lowest
        for queue_idx in (0..4).rev() {
            if remaining_slots == 0 || remaining_budget == 0 {
                break;
            }
            let priority_budget = budgets[queue_idx].min(remaining_budget);
            self.sort_priority_queue_by_urgency(queue_idx);
            let scheduled_ids = self.allocate_from_priority(
                queue_idx,
                priority_budget,
                &mut remaining_slots,
                &mut remaining_budget,
                &mut scheduled,
            );
            self.promote_scheduled_to_running(queue_idx, &scheduled_ids);
        }

        self.update_queue_depths();
        scheduled
    }

    /// Compute per-priority token budgets for a single schedule tick.
    /// Fair-share mode splits the batch budget via `priority_budgets` ratios;
    /// otherwise each priority may claim the full batch budget.
    fn priority_budgets(&self) -> [usize; 4] {
        if self.config.enable_fair_share {
            self.config
                .priority_budgets
                .map(|b| (b * self.batch_token_budget as f64) as usize)
        } else {
            [self.batch_token_budget; 4]
        }
    }

    /// Reorder a priority queue by urgency score (descending) when
    /// deadline-aware scheduling is enabled.
    fn sort_priority_queue_by_urgency(&mut self, queue_idx: usize) {
        if !self.config.enable_deadline_scheduling {
            return;
        }
        let mut sorted: Vec<_> = self.priority_queues[queue_idx].iter().copied().collect();
        sorted.sort_by(|&a, &b| {
            match (self.requests.get(&a), self.requests.get(&b)) {
                (Some(ra), Some(rb)) => rb
                    .urgency_score()
                    .partial_cmp(&ra.urgency_score())
                    .unwrap_or(std::cmp::Ordering::Equal),
                _ => std::cmp::Ordering::Equal,
            }
        });
        self.priority_queues[queue_idx] = sorted.into_iter().collect();
    }

    /// Walk a priority queue and allocate tokens to requests, appending
    /// `(id, tokens)` pairs to `scheduled`. Returns the ids scheduled from
    /// this queue so callers can promote them to running.
    fn allocate_from_priority(
        &mut self,
        queue_idx: usize,
        mut priority_budget: usize,
        remaining_slots: &mut usize,
        remaining_budget: &mut usize,
        scheduled: &mut Vec<(u64, usize)>,
    ) -> Vec<u64> {
        let min_tokens = self.config.min_tokens_per_request;
        let mut scheduled_from_queue = Vec::new();
        let queue_snapshot: Vec<u64> = self.priority_queues[queue_idx].iter().copied().collect();

        for request_id in queue_snapshot {
            if *remaining_slots == 0 || priority_budget < min_tokens {
                break;
            }
            let Some(tokens_to_allocate) =
                self.tokens_to_allocate_for(request_id, priority_budget, min_tokens)
            else {
                continue;
            };

            scheduled.push((request_id, tokens_to_allocate));
            scheduled_from_queue.push(request_id);
            priority_budget = priority_budget.saturating_sub(tokens_to_allocate);
            *remaining_budget = remaining_budget.saturating_sub(tokens_to_allocate);
            *remaining_slots -= 1;
            self.stats.tokens_by_priority[queue_idx] += tokens_to_allocate as u64;
        }
        scheduled_from_queue
    }

    /// Compute tokens to allocate to a single request, clamped to the priority
    /// budget and `min_tokens_per_request` floor. Returns `None` if the
    /// request id is no longer present.
    fn tokens_to_allocate_for(
        &self,
        request_id: u64,
        priority_budget: usize,
        min_tokens: usize,
    ) -> Option<usize> {
        let request = self.requests.get(&request_id)?;
        let tokens_needed = request.remaining_tokens().max(1);
        let allocated = tokens_needed.min(priority_budget).max(min_tokens);
        (allocated > 0).then_some(allocated)
    }

    /// Move scheduled ids out of their priority queue into `running` and
    /// record TTFT on first promotion.
    fn promote_scheduled_to_running(&mut self, queue_idx: usize, scheduled_ids: &[u64]) {
        for &request_id in scheduled_ids {
            self.priority_queues[queue_idx].retain(|&id| id != request_id);
            let Some(request) = self.requests.get_mut(&request_id) else {
                continue;
            };
            request.state = SequenceState::Running;
            self.running.push(request_id);
            if request.ttft_ms.is_none() {
                let ttft = request.wait_time_ms() as f64;
                request.ttft_ms = Some(ttft);
                self.ttft_samples.push(ttft);
            }
        }
    }

    /// Complete a request and update statistics
    pub fn complete_request(&mut self, request_id: u64) -> Option<DynamicRequest> {
        // Remove from running
        self.running.retain(|&id| id != request_id);

        if let Some(mut request) = self.requests.remove(&request_id) {
            request.state = SequenceState::Completed;
            self.stats.completed_requests += 1;

            // Check SLA compliance
            if let Some(deadline) = &request.deadline {
                let elapsed = request.wait_time_ms();
                if elapsed <= deadline.target_latency_ms {
                    self.stats.sla_met += 1;
                } else {
                    self.stats.sla_missed += 1;
                }
            }

            // Update average TTFT
            self.update_ttft_stats();

            Some(request)
        } else {
            None
        }
    }

    /// Update TTFT statistics
    fn update_ttft_stats(&mut self) {
        if self.ttft_samples.is_empty() {
            return;
        }

        // Average
        let sum: f64 = self.ttft_samples.iter().sum();
        self.stats.avg_ttft_ms = sum / self.ttft_samples.len() as f64;

        // P99
        let mut sorted = self.ttft_samples.clone();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        let p99_idx = ((sorted.len() as f64) * 0.99) as usize;
        self.stats.p99_ttft_ms = sorted
            .get(p99_idx.min(sorted.len() - 1))
            .copied()
            .unwrap_or(0.0);
    }

    /// Update queue depth statistics
    fn update_queue_depths(&mut self) {
        for (i, queue) in self.priority_queues.iter().enumerate() {
            self.stats.queue_depth_by_priority[i] = queue.len();
        }
    }

    /// Get a request by ID
    pub fn get_request(&self, request_id: u64) -> Option<&DynamicRequest> {
        self.requests.get(&request_id)
    }

    /// Get statistics
    pub fn stats(&self) -> &DynamicSchedulerStats {
        &self.stats
    }

    /// Get configuration
    pub fn config(&self) -> &DynamicPriorityConfig {
        &self.config
    }

    /// Total waiting requests
    pub fn waiting_count(&self) -> usize {
        self.priority_queues.iter().map(VecDeque::len).sum()
    }

    /// Running requests
    pub fn running_count(&self) -> usize {
        self.running.len()
    }

    /// SLA compliance rate (0.0 to 1.0)
    pub fn sla_compliance_rate(&self) -> f64 {
        let total = self.stats.sla_met + self.stats.sla_missed;
        if total == 0 {
            1.0
        } else {
            self.stats.sla_met as f64 / total as f64
        }
    }

    /// Get queue depth for a priority level
    pub fn queue_depth(&self, priority: Priority) -> usize {
        self.priority_queues[priority as usize].len()
    }
}

// ============================================================================
// Tests
// ============================================================================

// Tests extracted to tests.rs (PMAT-802)
#[cfg(test)]
#[path = "tests.rs"]
mod scheduler_tests;

// Additional tests for coverage (Part 02)
#[cfg(test)]
#[path = "tests_dynamic_request.rs"]
mod scheduler_tests_part_02;

// Deep coverage tests (Part 03)
#[cfg(test)]
#[path = "tests_sequence_batch.rs"]
mod scheduler_tests_part_03;
