//! High-performance experience replay buffer with in-memory storage

use std::collections::HashSet;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};

use parking_lot::RwLock;
use rand::prelude::*;

use super::{Experience, ExperienceBatch};
use ml_core::MLError;

/// Configuration for replay buffer
#[derive(Debug, Clone)]
pub struct ReplayBufferConfig {
    /// Maximum number of experiences to store
    pub capacity: usize,
    /// Batch size for sampling
    pub batch_size: usize,
    /// Minimum experiences before sampling
    pub min_experiences: usize,
}

impl Default for ReplayBufferConfig {
    fn default() -> Self {
        Self {
            capacity: 1_000_000,
            batch_size: 32,
            min_experiences: 1000,
        }
    }
}

/// Statistics about the replay buffer
#[derive(Debug, Clone)]
pub struct ReplayBufferStats {
    /// Current number of experiences stored
    pub size: usize,
    /// Total capacity of the buffer
    pub capacity: usize,
    /// Number of samples taken
    pub samples_taken: u64,
    /// Number of experiences added
    pub experiences_added: u64,
}

/// High-performance in-memory experience replay buffer
#[derive(Debug)]
pub struct ReplayBuffer {
    /// Buffer configuration
    config: ReplayBufferConfig,
    /// Circular buffer of experiences
    buffer: RwLock<Vec<Option<Experience>>>,
    /// Current write position
    write_pos: AtomicUsize,
    /// Current size of buffer
    size: AtomicUsize,
    /// Statistics counters
    samples_taken: AtomicU64,
    experiences_added: AtomicU64,
}

impl ReplayBuffer {
    /// Create a new replay buffer
    pub fn new(config: ReplayBufferConfig) -> Result<Self, MLError> {
        let buffer = vec![None; config.capacity];

        Ok(Self {
            config,
            buffer: RwLock::new(buffer),
            write_pos: AtomicUsize::new(0),
            size: AtomicUsize::new(0),
            samples_taken: AtomicU64::new(0),
            experiences_added: AtomicU64::new(0),
        })
    }

    /// Add an experience to the buffer
    pub fn push(&self, experience: Experience) -> Result<(), MLError> {
        let mut buffer = self.buffer.write();
        let pos = self.write_pos.load(Ordering::Relaxed);

        buffer[pos] = Some(experience);

        // Update position (circular)
        let new_pos = (pos + 1) % self.config.capacity;
        self.write_pos.store(new_pos, Ordering::Relaxed);

        // Update size (max is capacity)
        let current_size = self.size.load(Ordering::Relaxed);
        if current_size < self.config.capacity {
            self.size.store(current_size + 1, Ordering::Relaxed);
        }

        self.experiences_added.fetch_add(1, Ordering::Relaxed);

        Ok(())
    }

    /// Sample a batch of experiences
    pub fn sample(&self, batch_size: Option<usize>) -> Result<ExperienceBatch, MLError> {
        let batch_size = batch_size.unwrap_or(self.config.batch_size);
        let current_size = self.size.load(Ordering::Relaxed);

        if current_size < self.config.min_experiences {
            return Err(MLError::InvalidInput(format!(
                "Not enough experiences: {} < {}",
                current_size, self.config.min_experiences
            )));
        }

        if batch_size > current_size {
            return Err(MLError::InvalidInput(format!(
                "Batch size {} exceeds buffer size {}",
                batch_size, current_size
            )));
        }

        let buffer = self.buffer.read();
        let mut experiences = Vec::with_capacity(batch_size);

        // Efficient random sampling: generate batch_size random indices directly
        // instead of allocating and shuffling an array of current_size elements
        let mut rng = thread_rng();
        let mut selected = HashSet::with_capacity(batch_size);
        while selected.len() < batch_size {
            selected.insert(rng.gen_range(0..current_size));
        }

        for idx in &selected {
            if let Some(Some(experience)) = buffer.get(*idx) {
                experiences.push(experience.clone());
            }
        }

        self.samples_taken.fetch_add(1, Ordering::Relaxed);

        Ok(ExperienceBatch::new(experiences))
    }

    /// Get buffer statistics
    pub fn stats(&self) -> ReplayBufferStats {
        ReplayBufferStats {
            size: self.size.load(Ordering::Relaxed),
            capacity: self.config.capacity,
            samples_taken: self.samples_taken.load(Ordering::Relaxed),
            experiences_added: self.experiences_added.load(Ordering::Relaxed),
        }
    }

    /// Get current size
    pub fn size(&self) -> usize {
        self.size.load(Ordering::Relaxed)
    }

    /// Check if buffer can be sampled
    pub fn can_sample(&self) -> bool {
        self.size.load(Ordering::Relaxed) >= self.config.min_experiences
    }

    /// Resize buffer capacity (adaptive buffer sizing)
    ///
    /// Keeps most recent experiences, drops oldest if shrinking.
    ///
    /// # Arguments
    ///
    /// * `new_capacity` - New buffer capacity
    ///
    /// # Returns
    ///
    /// * `Ok(())` - Resize successful
    /// * `Err` - Resize failed
    pub fn resize(&mut self, new_capacity: usize) -> Result<(), MLError> {
        if new_capacity == self.config.capacity {
            return Ok(()); // No change needed
        }
        
        let mut buffer = self.buffer.write();
        let current_size = self.size.load(Ordering::Relaxed);
        
        // Create new buffer with updated capacity
        let mut new_buffer = vec![None; new_capacity];
        
        // Transfer experiences (keep most recent if shrinking)
        let to_keep = current_size.min(new_capacity);
        let start = current_size.saturating_sub(new_capacity);
        
        for i in 0..to_keep {
            new_buffer[i] = buffer[(start + i) % self.config.capacity].clone();
        }
        
        *buffer = new_buffer;
        self.config.capacity = new_capacity;
        self.write_pos.store(to_keep % new_capacity, Ordering::Relaxed);
        self.size.store(to_keep, Ordering::Relaxed);
        
        Ok(())
    }
    
    /// Get current buffer capacity
    pub const fn capacity(&self) -> usize {
        self.config.capacity
    }
}

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


    fn create_test_experience(reward: f32) -> Experience {
        Experience::new(vec![1.0, 2.0, 3.0], 1, reward, vec![1.1, 2.1, 3.1], false)
    }

    #[test]
    fn test_replay_buffer_creation() -> Result<(), Box<dyn std::error::Error>> {
        let config = ReplayBufferConfig::default();

        let buffer = ReplayBuffer::new(config)?;
        let stats = buffer.stats();

        assert_eq!(stats.size, 0);
        assert_eq!(stats.capacity, 1_000_000);
        Ok(())
    }

    #[test]
    fn test_experience_storage() -> Result<(), Box<dyn std::error::Error>> {
        let mut config = ReplayBufferConfig::default();
        config.capacity = 100;

        let buffer = ReplayBuffer::new(config)?;

        // Add experiences
        for i in 0..10 {
            let experience = create_test_experience(i as f32 * 0.1);
            buffer.push(experience)?;
        }

        let stats = buffer.stats();
        assert_eq!(stats.size, 10);
        Ok(())
    }

    #[test]
    fn test_batch_sampling() -> Result<(), Box<dyn std::error::Error>> {
        let mut config = ReplayBufferConfig::default();
        config.capacity = 100;
        config.batch_size = 5;
        config.min_experiences = 10; // Lower threshold for testing

        let buffer = ReplayBuffer::new(config)?;

        // Add experiences
        for i in 0..20 {
            let experience = create_test_experience(i as f32 * 0.1);
            buffer.push(experience)?;
        }

        let batch = buffer.sample(Some(5))?;
        assert_eq!(batch.batch_size, 5);
        assert!(batch.is_valid());
        Ok(())
    }
}
