shardex 0.1.0

//! Transaction recording and batching for WAL operations
//!
//! This module provides transaction-based recording of WAL operations with efficient
//! batching and timer-based flushing for improved write throughput while maintaining
//! ACID properties.

use crate::error::ShardexError;
use crate::identifiers::{DocumentId, TransactionId};
use bytemuck::{Pod, Zeroable};
use serde::{Deserialize, Serialize};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::sync::mpsc;
use tokio::time::{interval, Interval};

/// Operations that can be recorded in the WAL
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum WalOperation {
    /// Add a posting to the index
    AddPosting {
        document_id: DocumentId,
        start: u32,
        length: u32,
        vector: Vec<f32>,
    },
    /// Remove all postings for a document
    RemoveDocument { document_id: DocumentId },
    /// Store document text (part of atomic replace operation)
    StoreDocumentText { document_id: DocumentId, text: String },
    /// Delete document text (cleanup operation)
    DeleteDocumentText { document_id: DocumentId },
}

/// A transaction containing batched WAL operations
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct WalTransaction {
    /// Unique transaction identifier
    pub id: TransactionId,
    /// Timestamp when transaction was created
    pub timestamp: SystemTime,
    /// List of operations in this transaction
    pub operations: Vec<WalOperation>,
    /// CRC32 checksum of the transaction data
    pub checksum: u32,
}

/// Binary header for WAL transactions in memory-mapped storage
#[derive(Debug, Clone, Copy, PartialEq)]
#[repr(C)]
pub struct WalTransactionHeader {
    /// Transaction identifier
    pub id: TransactionId,
    /// Unix timestamp in microseconds
    pub timestamp_micros: u64,
    /// Number of operations in the transaction
    pub operation_count: u32,
    /// Total size of serialized operations data in bytes
    pub operations_data_size: u32,
    /// CRC32 checksum of the operations data
    pub checksum: u32,
    /// Reserved bytes for future use
    pub reserved: [u8; 4],
}

// SAFETY: WalTransactionHeader contains only Pod types and has repr(C) layout
unsafe impl Pod for WalTransactionHeader {}
// SAFETY: WalTransactionHeader can be safely zero-initialized
unsafe impl Zeroable for WalTransactionHeader {}

impl WalOperation {
    /// Get the document ID associated with this operation
    pub fn document_id(&self) -> DocumentId {
        match self {
            WalOperation::AddPosting { document_id, .. } => *document_id,
            WalOperation::RemoveDocument { document_id } => *document_id,
            WalOperation::StoreDocumentText { document_id, .. } => *document_id,
            WalOperation::DeleteDocumentText { document_id } => *document_id,
        }
    }

    /// Check if this is an AddPosting operation
    pub fn is_add_posting(&self) -> bool {
        matches!(self, WalOperation::AddPosting { .. })
    }

    /// Check if this is a RemoveDocument operation
    pub fn is_remove_document(&self) -> bool {
        matches!(self, WalOperation::RemoveDocument { .. })
    }

    /// Check if this is a StoreDocumentText operation
    pub fn is_store_document_text(&self) -> bool {
        matches!(self, WalOperation::StoreDocumentText { .. })
    }

    /// Check if this is a DeleteDocumentText operation
    pub fn is_delete_document_text(&self) -> bool {
        matches!(self, WalOperation::DeleteDocumentText { .. })
    }

    /// Estimate the serialized size of this operation in bytes
    pub fn estimated_serialized_size(&self) -> usize {
        match self {
            WalOperation::AddPosting { vector, .. } => {
                // operation tag (1) + document_id (16) + start (4) + length (4) + vector length (4) + vector data
                1 + 16 + 4 + 4 + 4 + (vector.len() * 4)
            }
            WalOperation::RemoveDocument { .. } => {
                // operation tag (1) + document_id (16)
                1 + 16
            }
            WalOperation::StoreDocumentText { text, .. } => {
                // operation tag (1) + document_id (16) + text length (4) + text data (UTF-8 bytes)
                1 + 16 + 4 + text.len()
            }
            WalOperation::DeleteDocumentText { .. } => {
                // operation tag (1) + document_id (16)
                1 + 16
            }
        }
    }

    /// Validate the operation data
    pub fn validate(
        &self,
        expected_vector_dimension: Option<usize>,
        max_document_text_size: usize,
    ) -> Result<(), ShardexError> {
        match self {
            WalOperation::AddPosting {
                vector, start, length, ..
            } => {
                if *start > u32::MAX - *length {
                    return Err(ShardexError::Wal(
                        "AddPosting start + length would overflow u32".to_string(),
                    ));
                }

                if *length == 0 {
                    return Err(ShardexError::Wal("AddPosting length cannot be zero".to_string()));
                }

                if vector.is_empty() {
                    return Err(ShardexError::Wal("AddPosting vector cannot be empty".to_string()));
                }

                if let Some(expected_dim) = expected_vector_dimension {
                    if vector.len() != expected_dim {
                        return Err(ShardexError::InvalidDimension {
                            expected: expected_dim,
                            actual: vector.len(),
                        });
                    }
                }

                // Check for invalid float values
                for (i, &value) in vector.iter().enumerate() {
                    if !value.is_finite() {
                        return Err(ShardexError::Wal(format!(
                            "Invalid vector value at index {}: {} (must be finite)",
                            i, value
                        )));
                    }
                }
            }
            WalOperation::RemoveDocument { .. } => {
                // RemoveDocument operations are always valid if document_id is valid
                // DocumentId validation is handled by the type system
            }
            WalOperation::StoreDocumentText { text, .. } => {
                // String type in Rust guarantees valid UTF-8, so we don't need to check
                // But we should validate the text is not empty
                if text.is_empty() {
                    return Err(ShardexError::Wal("StoreDocumentText text cannot be empty".to_string()));
                }

                // Use the configured maximum document text size
                if text.len() > max_document_text_size {
                    return Err(ShardexError::DocumentTooLarge {
                        size: text.len(),
                        max_size: max_document_text_size,
                    });
                }
            }
            WalOperation::DeleteDocumentText { .. } => {
                // DeleteDocumentText operations are always valid if document_id is valid
                // DocumentId validation is handled by the type system
            }
        }
        Ok(())
    }
}

impl WalTransaction {
    /// Create a new transaction with the given operations
    pub fn new(operations: Vec<WalOperation>) -> Result<Self, ShardexError> {
        if operations.is_empty() {
            return Err(ShardexError::Wal("Transaction cannot have zero operations".to_string()));
        }

        let id = TransactionId::new();
        let timestamp = SystemTime::now();

        // Serialize operations to calculate checksum
        let operations_data = Self::serialize_operations(&operations)?;
        let checksum = crc32fast::hash(&operations_data);

        Ok(Self {
            id,
            timestamp,
            operations,
            checksum,
        })
    }

    /// Create a transaction with a specific ID and timestamp (for testing/recovery)
    pub fn with_id_and_timestamp(
        id: TransactionId,
        timestamp: SystemTime,
        operations: Vec<WalOperation>,
    ) -> Result<Self, ShardexError> {
        if operations.is_empty() {
            return Err(ShardexError::Wal("Transaction cannot have zero operations".to_string()));
        }

        // Serialize operations to calculate checksum
        let operations_data = Self::serialize_operations(&operations)?;
        let checksum = crc32fast::hash(&operations_data);

        Ok(Self {
            id,
            timestamp,
            operations,
            checksum,
        })
    }

    /// Get the number of operations in this transaction
    pub fn operation_count(&self) -> usize {
        self.operations.len()
    }

    /// Get all document IDs affected by this transaction
    pub fn affected_document_ids(&self) -> Vec<DocumentId> {
        let mut doc_ids: Vec<DocumentId> = self.operations.iter().map(|op| op.document_id()).collect();
        doc_ids.sort();
        doc_ids.dedup();
        doc_ids
    }

    /// Estimate the total serialized size of this transaction
    pub fn estimated_serialized_size(&self) -> usize {
        std::mem::size_of::<WalTransactionHeader>()
            + self
                .operations
                .iter()
                .map(|op| op.estimated_serialized_size())
                .sum::<usize>()
    }

    /// Validate all operations in the transaction
    pub fn validate(
        &self,
        expected_vector_dimension: Option<usize>,
        max_document_text_size: usize,
    ) -> Result<(), ShardexError> {
        if self.operations.is_empty() {
            return Err(ShardexError::Wal(
                "Transaction must contain at least one operation".to_string(),
            ));
        }

        // Validate each operation
        for (i, operation) in self.operations.iter().enumerate() {
            operation
                .validate(expected_vector_dimension, max_document_text_size)
                .map_err(|e| {
                    ShardexError::Wal(format!("Operation {} in transaction {} is invalid: {}", i, self.id, e))
                })?;
        }

        // Validate timestamp is not in the future (with some tolerance)
        let now = SystemTime::now();
        const FUTURE_TOLERANCE_SECONDS: u64 = 60; // 1 minute tolerance for clock skew

        if let Ok(duration_since_epoch) = self.timestamp.duration_since(UNIX_EPOCH) {
            if let Ok(now_duration) = now.duration_since(UNIX_EPOCH) {
                if duration_since_epoch.as_secs() > now_duration.as_secs() + FUTURE_TOLERANCE_SECONDS {
                    return Err(ShardexError::Wal(format!(
                        "Transaction timestamp is too far in the future: transaction time {:?}, current time {:?}",
                        self.timestamp, now
                    )));
                }
            }
        }

        Ok(())
    }

    /// Verify the transaction's checksum
    pub fn verify_checksum(&self) -> Result<(), ShardexError> {
        let operations_data = Self::serialize_operations(&self.operations)?;
        let calculated_checksum = crc32fast::hash(&operations_data);

        if calculated_checksum != self.checksum {
            return Err(ShardexError::Wal(format!(
                "Transaction {} checksum mismatch: expected {}, calculated {}",
                self.id, self.checksum, calculated_checksum
            )));
        }

        Ok(())
    }

    /// Serialize operations to bytes for checksum calculation and storage
    fn serialize_operations(operations: &[WalOperation]) -> Result<Vec<u8>, ShardexError> {
        bincode::serialize(operations).map_err(|e| ShardexError::Wal(format!("Failed to serialize operations: {}", e)))
    }

    /// Create a binary header for this transaction
    pub fn to_header(&self) -> Result<WalTransactionHeader, ShardexError> {
        let operations_data = Self::serialize_operations(&self.operations)?;

        let timestamp_micros = self
            .timestamp
            .duration_since(UNIX_EPOCH)
            .map_err(|e| ShardexError::Wal(format!("Invalid timestamp: {}", e)))?
            .as_micros() as u64;

        Ok(WalTransactionHeader {
            id: self.id,
            timestamp_micros,
            operation_count: self.operations.len() as u32,
            operations_data_size: operations_data.len() as u32,
            checksum: self.checksum,
            reserved: [0; 4],
        })
    }

    /// Serialize the entire transaction to binary format
    pub fn serialize(&self) -> Result<Vec<u8>, ShardexError> {
        let header = self.to_header()?;
        let operations_data = Self::serialize_operations(&self.operations)?;

        let mut result = Vec::with_capacity(std::mem::size_of::<WalTransactionHeader>() + operations_data.len());

        // Write header
        result.extend_from_slice(bytemuck::bytes_of(&header));

        // Write operations data
        result.extend_from_slice(&operations_data);

        Ok(result)
    }

    /// Deserialize a transaction from binary format
    pub fn deserialize(data: &[u8]) -> Result<Self, ShardexError> {
        if data.len() < std::mem::size_of::<WalTransactionHeader>() {
            return Err(ShardexError::Wal("Transaction data too short for header".to_string()));
        }

        // Read header
        let header_bytes = &data[0..std::mem::size_of::<WalTransactionHeader>()];
        let header: WalTransactionHeader = bytemuck::pod_read_unaligned(header_bytes);

        // Validate header consistency
        let expected_total_size = std::mem::size_of::<WalTransactionHeader>() + header.operations_data_size as usize;
        if data.len() != expected_total_size {
            return Err(ShardexError::Wal(format!(
                "Transaction data size mismatch: expected {}, got {}",
                expected_total_size,
                data.len()
            )));
        }

        // Read operations data
        let operations_data_start = std::mem::size_of::<WalTransactionHeader>();
        let operations_data = &data[operations_data_start..];

        // Verify checksum before deserializing
        let calculated_checksum = crc32fast::hash(operations_data);
        if calculated_checksum != header.checksum {
            return Err(ShardexError::Wal(format!(
                "Transaction checksum mismatch: expected {}, calculated {}",
                header.checksum, calculated_checksum
            )));
        }

        // Deserialize operations
        let operations: Vec<WalOperation> = bincode::deserialize(operations_data)
            .map_err(|e| ShardexError::Wal(format!("Failed to deserialize operations: {}", e)))?;

        // Validate operation count matches header
        if operations.len() != header.operation_count as usize {
            return Err(ShardexError::Wal(format!(
                "Operation count mismatch: header says {}, found {}",
                header.operation_count,
                operations.len()
            )));
        }

        // Convert timestamp
        let timestamp = UNIX_EPOCH + std::time::Duration::from_micros(header.timestamp_micros);

        Ok(Self {
            id: header.id,
            timestamp,
            operations,
            checksum: header.checksum,
        })
    }
}

impl WalTransactionHeader {
    /// Create a zero-initialized transaction header
    pub fn new_zero() -> Self {
        Self::zeroed()
    }

    /// Check if this header represents a valid transaction
    pub fn is_valid(&self) -> bool {
        self.operation_count > 0 && self.operations_data_size > 0
    }

    /// Get the total size this transaction will occupy including header and data
    pub fn total_size(&self) -> usize {
        std::mem::size_of::<WalTransactionHeader>() + self.operations_data_size as usize
    }
}

/// Configuration for batch processing
#[derive(Debug, Clone)]
pub struct BatchConfig {
    /// Maximum interval between flushes in milliseconds
    pub batch_write_interval_ms: u64,
    /// Maximum number of operations to batch before forcing a flush
    pub max_operations_per_batch: usize,
    /// Maximum size in bytes for a batch before forcing a flush
    pub max_batch_size_bytes: usize,
    /// Maximum size for document text in bytes
    pub max_document_text_size: usize,
}

impl Default for BatchConfig {
    fn default() -> Self {
        Self {
            batch_write_interval_ms: 100,
            max_operations_per_batch: 1000,
            max_batch_size_bytes: 1024 * 1024,        // 1MB
            max_document_text_size: 10 * 1024 * 1024, // 10MB - matches ShardexConfig default
        }
    }
}

/// Batch command for communication with the batch manager
#[derive(Debug)]
pub enum BatchCommand {
    /// Add an operation to the current batch
    AddOperation(WalOperation),
    /// Force flush the current batch
    Flush,
    /// Shutdown the batch manager
    Shutdown,
}

/// Response from batch operations
#[derive(Debug)]
pub enum BatchResponse {
    /// Operation was added successfully
    OperationAdded,
    /// Batch was flushed with the given transaction ID
    BatchFlushed(TransactionId),
    /// Error occurred during processing
    Error(ShardexError),
    /// Shutdown acknowledged
    Shutdown,
}

/// Batch manager for WAL operations
///
/// The batch manager collects WAL operations and flushes them as transactions
/// at regular intervals or when batch limits are reached.
pub struct WalBatchManager {
    /// Current batch of operations being accumulated
    current_batch: Vec<WalOperation>,
    /// Configuration for batch processing
    config: BatchConfig,
    /// Estimated size of current batch in bytes
    current_batch_size: usize,
    /// Timer for periodic flushes
    flush_timer: Interval,
    /// Expected vector dimension for validation
    expected_vector_dimension: Option<usize>,
}

impl WalBatchManager {
    /// Create a new batch manager with the given configuration
    pub fn new(config: BatchConfig, expected_vector_dimension: Option<usize>) -> Self {
        let flush_timer = interval(Duration::from_millis(config.batch_write_interval_ms));

        Self {
            current_batch: Vec::new(),
            config,
            current_batch_size: 0,
            flush_timer,
            expected_vector_dimension,
        }
    }

    /// Add an operation to the current batch
    /// Returns true if a flush is required due to batch size limits
    pub fn add_operation(&mut self, operation: WalOperation) -> Result<bool, ShardexError> {
        // Validate operation before adding
        operation.validate(self.expected_vector_dimension, self.config.max_document_text_size)?;

        let operation_size = operation.estimated_serialized_size();

        // Add the operation first
        self.current_batch.push(operation);
        self.current_batch_size += operation_size;

        // Check if we've reached the limits after adding
        let should_flush_count = self.current_batch.len() >= self.config.max_operations_per_batch;
        let should_flush_size = self.current_batch_size > self.config.max_batch_size_bytes;

        Ok(should_flush_count || should_flush_size)
    }

    /// Flush the current batch as a transaction
    /// Returns the transaction ID if a flush occurred, None if batch was empty
    pub async fn flush_batch<F>(&mut self, write_fn: F) -> Result<Option<TransactionId>, ShardexError>
    where
        F: Fn(&WalTransaction) -> Result<(), ShardexError>,
    {
        if self.current_batch.is_empty() {
            return Ok(None);
        }

        // Create transaction from current batch
        let operations = std::mem::take(&mut self.current_batch);
        let transaction = WalTransaction::new(operations)?;
        let transaction_id = transaction.id;

        // Validate transaction
        transaction.validate(self.expected_vector_dimension, self.config.max_document_text_size)?;

        // Write transaction using the provided function
        write_fn(&transaction)?;

        // Reset batch state
        self.current_batch_size = 0;

        Ok(Some(transaction_id))
    }

    /// Get the current batch statistics
    pub fn batch_stats(&self) -> BatchStats {
        BatchStats {
            operation_count: self.current_batch.len(),
            estimated_size_bytes: self.current_batch_size,
            is_empty: self.current_batch.is_empty(),
        }
    }

    /// Check if the batch should be flushed due to time interval
    pub async fn should_flush_due_to_timer(&mut self) -> bool {
        // This is a non-blocking check - it returns true if the timer has fired
        self.flush_timer.tick().await;
        !self.current_batch.is_empty()
    }

    /// Run the batch manager event loop
    pub async fn run_event_loop<F>(
        mut self,
        mut receiver: mpsc::Receiver<BatchCommand>,
        response_sender: mpsc::Sender<BatchResponse>,
        write_fn: F,
    ) where
        F: Fn(&WalTransaction) -> Result<(), ShardexError> + Send + 'static,
    {
        loop {
            tokio::select! {
                // Handle incoming commands
                command = receiver.recv() => {
                    match command {
                        Some(BatchCommand::AddOperation(operation)) => {
                            match self.add_operation(operation) {
                                Ok(should_flush) => {
                                    let _ = response_sender.send(BatchResponse::OperationAdded).await;

                                    if should_flush {
                                        match self.flush_batch(&write_fn).await {
                                            Ok(Some(transaction_id)) => {
                                                let _ = response_sender.send(BatchResponse::BatchFlushed(transaction_id)).await;
                                            }
                                            Ok(None) => {
                                                // Empty batch, nothing to flush
                                            }
                                            Err(e) => {
                                                let _ = response_sender.send(BatchResponse::Error(e)).await;
                                            }
                                        }
                                    }
                                }
                                Err(e) => {
                                    let _ = response_sender.send(BatchResponse::Error(e)).await;
                                }
                            }
                        }
                        Some(BatchCommand::Flush) => {
                            match self.flush_batch(&write_fn).await {
                                Ok(Some(transaction_id)) => {
                                    let _ = response_sender.send(BatchResponse::BatchFlushed(transaction_id)).await;
                                }
                                Ok(None) => {
                                    let _ = response_sender.send(BatchResponse::BatchFlushed(TransactionId::new())).await; // Empty flush
                                }
                                Err(e) => {
                                    let _ = response_sender.send(BatchResponse::Error(e)).await;
                                }
                            }
                        }
                        Some(BatchCommand::Shutdown) => {
                            // Flush any remaining operations
                            let _ = self.flush_batch(&write_fn).await;
                            let _ = response_sender.send(BatchResponse::Shutdown).await;
                            break;
                        }
                        None => {
                            // Channel closed, shutdown
                            let _ = self.flush_batch(&write_fn).await;
                            break;
                        }
                    }
                }

                // Handle timer-based flushes
                _ = self.should_flush_due_to_timer() => {
                    match self.flush_batch(&write_fn).await {
                        Ok(Some(transaction_id)) => {
                            let _ = response_sender.send(BatchResponse::BatchFlushed(transaction_id)).await;
                        }
                        Ok(None) => {
                            // Empty batch, nothing to flush
                        }
                        Err(e) => {
                            let _ = response_sender.send(BatchResponse::Error(e)).await;
                        }
                    }
                }
            }
        }
    }
}

/// Statistics about the current batch
#[derive(Debug, Clone, PartialEq)]
pub struct BatchStats {
    /// Number of operations in the current batch
    pub operation_count: usize,
    /// Estimated size of the batch in bytes
    pub estimated_size_bytes: usize,
    /// Whether the batch is empty
    pub is_empty: bool,
}

/// Handle for communicating with the batch manager
pub struct WalBatchHandle {
    /// Sender for batch commands
    command_sender: mpsc::Sender<BatchCommand>,
    /// Receiver for batch responses
    response_receiver: mpsc::Receiver<BatchResponse>,
}

impl WalBatchHandle {
    /// Create a new batch handle and manager
    pub fn new(config: BatchConfig, expected_vector_dimension: Option<usize>) -> (Self, WalBatchManager) {
        let (command_sender, _command_receiver) = mpsc::channel(1000);
        let (_response_sender, response_receiver) = mpsc::channel(1000);

        let manager = WalBatchManager::new(config, expected_vector_dimension);

        let handle = Self {
            command_sender,
            response_receiver,
        };

        (handle, manager)
    }

    /// Add an operation to the batch
    pub async fn add_operation(&mut self, operation: WalOperation) -> Result<(), ShardexError> {
        self.command_sender
            .send(BatchCommand::AddOperation(operation))
            .await
            .map_err(|_| ShardexError::Wal("Batch manager channel closed".to_string()))?;

        // Wait for response
        match self.response_receiver.recv().await {
            Some(BatchResponse::OperationAdded) => Ok(()),
            Some(BatchResponse::Error(e)) => Err(e),
            Some(response) => Err(ShardexError::Wal(format!("Unexpected response: {:?}", response))),
            None => Err(ShardexError::Wal("Batch manager response channel closed".to_string())),
        }
    }

    /// Flush the current batch
    pub async fn flush(&mut self) -> Result<TransactionId, ShardexError> {
        self.command_sender
            .send(BatchCommand::Flush)
            .await
            .map_err(|_| ShardexError::Wal("Batch manager channel closed".to_string()))?;

        // Wait for response
        match self.response_receiver.recv().await {
            Some(BatchResponse::BatchFlushed(transaction_id)) => Ok(transaction_id),
            Some(BatchResponse::Error(e)) => Err(e),
            Some(response) => Err(ShardexError::Wal(format!("Unexpected response: {:?}", response))),
            None => Err(ShardexError::Wal("Batch manager response channel closed".to_string())),
        }
    }

    /// Shutdown the batch manager
    pub async fn shutdown(&mut self) -> Result<(), ShardexError> {
        self.command_sender
            .send(BatchCommand::Shutdown)
            .await
            .map_err(|_| ShardexError::Wal("Batch manager channel closed".to_string()))?;

        // Wait for response
        match self.response_receiver.recv().await {
            Some(BatchResponse::Shutdown) => Ok(()),
            Some(BatchResponse::Error(e)) => Err(e),
            Some(response) => Err(ShardexError::Wal(format!("Unexpected response: {:?}", response))),
            None => Err(ShardexError::Wal("Batch manager response channel closed".to_string())),
        }
    }
}

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

    #[test]
    fn test_wal_operation_document_id() {
        let doc_id = DocumentId::new();
        let vector = vec![1.0, 2.0, 3.0];

        let add_op = WalOperation::AddPosting {
            document_id: doc_id,
            start: 0,
            length: 100,
            vector,
        };

        let remove_op = WalOperation::RemoveDocument { document_id: doc_id };

        assert_eq!(add_op.document_id(), doc_id);
        assert_eq!(remove_op.document_id(), doc_id);
        assert!(add_op.is_add_posting());
        assert!(!add_op.is_remove_document());
        assert!(!remove_op.is_add_posting());
        assert!(remove_op.is_remove_document());
    }

    #[test]
    fn test_wal_operation_validation() {
        let doc_id = DocumentId::new();

        // Valid AddPosting
        let valid_add = WalOperation::AddPosting {
            document_id: doc_id,
            start: 100,
            length: 50,
            vector: vec![1.0, 2.0, 3.0],
        };
        assert!(valid_add.validate(Some(3), 10 * 1024 * 1024).is_ok());
        assert!(valid_add.validate(None, 10 * 1024 * 1024).is_ok());

        // Invalid dimension
        assert!(valid_add.validate(Some(4), 10 * 1024 * 1024).is_err());

        // Invalid vector with NaN
        let invalid_nan = WalOperation::AddPosting {
            document_id: doc_id,
            start: 100,
            length: 50,
            vector: vec![1.0, f32::NAN, 3.0],
        };
        assert!(invalid_nan.validate(None, 10 * 1024 * 1024).is_err());

        // Invalid vector with infinity
        let invalid_inf = WalOperation::AddPosting {
            document_id: doc_id,
            start: 100,
            length: 50,
            vector: vec![1.0, f32::INFINITY, 3.0],
        };
        assert!(invalid_inf.validate(None, 10 * 1024 * 1024).is_err());

        // Zero length
        let zero_length = WalOperation::AddPosting {
            document_id: doc_id,
            start: 100,
            length: 0,
            vector: vec![1.0, 2.0, 3.0],
        };
        assert!(zero_length.validate(None, 10 * 1024 * 1024).is_err());

        // Empty vector
        let empty_vector = WalOperation::AddPosting {
            document_id: doc_id,
            start: 100,
            length: 50,
            vector: vec![],
        };
        assert!(empty_vector.validate(None, 10 * 1024 * 1024).is_err());

        // Overflow in start + length
        let overflow = WalOperation::AddPosting {
            document_id: doc_id,
            start: u32::MAX,
            length: 1,
            vector: vec![1.0, 2.0, 3.0],
        };
        assert!(overflow.validate(None, 10 * 1024 * 1024).is_err());

        // Valid RemoveDocument (always valid)
        let remove_doc = WalOperation::RemoveDocument { document_id: doc_id };
        assert!(remove_doc.validate(None, 10 * 1024 * 1024).is_ok());
        assert!(remove_doc.validate(Some(128), 10 * 1024 * 1024).is_ok());
    }

    #[test]
    fn test_wal_transaction_creation() {
        let doc_id = DocumentId::new();
        let operations = vec![
            WalOperation::AddPosting {
                document_id: doc_id,
                start: 0,
                length: 100,
                vector: vec![1.0, 2.0, 3.0],
            },
            WalOperation::RemoveDocument { document_id: doc_id },
        ];

        let transaction = WalTransaction::new(operations.clone()).unwrap();

        assert_eq!(transaction.operations, operations);
        assert_eq!(transaction.operation_count(), 2);
        assert!(transaction.checksum != 0);
    }

    #[test]
    fn test_wal_transaction_empty_operations() {
        let result = WalTransaction::new(vec![]);
        // assert!(result.is_err());
        if let Err(ShardexError::Wal(msg)) = result {
            assert!(msg.contains("cannot have zero operations"));
        } else {
            panic!("Expected Wal error");
        }
    }

    #[test]
    fn test_wal_transaction_affected_documents() {
        let doc_id1 = DocumentId::new();
        let doc_id2 = DocumentId::new();

        let operations = vec![
            WalOperation::AddPosting {
                document_id: doc_id1,
                start: 0,
                length: 100,
                vector: vec![1.0, 2.0, 3.0],
            },
            WalOperation::AddPosting {
                document_id: doc_id2,
                start: 50,
                length: 75,
                vector: vec![4.0, 5.0, 6.0],
            },
            WalOperation::RemoveDocument { document_id: doc_id1 },
        ];

        let transaction = WalTransaction::new(operations).unwrap();
        let affected = transaction.affected_document_ids();

        // Should have both document IDs, sorted and deduplicated
        assert_eq!(affected.len(), 2);
        assert!(affected.contains(&doc_id1));
        assert!(affected.contains(&doc_id2));
        // Check sorting
        assert!(affected[0] < affected[1]);
    }

    #[test]
    fn test_wal_transaction_checksum_verification() {
        let doc_id = DocumentId::new();
        let operations = vec![WalOperation::AddPosting {
            document_id: doc_id,
            start: 0,
            length: 100,
            vector: vec![1.0, 2.0, 3.0],
        }];

        let transaction = WalTransaction::new(operations).unwrap();

        // Should verify successfully
        assert!(transaction.verify_checksum().is_ok());

        // Manually create transaction with wrong checksum
        let mut bad_transaction = transaction.clone();
        bad_transaction.checksum = 12345; // Wrong checksum

        assert!(bad_transaction.verify_checksum().is_err());
    }

    #[test]
    fn test_wal_transaction_serialization() {
        let doc_id = DocumentId::new();
        let operations = vec![
            WalOperation::AddPosting {
                document_id: doc_id,
                start: 0,
                length: 100,
                vector: vec![1.0, 2.0, 3.0],
            },
            WalOperation::RemoveDocument { document_id: doc_id },
        ];

        let transaction = WalTransaction::new(operations).unwrap();

        // Test serialization
        let serialized = transaction.serialize().unwrap();
        assert!(!serialized.is_empty());

        // Test deserialization
        let deserialized = WalTransaction::deserialize(&serialized).unwrap();

        assert_eq!(transaction.id, deserialized.id);
        assert_eq!(transaction.operations, deserialized.operations);
        assert_eq!(transaction.checksum, deserialized.checksum);

        // Timestamps should be very close (within a microsecond or two)
        let time_diff = if transaction.timestamp > deserialized.timestamp {
            transaction
                .timestamp
                .duration_since(deserialized.timestamp)
                .unwrap()
        } else {
            deserialized
                .timestamp
                .duration_since(transaction.timestamp)
                .unwrap()
        };
        assert!(time_diff.as_micros() < 10); // Should be identical or very close
    }

    #[test]
    fn test_wal_transaction_header_operations() {
        let doc_id = DocumentId::new();
        let operations = vec![WalOperation::AddPosting {
            document_id: doc_id,
            start: 0,
            length: 100,
            vector: vec![1.0, 2.0, 3.0],
        }];

        let transaction = WalTransaction::new(operations).unwrap();
        let header = transaction.to_header().unwrap();

        assert_eq!(header.id, transaction.id);
        assert_eq!(header.operation_count, 1);
        assert!(header.operations_data_size > 0);
        assert_eq!(header.checksum, transaction.checksum);
        assert!(header.is_valid());

        let total_size = header.total_size();
        assert_eq!(
            total_size,
            std::mem::size_of::<WalTransactionHeader>() + header.operations_data_size as usize
        );
    }

    #[test]
    fn test_wal_transaction_header_bytemuck() {
        let header = WalTransactionHeader {
            id: TransactionId::new(),
            timestamp_micros: 1640995200000000, // 2022-01-01T00:00:00Z
            operation_count: 5,
            operations_data_size: 1024,
            checksum: 0x12345678,
            reserved: [0; 4],
        };

        // Test Pod trait - should be able to cast to bytes
        let bytes: &[u8] = bytemuck::bytes_of(&header);
        assert_eq!(bytes.len(), std::mem::size_of::<WalTransactionHeader>());

        // Test round-trip
        let header_restored: WalTransactionHeader = bytemuck::pod_read_unaligned(bytes);
        assert_eq!(header, header_restored);
    }

    #[tokio::test]
    async fn test_invalid_transaction_data() {
        // Test invalid transaction ID
        let invalid_header = WalTransactionHeader {
            id: TransactionId::from_bytes([0xFF; 16]),
            timestamp_micros: 0,
            operation_count: 1,
            operations_data_size: 100,
            checksum: 0,
            reserved: [0; 4],
        };

        let bytes = bytemuck::bytes_of(&invalid_header);
        let result: Result<WalTransactionHeader, _> = bytemuck::try_pod_read_unaligned(bytes);
        assert!(result.is_ok()); // bytemuck doesn't validate values, just structure

        // Test extreme timestamp values
        let timestamp_header = WalTransactionHeader {
            id: TransactionId::new(),
            timestamp_micros: u64::MAX,
            operation_count: 1,
            operations_data_size: 100,
            checksum: 0,
            reserved: [0; 4],
        };

        let bytes = bytemuck::bytes_of(&timestamp_header);
        let header: WalTransactionHeader = bytemuck::pod_read_unaligned(bytes);
        assert_eq!(header.timestamp_micros, u64::MAX);

        // Test operation count overflow
        let overflow_header = WalTransactionHeader {
            id: TransactionId::new(),
            timestamp_micros: 1234567890,
            operation_count: u32::MAX,
            operations_data_size: 100,
            checksum: 0,
            reserved: [0; 4],
        };

        let bytes = bytemuck::bytes_of(&overflow_header);
        let header: WalTransactionHeader = bytemuck::pod_read_unaligned(bytes);
        assert_eq!(header.operation_count, u32::MAX);

        // Test size mismatch scenarios
        let size_mismatch_header = WalTransactionHeader {
            id: TransactionId::new(),
            timestamp_micros: 1234567890,
            operation_count: 1,
            operations_data_size: u32::MAX,
            checksum: 0,
            reserved: [0; 4],
        };

        let bytes = bytemuck::bytes_of(&size_mismatch_header);
        let header: WalTransactionHeader = bytemuck::pod_read_unaligned(bytes);
        assert_eq!(header.operations_data_size, u32::MAX);
    }

    #[tokio::test]
    async fn test_invalid_batch_operations() {
        let config = BatchConfig::default();
        let mut manager = WalBatchManager::new(config, None); // No dimension constraint for testing edge cases

        // Test with zero-length vector (edge case)
        let operation = WalOperation::AddPosting {
            document_id: DocumentId::new(),
            start: 0,
            length: 100,
            vector: vec![], // Empty vector
        };

        let result = manager.add_operation(operation);
        assert!(result.is_err()); // Empty vectors should be invalid

        // Test with extremely large vector
        let large_vector = vec![1.0; 1_000_000]; // 1M elements
        let operation = WalOperation::AddPosting {
            document_id: DocumentId::new(),
            start: 0,
            length: 100,
            vector: large_vector,
        };

        let result = manager.add_operation(operation);
        assert!(result.is_ok()); // Large vectors should be valid

        // Test extreme start/length combinations
        let operation = WalOperation::AddPosting {
            document_id: DocumentId::new(),
            start: u32::MAX,
            length: u32::MAX,
            vector: vec![1.0, 2.0, 3.0],
        };

        let result = manager.add_operation(operation);
        assert!(result.is_err()); // start + length would overflow u32

        // Test document removal operation
        let operation = WalOperation::RemoveDocument {
            document_id: DocumentId::new(),
        };

        let result = manager.add_operation(operation);
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_batch_manager_basic_operations() {
        let config = BatchConfig::default();
        let mut manager = WalBatchManager::new(config, Some(3));

        let doc_id = DocumentId::new();
        let operation = WalOperation::AddPosting {
            document_id: doc_id,
            start: 0,
            length: 100,
            vector: vec![1.0, 2.0, 3.0],
        };

        // Add operation should not trigger flush for first operation
        let should_flush = manager.add_operation(operation).unwrap();
        assert!(!should_flush);

        let stats = manager.batch_stats();
        assert_eq!(stats.operation_count, 1);
        assert!(!stats.is_empty);
        assert!(stats.estimated_size_bytes > 0);

        // Flush batch manually
        use std::sync::atomic::{AtomicBool, Ordering};
        use std::sync::Arc;

        let transaction_written = Arc::new(AtomicBool::new(false));
        let transaction_written_clone = transaction_written.clone();

        let write_fn = move |transaction: &WalTransaction| -> Result<(), ShardexError> {
            assert_eq!(transaction.operations.len(), 1);
            transaction_written_clone.store(true, Ordering::SeqCst);
            Ok(())
        };

        let result = manager.flush_batch(write_fn).await.unwrap();
        assert!(result.is_some());
        assert!(transaction_written.load(Ordering::SeqCst));

        // Batch should be empty after flush
        let stats = manager.batch_stats();
        assert_eq!(stats.operation_count, 0);
        assert!(stats.is_empty);
    }

    #[tokio::test]
    async fn test_batch_manager_size_limits() {
        let config = BatchConfig {
            batch_write_interval_ms: 1000, // Long interval
            max_operations_per_batch: 2,   // Small limit for testing
            max_batch_size_bytes: 1024 * 1024,
            max_document_text_size: 10 * 1024 * 1024,
        };
        let mut manager = WalBatchManager::new(config, Some(3));

        let doc_id = DocumentId::new();
        let operation = WalOperation::AddPosting {
            document_id: doc_id,
            start: 0,
            length: 100,
            vector: vec![1.0, 2.0, 3.0],
        };

        // Add first operation - should not trigger flush
        let should_flush = manager.add_operation(operation.clone()).unwrap();
        assert!(!should_flush);

        // Add second operation - should trigger flush due to count limit
        let should_flush = manager.add_operation(operation).unwrap();
        let stats = manager.batch_stats();
        assert!(
            should_flush,
            "Second operation should trigger flush. Count: {}, should_flush: {}",
            stats.operation_count, should_flush
        );
    }

    #[tokio::test]
    async fn test_batch_config_defaults() {
        let config = BatchConfig::default();
        assert_eq!(config.batch_write_interval_ms, 100);
        assert_eq!(config.max_operations_per_batch, 1000);
        assert_eq!(config.max_batch_size_bytes, 1024 * 1024);
    }

    #[tokio::test]
    async fn test_empty_batch_flush() {
        let config = BatchConfig::default();
        let mut manager = WalBatchManager::new(config, None);

        let write_fn = |_: &WalTransaction| -> Result<(), ShardexError> {
            panic!("Write function should not be called for empty batch");
        };

        let result = manager.flush_batch(write_fn).await.unwrap();
        assert!(result.is_none());
    }

    #[tokio::test]
    async fn test_batch_manager_validation() {
        let config = BatchConfig::default();
        let mut manager = WalBatchManager::new(config, Some(3));

        let doc_id = DocumentId::new();

        // Invalid operation - wrong vector dimension
        let invalid_operation = WalOperation::AddPosting {
            document_id: doc_id,
            start: 0,
            length: 100,
            vector: vec![1.0, 2.0], // Wrong dimension - expected 3
        };

        let result = manager.add_operation(invalid_operation);
        assert!(result.is_err());

        if let Err(ShardexError::InvalidDimension { expected, actual }) = result {
            assert_eq!(expected, 3);
            assert_eq!(actual, 2);
        } else {
            panic!("Expected InvalidDimension error");
        }

        // Batch should still be empty after failed operation
        let stats = manager.batch_stats();
        assert!(stats.is_empty);
    }

    #[tokio::test]
    async fn test_batch_manager_with_wal_integration() {
        use crate::layout::DirectoryLayout;
        use crate::test_utils::TestEnvironment;
        use crate::wal::WalSegment;

        let _test_env = TestEnvironment::new("test_batch_wal_integration");
        let layout = DirectoryLayout::new(_test_env.path());
        let segment_path = layout.wal_segment_path(1);
        let capacity = 8192; // 8KB for testing

        // Create a WAL segment
        let segment = std::sync::Arc::new(WalSegment::create(1, segment_path, capacity).unwrap());

        // Create batch configuration
        let config = BatchConfig {
            batch_write_interval_ms: 50,
            max_operations_per_batch: 3,
            max_batch_size_bytes: 1024,
            max_document_text_size: 10 * 1024 * 1024,
        };

        let mut manager = WalBatchManager::new(config, Some(3));

        let doc_id = DocumentId::new();
        let operations = vec![
            WalOperation::AddPosting {
                document_id: doc_id,
                start: 0,
                length: 100,
                vector: vec![1.0, 2.0, 3.0],
            },
            WalOperation::RemoveDocument { document_id: doc_id },
        ];

        // Add operations to batch
        for operation in operations {
            let should_flush = manager.add_operation(operation).unwrap();
            assert!(!should_flush); // Should not trigger flush yet
        }

        // Create a write function that uses the WAL segment
        use std::sync::atomic::{AtomicUsize, Ordering};
        use std::sync::Arc;

        let write_count = Arc::new(AtomicUsize::new(0));
        let write_count_clone = write_count.clone();
        let segment_clone = segment.clone();

        let write_fn = move |transaction: &WalTransaction| -> Result<(), ShardexError> {
            segment_clone.append_transaction(transaction)?;
            segment_clone.sync()?;
            write_count_clone.fetch_add(1, Ordering::SeqCst);
            Ok(())
        };

        // Flush the batch
        let result = manager.flush_batch(write_fn).await.unwrap();
        assert!(result.is_some());
        assert_eq!(write_count.load(Ordering::SeqCst), 1);

        // Verify the transaction was written to the segment
        assert!(segment.write_pointer() > crate::wal::initial_write_position());
    }

    #[tokio::test]
    async fn test_atomic_batch_commits() {
        let config = BatchConfig::default();
        let mut manager = WalBatchManager::new(config, Some(3));

        let doc_id1 = DocumentId::new();
        let doc_id2 = DocumentId::new();

        // Create a batch with multiple operations affecting different documents
        let operations = vec![
            WalOperation::AddPosting {
                document_id: doc_id1,
                start: 0,
                length: 100,
                vector: vec![1.0, 2.0, 3.0],
            },
            WalOperation::AddPosting {
                document_id: doc_id2,
                start: 50,
                length: 75,
                vector: vec![4.0, 5.0, 6.0],
            },
            WalOperation::RemoveDocument { document_id: doc_id1 },
        ];

        for operation in operations {
            manager.add_operation(operation).unwrap();
        }

        // Test successful atomic commit
        use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
        use std::sync::Arc;

        let commit_called = Arc::new(AtomicBool::new(false));
        let operation_count = Arc::new(AtomicUsize::new(0));
        let commit_called_clone = commit_called.clone();
        let operation_count_clone = operation_count.clone();

        let write_fn = move |transaction: &WalTransaction| -> Result<(), ShardexError> {
            // Verify all operations are committed together
            commit_called_clone.store(true, Ordering::SeqCst);
            operation_count_clone.store(transaction.operations.len(), Ordering::SeqCst);

            // Verify transaction integrity
            transaction.validate(Some(3), 10 * 1024 * 1024)?;
            transaction.verify_checksum()?;

            Ok(())
        };

        let result = manager.flush_batch(write_fn).await.unwrap();
        assert!(result.is_some());
        assert!(commit_called.load(Ordering::SeqCst));
        assert_eq!(operation_count.load(Ordering::SeqCst), 3);

        // Batch should be empty after successful commit
        let stats = manager.batch_stats();
        assert!(stats.is_empty);
    }

    #[tokio::test]
    async fn test_failed_batch_commit_rollback() {
        let config = BatchConfig::default();
        let mut manager = WalBatchManager::new(config, Some(3));

        let doc_id = DocumentId::new();
        let operation = WalOperation::AddPosting {
            document_id: doc_id,
            start: 0,
            length: 100,
            vector: vec![1.0, 2.0, 3.0],
        };

        manager.add_operation(operation).unwrap();

        // Verify batch has one operation before flush
        assert_eq!(manager.batch_stats().operation_count, 1);

        // Test failed commit (simulate write failure)
        let write_fn = |_transaction: &WalTransaction| -> Result<(), ShardexError> {
            Err(ShardexError::Wal("Simulated write failure".to_string()))
        };

        let result = manager.flush_batch(write_fn).await;
        assert!(result.is_err());

        if let Err(ShardexError::Wal(msg)) = result {
            assert!(msg.contains("Simulated write failure"));
        } else {
            panic!("Expected Wal error");
        }

        // On flush failure, the batch should still be cleared to prevent infinite retry
        // This is the current behavior - in production you might want different behavior
        let stats = manager.batch_stats();
        assert!(stats.is_empty);
    }

    #[test]
    fn test_operation_estimated_size() {
        let doc_id = DocumentId::new();

        let add_op = WalOperation::AddPosting {
            document_id: doc_id,
            start: 0,
            length: 100,
            vector: vec![1.0, 2.0, 3.0, 4.0],
        };

        let remove_op = WalOperation::RemoveDocument { document_id: doc_id };

        // AddPosting: tag(1) + doc_id(16) + start(4) + length(4) + vec_len(4) + vec_data(16) = 45
        assert_eq!(add_op.estimated_serialized_size(), 45);

        // RemoveDocument: tag(1) + doc_id(16) = 17
        assert_eq!(remove_op.estimated_serialized_size(), 17);
    }

    #[test]
    fn test_transaction_validation() {
        let doc_id = DocumentId::new();

        // Valid transaction
        let valid_ops = vec![WalOperation::AddPosting {
            document_id: doc_id,
            start: 0,
            length: 100,
            vector: vec![1.0, 2.0, 3.0],
        }];
        let valid_transaction = WalTransaction::new(valid_ops).unwrap();
        assert!(valid_transaction
            .validate(Some(3), 10 * 1024 * 1024)
            .is_ok());

        // Invalid vector dimension
        assert!(valid_transaction
            .validate(Some(4), 10 * 1024 * 1024)
            .is_err());

        // Future timestamp should be rejected (create with manual timestamp)
        let future_time = SystemTime::now() + std::time::Duration::from_secs(3600); // 1 hour in future
        let future_ops = vec![WalOperation::RemoveDocument { document_id: doc_id }];
        let future_transaction =
            WalTransaction::with_id_and_timestamp(TransactionId::new(), future_time, future_ops).unwrap();
        assert!(future_transaction.validate(None, 10 * 1024 * 1024).is_err());
    }

    #[test]
    fn test_document_text_wal_operations() {
        let doc_id = DocumentId::new();
        let text = "The quick brown fox jumps over the lazy dog.".to_string();

        // Test StoreDocumentText operation
        let store_op = WalOperation::StoreDocumentText {
            document_id: doc_id,
            text: text.clone(),
        };

        assert_eq!(store_op.document_id(), doc_id);
        assert!(store_op.is_store_document_text());
        assert!(!store_op.is_delete_document_text());
        assert!(!store_op.is_add_posting());
        assert!(!store_op.is_remove_document());

        // Test DeleteDocumentText operation
        let delete_op = WalOperation::DeleteDocumentText { document_id: doc_id };

        assert_eq!(delete_op.document_id(), doc_id);
        assert!(delete_op.is_delete_document_text());
        assert!(!delete_op.is_store_document_text());
        assert!(!delete_op.is_add_posting());
        assert!(!delete_op.is_remove_document());

        // Test size estimation includes text content
        let store_size = store_op.estimated_serialized_size();
        let delete_size = delete_op.estimated_serialized_size();

        assert!(store_size > delete_size);
        assert!(store_size > text.len()); // Should include text plus overhead
    }

    #[test]
    fn test_document_text_operation_validation() {
        let doc_id = DocumentId::new();

        // Valid text operation
        let valid_text = "Hello, world!".to_string();
        let valid_op = WalOperation::StoreDocumentText {
            document_id: doc_id,
            text: valid_text,
        };
        assert!(valid_op.validate(None, 10 * 1024 * 1024).is_ok());

        // Test with configured max size
        let large_text = "x".repeat(1000);
        let large_op = WalOperation::StoreDocumentText {
            document_id: doc_id,
            text: large_text,
        };

        // Should be valid by default (no max size set)
        assert!(large_op.validate(None, 10 * 1024 * 1024).is_ok());

        // Delete operation should always be valid
        let delete_op = WalOperation::DeleteDocumentText { document_id: doc_id };
        assert!(delete_op.validate(None, 10 * 1024 * 1024).is_ok());
        assert!(delete_op.validate(Some(128), 10 * 1024 * 1024).is_ok());
    }

    #[test]
    fn test_document_text_transaction_atomicity() {
        let doc_id = DocumentId::new();
        let text = "Updated document content".to_string();

        // Create atomic replace transaction
        let operations = vec![
            WalOperation::StoreDocumentText {
                document_id: doc_id,
                text,
            },
            WalOperation::RemoveDocument { document_id: doc_id },
            WalOperation::AddPosting {
                document_id: doc_id,
                start: 0,
                length: 7,
                vector: vec![1.0, 2.0, 3.0],
            },
        ];

        let transaction = WalTransaction::new(operations).unwrap();
        assert_eq!(transaction.operation_count(), 3);

        let affected_docs = transaction.affected_document_ids();
        assert_eq!(affected_docs.len(), 1);
        assert_eq!(affected_docs[0], doc_id);

        // Should validate successfully
        assert!(transaction.validate(Some(3), 10 * 1024 * 1024).is_ok());

        // Should serialize and deserialize correctly
        let serialized = transaction.serialize().unwrap();
        let deserialized = WalTransaction::deserialize(&serialized).unwrap();
        assert_eq!(transaction.operations, deserialized.operations);
    }
}