//! Priority staleness tracking for PER with N-step returns
//!
//! When using N-step returns, TD errors computed at insertion time become stale
//! as the network trains. This module tracks per-slot update timestamps and
//! provides decay/refresh mechanisms to keep priorities fresh.

use std::sync::Arc;

use parking_lot::Mutex;

/// Tracks when each slot's priority was last updated.
///
/// Used by `PrioritizedReplayBuffer` to identify stale priorities
/// that should be refreshed by recomputing TD errors with the current network.
#[derive(Debug)]
pub struct StalenessTracker {
    /// Step at which each slot's priority was last updated (0 = never sampled/updated)
    last_updated: Arc<Mutex<Vec<usize>>>,
    /// Current training step (mirrors `PrioritizedReplayBuffer::training_step`)
    capacity: usize,
}

impl StalenessTracker {
    /// Create a new staleness tracker for a buffer of the given capacity.
    pub fn new(capacity: usize) -> Self {
        Self {
            last_updated: Arc::new(Mutex::new(vec![0; capacity])),
            capacity,
        }
    }

    /// Mark a slot as freshly updated at the given training step.
    pub fn mark_updated(&self, index: usize, step: usize) {
        if index < self.capacity {
            let mut last = self.last_updated.lock();
            if let Some(slot) = last.get_mut(index) {
                *slot = step;
            }
        }
    }

    /// Mark multiple slots as freshly updated.
    pub fn mark_updated_batch(&self, indices: &[usize], step: usize) {
        let mut last = self.last_updated.lock();
        for &idx in indices {
            if idx < self.capacity {
                if let Some(s) = last.get_mut(idx) {
                    *s = step;
                }
            }
        }
    }

    /// Get indices of stale experiences (priority not updated for > `max_age` steps).
    ///
    /// Returns up to `limit` indices sorted by staleness (oldest first).
    /// Only considers indices in `[0, buffer_size)`.
    pub fn get_stale_indices(
        &self,
        current_step: usize,
        max_age: usize,
        buffer_size: usize,
        limit: usize,
    ) -> Vec<usize> {
        let last = self.last_updated.lock();
        let effective_size = buffer_size.min(self.capacity);

        // Collect (index, age) pairs for stale entries
        let mut stale: Vec<(usize, usize)> = (0..effective_size)
            .filter_map(|i| {
                let last_step = last.get(i).copied().unwrap_or(0);
                let age = current_step.saturating_sub(last_step);
                (age > max_age).then_some((i, age))
            })
            .collect();

        // Sort by age descending (oldest first = highest priority for refresh)
        stale.sort_unstable_by(|a, b| b.1.cmp(&a.1));
        stale.truncate(limit);
        stale.into_iter().map(|(idx, _)| idx).collect()
    }

    /// Reset all timestamps (e.g., on buffer clear).
    pub fn reset(&self) {
        let mut last = self.last_updated.lock();
        last.fill(0);
    }
}

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

    #[test]
    fn test_staleness_tracker_new() {
        let tracker = StalenessTracker::new(100);
        // All slots should start at step 0
        let stale = tracker.get_stale_indices(10, 5, 100, 200);
        assert_eq!(stale.len(), 100, "All slots should be stale at step 10 with max_age=5");
    }

    #[test]
    fn test_mark_updated_makes_fresh() {
        let tracker = StalenessTracker::new(10);
        tracker.mark_updated(3, 100);
        // At step 105 with max_age=10, slot 3 is fresh (age=5 <= 10)
        let stale = tracker.get_stale_indices(105, 10, 10, 100);
        assert!(!stale.contains(&3), "Slot 3 should not be stale");
        // Other slots should be stale (age=105 > 10)
        assert_eq!(stale.len(), 9);
    }

    #[test]
    fn test_batch_update() {
        let tracker = StalenessTracker::new(10);
        tracker.mark_updated_batch(&[0, 2, 4, 6, 8], 50);
        let stale = tracker.get_stale_indices(55, 10, 10, 100);
        // Slots 0,2,4,6,8 have age=5 (fresh), slots 1,3,5,7,9 have age=55 (stale)
        assert_eq!(stale.len(), 5);
        assert!(stale.contains(&1));
        assert!(stale.contains(&3));
    }

    #[test]
    fn test_limit_caps_results() {
        let tracker = StalenessTracker::new(100);
        let stale = tracker.get_stale_indices(100, 5, 100, 10);
        assert_eq!(stale.len(), 10, "Should be capped at limit=10");
    }

    #[test]
    fn test_oldest_first_ordering() {
        let tracker = StalenessTracker::new(5);
        tracker.mark_updated(0, 10);
        tracker.mark_updated(1, 5);
        tracker.mark_updated(2, 20);
        // At step 30: ages are [20, 25, 10, 30, 30]
        let stale = tracker.get_stale_indices(30, 0, 5, 5);
        // slot 3 and 4 (age=30), then 1 (25), then 0 (20), then 2 (10)
        assert_eq!(stale[0], 3); // age 30 (or 4, both are 30)
    }

    #[test]
    fn test_respects_buffer_size() {
        let tracker = StalenessTracker::new(100);
        // Only 20 slots are occupied
        let stale = tracker.get_stale_indices(100, 5, 20, 200);
        assert_eq!(stale.len(), 20, "Should only check occupied slots");
    }

    #[test]
    fn test_reset_clears_timestamps() {
        let tracker = StalenessTracker::new(10);
        tracker.mark_updated_batch(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 100);
        tracker.reset();
        let stale = tracker.get_stale_indices(10, 5, 10, 100);
        assert_eq!(stale.len(), 10, "After reset, all should be stale");
    }
}
