//! WAL Reader for crash recovery and replay
//!
//! Reads WAL files and provides iterators for replaying entries.

use std::path::{Path, PathBuf};
use std::io::{self, Read, BufReader};
use std::fs::File;
use std::time::Instant;

use rayon::prelude::*;

use super::format::{WALEntry, WALHeader, WALFooter};

/// WAL Reader for reading and replaying WAL files
pub struct WALReader {
    /// Directory containing WAL files
    path: PathBuf,
    
    /// Shard ID filter (None = all shards)
    shard_id: Option<usize>,
}

impl WALReader {
    /// Create a new WAL reader
    pub fn new(path: impl Into<PathBuf>) -> Self {
        Self {
            path: path.into(),
            shard_id: None,
        }
    }
    
    /// Filter by shard ID
    pub fn for_shard(mut self, shard_id: usize) -> Self {
        self.shard_id = Some(shard_id);
        self
    }
    
    /// List all WAL files in order
    pub fn list_files(&self) -> io::Result<Vec<PathBuf>> {
        let mut files: Vec<PathBuf> = Vec::new();
        
        if !self.path.exists() {
            return Ok(files);
        }
        
        for entry in std::fs::read_dir(&self.path)? {
            let entry = entry?;
            let path = entry.path();
            
            if path.extension().map(|e| e == "wal").unwrap_or(false) {
                // Filter by shard if specified
                if let Some(shard_id) = self.shard_id {
                    if let Some(file_shard) = Self::parse_shard_id(&path) {
                        if file_shard != shard_id {
                            continue;
                        }
                    }
                }
                
                files.push(path);
            }
        }
        
        // Sort by sequence number (encoded in filename)
        files.sort_by(|a, b| {
            let seq_a = Self::parse_sequence(a).unwrap_or(0);
            let seq_b = Self::parse_sequence(b).unwrap_or(0);
            seq_a.cmp(&seq_b)
        });
        
        Ok(files)
    }
    
    /// Parse shard ID from filename
    fn parse_shard_id(path: &Path) -> Option<usize> {
        let stem = path.file_stem()?.to_str()?;
        let parts: Vec<&str> = stem.split('_').collect();
        if parts.len() >= 2 {
            parts[1].parse().ok()
        } else {
            None
        }
    }
    
    /// Parse sequence number from filename
    fn parse_sequence(path: &Path) -> Option<u64> {
        let stem = path.file_stem()?.to_str()?;
        let parts: Vec<&str> = stem.split('_').collect();
        if parts.len() >= 3 {
            parts[2].parse().ok()
        } else {
            None
        }
    }
    
    /// Read all entries from all WAL files (sequential)
    pub fn read_all(&self) -> io::Result<Vec<WALEntry>> {
        let mut entries = Vec::new();
        
        for file_path in self.list_files()? {
            let file_entries = self.read_file(&file_path)?;
            entries.extend(file_entries);
        }
        
        // Sort by sequence number
        entries.sort_by_key(|e| e.sequence);
        
        Ok(entries)
    }
    
    /// Read all entries from all WAL files in parallel
    /// 
    /// Uses rayon to read multiple WAL files concurrently, which can
    /// significantly speed up recovery when there are many WAL files.
    /// 
    /// Returns entries sorted by sequence number.
    pub fn read_all_parallel(&self) -> io::Result<Vec<WALEntry>> {
        let files = self.list_files()?;
        
        if files.is_empty() {
            return Ok(Vec::new());
        }
        
        let start = Instant::now();
        let file_count = files.len();
        
        // Read files in parallel
        let results: Vec<io::Result<Vec<WALEntry>>> = files
            .par_iter()
            .map(|file_path| self.read_file(file_path))
            .collect();
        
        // Merge results
        let mut all_entries = Vec::new();
        let mut total_entries = 0;
        
        for result in results {
            match result {
                Ok(entries) => {
                    total_entries += entries.len();
                    all_entries.extend(entries);
                }
                Err(e) => {
                    tracing::warn!("Error reading WAL file: {}", e);
                    // Continue with other files
                }
            }
        }
        
        // Sort by sequence number
        all_entries.par_sort_by_key(|e| e.sequence);
        
        let elapsed = start.elapsed();
        tracing::info!(
            "WAL parallel recovery: {} files, {} entries in {:?}",
            file_count, total_entries, elapsed
        );
        
        Ok(all_entries)
    }
    
    /// Read entries grouped by shard (parallel)
    /// 
    /// Returns a map from shard_id to entries for that shard.
    /// This is useful for parallel recovery where each shard can be
    /// recovered independently.
    pub fn read_by_shard_parallel(&self) -> io::Result<std::collections::HashMap<usize, Vec<WALEntry>>> {
        use std::collections::HashMap;
        
        let files = self.list_files()?;
        
        if files.is_empty() {
            return Ok(HashMap::new());
        }
        
        // Group files by shard
        let mut shard_files: HashMap<usize, Vec<PathBuf>> = HashMap::new();
        for file in files {
            let shard_id = Self::parse_shard_id(&file).unwrap_or(0);
            shard_files.entry(shard_id).or_default().push(file);
        }
        
        // Read each shard's files in parallel
        let shard_entries: Vec<(usize, Vec<WALEntry>)> = shard_files
            .into_par_iter()
            .map(|(shard_id, files)| {
                let mut entries = Vec::new();
                for file in files {
                    if let Ok(file_entries) = self.read_file(&file) {
                        entries.extend(file_entries);
                    }
                }
                entries.sort_by_key(|e| e.sequence);
                (shard_id, entries)
            })
            .collect();
        
        Ok(shard_entries.into_iter().collect())
    }
    
    /// Read entries from a single WAL file
    pub fn read_file(&self, path: &Path) -> io::Result<Vec<WALEntry>> {
        let file = File::open(path)?;
        let mut reader = BufReader::new(file);
        let mut entries = Vec::new();
        
        // Read header
        let _header = self.read_header(&mut reader)?;
        
        // Read entries until EOF or footer
        loop {
            match self.read_entry(&mut reader) {
                Ok(Some(entry)) => {
                    // Verify checksum
                    if entry.verify_checksum() {
                        entries.push(entry);
                    } else {
                        tracing::warn!("WAL entry checksum mismatch, skipping: seq={}", entry.sequence);
                    }
                }
                Ok(None) => break, // EOF or footer
                Err(e) => {
                    tracing::warn!("Error reading WAL entry: {}", e);
                    break;
                }
            }
        }
        
        Ok(entries)
    }
    
    /// Read header from reader
    fn read_header(&self, reader: &mut BufReader<File>) -> io::Result<WALHeader> {
        // Read length
        let mut len_buf = [0u8; 4];
        reader.read_exact(&mut len_buf)?;
        let len = u32::from_le_bytes(len_buf) as usize;
        
        // Read header data
        let mut data = vec![0u8; len];
        reader.read_exact(&mut data)?;
        
        let header: WALHeader = serde_json::from_slice(&data)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
        
        header.validate()
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
        
        Ok(header)
    }
    
    /// Read a single entry
    fn read_entry(&self, reader: &mut BufReader<File>) -> io::Result<Option<WALEntry>> {
        // Read length
        let mut len_buf = [0u8; 4];
        match reader.read_exact(&mut len_buf) {
            Ok(()) => {}
            Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
            Err(e) => return Err(e),
        }
        let len = u32::from_le_bytes(len_buf) as usize;
        
        // Check for footer (small size)
        if len < 50 {
            // Might be footer, try to parse
            let mut data = vec![0u8; len];
            reader.read_exact(&mut data)?;
            
            // Try to parse as footer
            if serde_json::from_slice::<WALFooter>(&data).is_ok() {
                return Ok(None); // Footer signals end
            }
            
            // Otherwise try as entry
            let entry = WALEntry::from_bytes(&data)
                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
            return Ok(Some(entry));
        }
        
        // Read entry data
        let mut data = vec![0u8; len];
        reader.read_exact(&mut data)?;
        
        let entry = WALEntry::from_bytes(&data)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
        
        Ok(Some(entry))
    }
    
    /// Create an iterator over all entries
    pub fn iter(&self) -> io::Result<WALIterator> {
        let files = self.list_files()?;
        Ok(WALIterator::new(files))
    }
    
    /// Get entries starting from a sequence number
    pub fn read_from_sequence(&self, start_seq: u64) -> io::Result<Vec<WALEntry>> {
        let entries = self.read_all()?;
        Ok(entries.into_iter().filter(|e| e.sequence >= start_seq).collect())
    }
    
    /// Get the highest sequence number in WAL files
    pub fn get_max_sequence(&self) -> io::Result<u64> {
        let files = self.list_files()?;
        let mut max_seq: u64 = 0;
        
        for file in files {
            if let Some(seq) = Self::parse_sequence(&file) {
                max_seq = max_seq.max(seq);
            }
        }
        
        Ok(max_seq)
    }
}

/// Iterator over WAL entries from multiple files
pub struct WALIterator {
    /// WAL files to iterate over
    files: Vec<PathBuf>,
    
    /// Current file index
    current_file: usize,
    
    /// Current file reader
    reader: Option<BufReader<File>>,
    
    /// Buffer of entries from current file
    buffer: Vec<WALEntry>,
    
    /// Current position in buffer
    buffer_pos: usize,
}

impl WALIterator {
    /// Create a new WAL iterator
    pub fn new(files: Vec<PathBuf>) -> Self {
        Self {
            files,
            current_file: 0,
            reader: None,
            buffer: Vec::new(),
            buffer_pos: 0,
        }
    }
    
    /// Load entries from next file
    fn load_next_file(&mut self) -> io::Result<bool> {
        if self.current_file >= self.files.len() {
            return Ok(false);
        }
        
        let path = &self.files[self.current_file];
        let reader = WALReader::new(path.parent().unwrap_or(Path::new(".")));
        
        self.buffer = reader.read_file(path)?;
        self.buffer_pos = 0;
        self.current_file += 1;
        
        Ok(!self.buffer.is_empty())
    }
}

impl Iterator for WALIterator {
    type Item = io::Result<WALEntry>;
    
    fn next(&mut self) -> Option<Self::Item> {
        loop {
            // Check if we have entries in buffer
            if self.buffer_pos < self.buffer.len() {
                let entry = self.buffer[self.buffer_pos].clone();
                self.buffer_pos += 1;
                return Some(Ok(entry));
            }
            
            // Load next file
            match self.load_next_file() {
                Ok(true) => continue,
                Ok(false) => return None,
                Err(e) => return Some(Err(e)),
            }
        }
    }
}

/// Recovery result from WAL replay
#[derive(Debug, Default)]
pub struct RecoveryResult {
    /// Number of entries replayed
    pub entries_replayed: usize,
    
    /// Number of entries skipped (checksum errors)
    pub entries_skipped: usize,
    
    /// Last sequence number seen
    pub last_sequence: u64,
    
    /// Number of files processed
    pub files_processed: usize,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::storage::wal::writer::{WALWriter, WALWriterConfig};
    use crate::storage::concurrent::metrics::WriteMetrics;
    use crate::types::MemoryNode;
    use tempfile::tempdir;
    use std::sync::Arc;
    
    #[test]
    fn test_wal_reader_basic() {
        let dir = tempdir().unwrap();
        let metrics = Arc::new(WriteMetrics::new());
        
        // Write some entries
        {
            let config = WALWriterConfig::new(dir.path());
            let mut writer = WALWriter::new(config, 0, metrics).unwrap();
            
            for i in 0..10 {
                let node = MemoryNode::new(&format!("content {}", i));
                let entry = WALEntry::insert(writer.next_sequence(), "default".to_string(), node);
                writer.append(&entry).unwrap();
            }
            
            writer.close().unwrap();
        }
        
        // Read entries
        let reader = WALReader::new(dir.path());
        let entries = reader.read_all().unwrap();
        
        assert_eq!(entries.len(), 10);
        
        // Verify order
        for (i, entry) in entries.iter().enumerate() {
            assert!(entry.verify_checksum());
        }
    }
    
    #[test]
    fn test_wal_reader_shard_filter() {
        let dir = tempdir().unwrap();
        let metrics = Arc::new(WriteMetrics::new());
        
        // Write entries to multiple shards
        for shard_id in 0..3 {
            let config = WALWriterConfig::new(dir.path());
            let mut writer = WALWriter::new(config, shard_id, metrics.clone()).unwrap();
            
            for i in 0..5 {
                let node = MemoryNode::new(&format!("shard {} content {}", shard_id, i));
                let entry = WALEntry::insert(writer.next_sequence(), "default".to_string(), node);
                writer.append(&entry).unwrap();
            }
            
            writer.close().unwrap();
        }
        
        // Read only shard 1
        let reader = WALReader::new(dir.path()).for_shard(1);
        let entries = reader.read_all().unwrap();
        
        // Should only have entries from shard 1
        assert_eq!(entries.len(), 5);
    }
    
    #[test]
    fn test_wal_reader_parallel() {
        let dir = tempdir().unwrap();
        let metrics = Arc::new(WriteMetrics::new());
        
        // Write entries to multiple shards (simulating multiple WAL files)
        for shard_id in 0..4 {
            let config = WALWriterConfig::new(dir.path());
            let mut writer = WALWriter::new(config, shard_id, metrics.clone()).unwrap();
            
            for i in 0..10 {
                let node = MemoryNode::new(&format!("shard {} content {}", shard_id, i));
                let entry = WALEntry::insert(writer.next_sequence(), "default".to_string(), node);
                writer.append(&entry).unwrap();
            }
            
            writer.close().unwrap();
        }
        
        // Read all entries in parallel
        let reader = WALReader::new(dir.path());
        let entries_parallel = reader.read_all_parallel().unwrap();
        
        // Read all entries sequentially for comparison
        let entries_sequential = reader.read_all().unwrap();
        
        // Both should have same number of entries
        assert_eq!(entries_parallel.len(), entries_sequential.len());
        assert_eq!(entries_parallel.len(), 40); // 4 shards * 10 entries
        
        // Both should be sorted by sequence
        for i in 1..entries_parallel.len() {
            assert!(entries_parallel[i].sequence >= entries_parallel[i-1].sequence);
        }
    }
    
    #[test]
    fn test_wal_reader_by_shard_parallel() {
        let dir = tempdir().unwrap();
        let metrics = Arc::new(WriteMetrics::new());
        
        // Write entries to multiple shards
        for shard_id in 0..3 {
            let config = WALWriterConfig::new(dir.path());
            let mut writer = WALWriter::new(config, shard_id, metrics.clone()).unwrap();
            
            for i in 0..5 {
                let node = MemoryNode::new(&format!("shard {} content {}", shard_id, i));
                let entry = WALEntry::insert(writer.next_sequence(), "default".to_string(), node);
                writer.append(&entry).unwrap();
            }
            
            writer.close().unwrap();
        }
        
        // Read by shard in parallel
        let reader = WALReader::new(dir.path());
        let shard_entries = reader.read_by_shard_parallel().unwrap();
        
        // Should have 3 shards
        assert_eq!(shard_entries.len(), 3);
        
        // Each shard should have 5 entries
        for (shard_id, entries) in &shard_entries {
            assert_eq!(entries.len(), 5, "Shard {} should have 5 entries", shard_id);
            
            // Entries should be sorted
            for i in 1..entries.len() {
                assert!(entries[i].sequence >= entries[i-1].sequence);
            }
        }
    }
}
