//! Comprehensive Bounds Checking for ML Operations
//!
//! This module provides bounds checking for all array and tensor operations
//! to prevent buffer overflows and memory safety violations.

use std::collections::HashMap;

use tracing::{debug, error, warn};

use super::{MLSafetyConfig, MLSafetyError, SafetyResult};

/// Bounds checker with comprehensive validation
#[derive(Debug, Clone)]
pub struct BoundsChecker {
    config: MLSafetyConfig,
    violation_counts: HashMap<String, usize>,
}

impl BoundsChecker {
    /// Create new bounds checker
    pub fn new(config: &MLSafetyConfig) -> Self {
        Self {
            config: config.clone(),
            violation_counts: HashMap::new(),
        }
    }

    /// Check array bounds for indexing operations
    pub fn check_array_bounds(
        &mut self,
        array_len: usize,
        index: usize,
        operation: &str,
    ) -> SafetyResult<()> {
        if !self.config.bounds_checking {
            return Ok(());
        }

        if index >= array_len {
            let violation_key = format!("array_bounds_{}", operation);
            let count = self
                .violation_counts
                .entry(violation_key.clone())
                .or_insert(0);
            *count += 1;

            error!(
                "Array bounds violation in {}: index {} >= length {} (violation #{} for this operation)",
                operation, index, array_len, count
            );

            return Err(MLSafetyError::BoundsCheck {
                index,
                length: array_len,
            });
        }

        debug!(
            "Array bounds check passed: {} index {} < length {}",
            operation, index, array_len
        );
        Ok(())
    }

    /// Check slice bounds for range operations
    pub fn check_slice_bounds(
        &mut self,
        array_len: usize,
        start: usize,
        end: usize,
        operation: &str,
    ) -> SafetyResult<()> {
        if !self.config.bounds_checking {
            return Ok(());
        }

        // Check start bounds
        if start > array_len {
            let violation_key = format!("slice_start_{}", operation);
            let count = self.violation_counts.entry(violation_key).or_insert(0);
            *count += 1;

            return Err(MLSafetyError::BoundsCheck {
                index: start,
                length: array_len,
            });
        }

        // Check end bounds
        if end > array_len {
            let violation_key = format!("slice_end_{}", operation);
            let count = self.violation_counts.entry(violation_key).or_insert(0);
            *count += 1;

            return Err(MLSafetyError::BoundsCheck {
                index: end,
                length: array_len,
            });
        }

        // Check start <= end
        if start > end {
            return Err(MLSafetyError::TensorSafety {
                reason: format!(
                    "Invalid slice range in {}: start {} > end {}",
                    operation, start, end
                ),
            });
        }

        debug!(
            "Slice bounds check passed: {} range [{}..{}] within length {}",
            operation, start, end, array_len
        );
        Ok(())
    }

    /// Check multi-dimensional tensor bounds
    pub fn check_tensor_bounds(
        &mut self,
        tensor_dims: &[usize],
        indices: &[usize],
        operation: &str,
    ) -> SafetyResult<()> {
        if !self.config.bounds_checking {
            return Ok(());
        }

        if indices.len() != tensor_dims.len() {
            return Err(MLSafetyError::TensorSafety {
                reason: format!(
                    "Dimension mismatch in {}: indices length {} != tensor dimensions {}",
                    operation,
                    indices.len(),
                    tensor_dims.len()
                ),
            });
        }

        for (dim, (&index, &dim_size)) in
            indices.into_iter().zip(tensor_dims.into_iter()).enumerate()
        {
            if index >= dim_size {
                let violation_key = format!("tensor_bounds_{}_{}", operation, dim);
                let count = self.violation_counts.entry(violation_key).or_insert(0);
                *count += 1;

                error!(
                    "Tensor bounds violation in {} dim {}: index {} >= size {} (violation #{})",
                    operation, dim, index, dim_size, count
                );

                return Err(MLSafetyError::BoundsCheck {
                    index,
                    length: dim_size,
                });
            }
        }

        debug!(
            "Tensor bounds check passed: {} indices {:?} within dims {:?}",
            operation, indices, tensor_dims
        );
        Ok(())
    }

    /// Check buffer size for data operations
    pub fn check_buffer_size(
        &mut self,
        buffer_size: usize,
        required_size: usize,
        operation: &str,
    ) -> SafetyResult<()> {
        if !self.config.bounds_checking {
            return Ok(());
        }

        if buffer_size < required_size {
            let violation_key = format!("buffer_size_{}", operation);
            let count = self.violation_counts.entry(violation_key).or_insert(0);
            *count += 1;

            error!(
                "Buffer size violation in {}: size {} < required {} (violation #{})",
                operation, buffer_size, required_size, count
            );

            return Err(MLSafetyError::MemorySafety {
                reason: format!(
                    "Buffer too small: {} bytes < {} required",
                    buffer_size, required_size
                ),
            });
        }

        debug!(
            "Buffer size check passed: {} size {} >= required {}",
            operation, buffer_size, required_size
        );
        Ok(())
    }

    /// Check matrix dimensions for multiplication
    pub fn check_matmul_dims(
        &mut self,
        lhs_dims: &[usize],
        rhs_dims: &[usize],
        operation: &str,
    ) -> SafetyResult<()> {
        if !self.config.bounds_checking {
            return Ok(());
        }

        // Check minimum dimensions
        if lhs_dims.len() < 2 {
            return Err(MLSafetyError::TensorSafety {
                reason: format!(
                    "Left matrix in {} has insufficient dimensions: {} < 2",
                    operation,
                    lhs_dims.len()
                ),
            });
        }

        if rhs_dims.len() < 2 {
            return Err(MLSafetyError::TensorSafety {
                reason: format!(
                    "Right matrix in {} has insufficient dimensions: {} < 2",
                    operation,
                    rhs_dims.len()
                ),
            });
        }

        // Check inner dimensions match
        let lhs_cols = lhs_dims[lhs_dims.len() - 1];
        let rhs_rows = rhs_dims[rhs_dims.len() - 2];

        if lhs_cols != rhs_rows {
            let violation_key = format!("matmul_dims_{}", operation);
            let count = self.violation_counts.entry(violation_key).or_insert(0);
            *count += 1;

            error!(
                "Matrix multiplication dimension mismatch in {}: {} cols != {} rows (violation #{})",
                operation, lhs_cols, rhs_rows, count
            );

            return Err(MLSafetyError::TensorSafety {
                reason: format!(
                    "Matrix multiplication incompatible: {} x {} cannot multiply with {} x {}",
                    lhs_dims[lhs_dims.len() - 2],
                    lhs_cols,
                    rhs_rows,
                    rhs_dims[rhs_dims.len() - 1]
                ),
            });
        }

        // Check batch dimensions match (if present)
        let min_batch_dims = (lhs_dims.len() - 2).min(rhs_dims.len() - 2);
        for i in 0..min_batch_dims {
            let lhs_batch = lhs_dims[i];
            let rhs_batch = rhs_dims[i];

            if lhs_batch != rhs_batch && lhs_batch != 1 && rhs_batch != 1 {
                return Err(MLSafetyError::TensorSafety {
                    reason: format!(
                        "Batch dimension mismatch in {}: dim {} has {} vs {}",
                        operation, i, lhs_batch, rhs_batch
                    ),
                });
            }
        }

        debug!(
            "Matrix multiplication dims check passed: {:?} x {:?}",
            lhs_dims, rhs_dims
        );
        Ok(())
    }

    /// Check broadcasting compatibility
    pub fn check_broadcast_dims(
        &mut self,
        dims1: &[usize],
        dims2: &[usize],
        operation: &str,
    ) -> SafetyResult<()> {
        if !self.config.bounds_checking {
            return Ok(());
        }

        let max_dims = dims1.len().max(dims2.len());

        for i in 0..max_dims {
            let dim1 = if i < dims1.len() {
                dims1[dims1.len() - 1 - i]
            } else {
                1
            };

            let dim2 = if i < dims2.len() {
                dims2[dims2.len() - 1 - i]
            } else {
                1
            };

            if dim1 != dim2 && dim1 != 1 && dim2 != 1 {
                let violation_key = format!("broadcast_dims_{}", operation);
                let count = self.violation_counts.entry(violation_key).or_insert(0);
                *count += 1;

                error!(
                    "Broadcasting incompatible in {}: dim {} has {} vs {} (violation #{})",
                    operation, i, dim1, dim2, count
                );

                return Err(MLSafetyError::TensorSafety {
                    reason: format!(
                        "Incompatible shapes for broadcasting: {:?} and {:?}",
                        dims1, dims2
                    ),
                });
            }
        }

        debug!("Broadcast dims check passed: {:?} with {:?}", dims1, dims2);
        Ok(())
    }

    /// Check window size for sliding operations
    pub fn check_window_bounds(
        &mut self,
        sequence_len: usize,
        window_size: usize,
        stride: usize,
        operation: &str,
    ) -> SafetyResult<()> {
        if !self.config.bounds_checking {
            return Ok(());
        }

        if window_size == 0 {
            return Err(MLSafetyError::TensorSafety {
                reason: format!("Zero window size in {}", operation),
            });
        }

        if stride == 0 {
            return Err(MLSafetyError::TensorSafety {
                reason: format!("Zero stride in {}", operation),
            });
        }

        if window_size > sequence_len {
            let violation_key = format!("window_bounds_{}", operation);
            let count = self.violation_counts.entry(violation_key).or_insert(0);
            *count += 1;

            error!(
                "Window size violation in {}: window {} > sequence {} (violation #{})",
                operation, window_size, sequence_len, count
            );

            return Err(MLSafetyError::TensorSafety {
                reason: format!(
                    "Window size {} exceeds sequence length {}",
                    window_size, sequence_len
                ),
            });
        }

        debug!(
            "Window bounds check passed: {} window {} stride {} in sequence {}",
            operation, window_size, stride, sequence_len
        );
        Ok(())
    }

    /// Safe array element access with bounds checking
    pub fn safe_get<T: Clone>(
        &mut self,
        array: &[T],
        index: usize,
        operation: &str,
    ) -> SafetyResult<T> {
        self.check_array_bounds(array.len(), index, operation)?;
        Ok(array[index].clone())
    }

    /// Safe mutable array element access with bounds checking
    pub fn safe_get_mut<'a, T>(
        &mut self,
        array: &'a mut [T],
        index: usize,
        operation: &str,
    ) -> SafetyResult<&'a mut T> {
        self.check_array_bounds(array.len(), index, operation)?;
        Ok(&mut array[index])
    }

    /// Safe slice creation with bounds checking
    pub fn safe_slice<'a, T>(
        &mut self,
        array: &'a [T],
        start: usize,
        end: usize,
        operation: &str,
    ) -> SafetyResult<&'a [T]> {
        self.check_slice_bounds(array.len(), start, end, operation)?;
        Ok(&array[start..end])
    }

    /// Safe mutable slice creation with bounds checking
    pub fn safe_slice_mut<'a, T>(
        &mut self,
        array: &'a mut [T],
        start: usize,
        end: usize,
        operation: &str,
    ) -> SafetyResult<&'a mut [T]> {
        self.check_slice_bounds(array.len(), start, end, operation)?;
        Ok(&mut array[start..end])
    }

    /// Get violation statistics
    pub fn get_violation_stats(&self) -> HashMap<String, usize> {
        self.violation_counts.clone()
    }

    /// Reset violation counts
    pub fn reset_violations(&mut self) {
        self.violation_counts.clear();
        debug!("Bounds checker violation counts reset");
    }

    /// Check if violations exceed threshold
    pub fn check_violation_threshold(&self, threshold: usize) -> Vec<String> {
        self.violation_counts
            .iter()
            .filter(|(_, &count)| count >= threshold)
            .map(|(operation, count)| format!("{}: {} violations", operation, count))
            .collect()
    }

    /// Enable or disable bounds checking
    pub fn set_enabled(&mut self, enabled: bool) {
        self.config.bounds_checking = enabled;
        if enabled {
            debug!("Bounds checking enabled");
        } else {
            warn!("Bounds checking DISABLED - use only for performance testing");
        }
    }

    /// Check if bounds checking is enabled
    pub fn is_enabled(&self) -> bool {
        self.config.bounds_checking
    }
}

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

    fn create_test_checker() -> BoundsChecker {
        BoundsChecker::new(&MLSafetyConfig::default())
    }

    #[test]
    fn test_array_bounds() {
        let mut checker = create_test_checker();

        // Valid access
        assert!(checker.check_array_bounds(10, 5, "test").is_ok());

        // Out of bounds
        assert!(checker.check_array_bounds(10, 10, "test").is_err());
        assert!(checker.check_array_bounds(10, 15, "test").is_err());
    }

    #[test]
    fn test_slice_bounds() {
        let mut checker = create_test_checker();

        // Valid slice
        assert!(checker.check_slice_bounds(10, 2, 8, "test").is_ok());

        // Invalid start
        assert!(checker.check_slice_bounds(10, 15, 20, "test").is_err());

        // Invalid end
        assert!(checker.check_slice_bounds(10, 2, 15, "test").is_err());

        // Start > end
        assert!(checker.check_slice_bounds(10, 8, 2, "test").is_err());
    }

    #[test]
    fn test_tensor_bounds() {
        let mut checker = create_test_checker();

        let dims = &[3, 4, 5];

        // Valid indices
        assert!(checker
            .check_tensor_bounds(dims, &[0, 0, 0], "test")
            .is_ok());
        assert!(checker
            .check_tensor_bounds(dims, &[2, 3, 4], "test")
            .is_ok());

        // Out of bounds
        assert!(checker
            .check_tensor_bounds(dims, &[3, 0, 0], "test")
            .is_err());
        assert!(checker
            .check_tensor_bounds(dims, &[0, 4, 0], "test")
            .is_err());
        assert!(checker
            .check_tensor_bounds(dims, &[0, 0, 5], "test")
            .is_err());

        // Wrong number of indices
        assert!(checker.check_tensor_bounds(dims, &[0, 0], "test").is_err());
    }

    #[test]
    fn test_matmul_dims() {
        let mut checker = create_test_checker();

        // Valid matrix multiplication
        let lhs = &[3, 4];
        let rhs = &[4, 5];
        assert!(checker.check_matmul_dims(lhs, rhs, "test").is_ok());

        // Dimension mismatch
        let lhs = &[3, 4];
        let rhs = &[5, 6];
        assert!(checker.check_matmul_dims(lhs, rhs, "test").is_err());

        // Insufficient dimensions
        let lhs = &[3];
        let rhs = &[3, 4];
        assert!(checker.check_matmul_dims(lhs, rhs, "test").is_err());
    }

    #[test]
    fn test_safe_array_access() {
        let mut checker = create_test_checker();
        let array = vec![1, 2, 3, 4, 5];

        // Valid access
        let result = checker.safe_get(&array, 2, "test");
        assert!(result.is_ok());
        if let Ok(value) = result {
            assert_eq!(value, 3);
        }

        // Out of bounds access
        assert!(checker.safe_get(&array, 10, "test").is_err());
    }

    #[test]
    fn test_violation_tracking() {
        let mut checker = create_test_checker();

        // Generate some violations
        let _ = checker.check_array_bounds(10, 15, "test_op");
        let _ = checker.check_array_bounds(10, 20, "test_op");
        let _ = checker.check_array_bounds(5, 10, "other_op");

        let stats = checker.get_violation_stats();
        assert_eq!(stats.get("array_bounds_test_op"), Some(&2));
        assert_eq!(stats.get("array_bounds_other_op"), Some(&1));

        // Check threshold
        let violations = checker.check_violation_threshold(2);
        assert_eq!(violations.len(), 1);
        assert!(violations[0].contains("test_op"));
    }

    #[test]
    fn test_enable_disable() {
        let mut checker = create_test_checker();

        // Enabled by default
        assert!(checker.is_enabled());
        assert!(checker.check_array_bounds(10, 15, "test").is_err());

        // Disable bounds checking
        checker.set_enabled(false);
        assert!(!checker.is_enabled());
        assert!(checker.check_array_bounds(10, 15, "test").is_ok());

        // Re-enable
        checker.set_enabled(true);
        assert!(checker.is_enabled());
        assert!(checker.check_array_bounds(10, 15, "test").is_err());
    }
}
