//! OOM (Out-Of-Memory) Error Detection Utilities
//!
//! This module provides robust OOM error detection for AutoBatchSizer retry logic.
//! It handles various OOM error patterns from CUDA, CPU allocators, and ML frameworks.
//!
//! # Error Patterns Detected
//!
//! - **CUDA OOM**: "cuda error 2", "out of memory", "CUDA_ERROR_OUT_OF_MEMORY"
//! - **CPU OOM**: "failed to allocate", "memory allocation", "allocate"
//! - **Framework OOM**: "oom", "out_of_memory", "cudaMalloc"
//!
//! # Usage
//!
//! ```rust,ignore
//! use ml_core::memory_optimization::oom_detection::{is_oom_error, extract_oom_size};
//!
//! fn handle_training_error(err: &dyn std::fmt::Debug) {
//!     if is_oom_error(err) {
//!         println!("OOM detected! Halving batch size...");
//!         if let Some(size) = extract_oom_size(err) {
//!             println!("Requested memory: {} bytes", size);
//!         }
//!     }
//! }
//! ```

/// Check if an error is an OOM (Out-Of-Memory) error
///
/// This function uses comprehensive pattern matching to detect OOM errors across
/// CUDA, CPU, and framework-specific allocators. It matches error messages from:
///
/// - **CUDA runtime**: "cuda error 2", "CUDA_ERROR_OUT_OF_MEMORY", "cudaMalloc"
/// - **CPU allocators**: "failed to allocate", "memory allocation"
/// - **Generic OOM**: "out of memory", "oom", "out_of_memory"
///
/// # Arguments
///
/// * `err` - Any error type implementing `Debug`
///
/// # Returns
///
/// `true` if the error is an OOM error, `false` otherwise
pub fn is_oom_error(err: &dyn std::fmt::Debug) -> bool {
    let error_msg = format!("{:?}", err).to_lowercase();

    // Pattern matching based on test_gpu_oom_handling.rs (lines 401-406)
    // and batch_size_finder.rs (lines 157-161)
    error_msg.contains("out of memory")
        || error_msg.contains("oom")
        || error_msg.contains("cuda error 2")
        || error_msg.contains("failed to allocate")
        || error_msg.contains("cudamalloc")
        || error_msg.contains("out_of_memory")
        || error_msg.contains("cuda_error_out_of_memory")
        || error_msg.contains("memory allocation")
        || error_msg.contains("allocate")
}

/// Extract requested memory size from OOM error message (if available)
///
/// This function attempts to parse the requested memory size from OOM error messages.
/// It looks for common patterns like:
///
/// - "tried to allocate 1.2GB"
/// - "failed to allocate 1024 MB"
/// - "out of memory (requested 512MB)"
///
/// # Arguments
///
/// * `err` - Any error type implementing `Debug`
///
/// # Returns
///
/// `Some(size_in_bytes)` if a memory size was found, `None` otherwise
pub fn extract_oom_size(err: &dyn std::fmt::Debug) -> Option<usize> {
    let error_msg = format!("{:?}", err);

    // Try to extract memory size patterns
    // Pattern 1: "1.2GB", "512MB", "1024KB"
    if let Some(size) = extract_size_with_unit(&error_msg) {
        return Some(size);
    }

    // Pattern 2: "allocate 1024 bytes"
    if let Some(size) = extract_raw_bytes(&error_msg) {
        return Some(size);
    }

    None
}

/// Extract memory size with unit (GB, MB, KB)
fn extract_size_with_unit(msg: &str) -> Option<usize> {
    // Regex-free approach: search for number followed by unit
    let lower = msg.to_lowercase();

    // Find patterns like "1.2gb", "512mb", "1024kb"
    for (i, c) in lower.char_indices() {
        if c.is_ascii_digit() || c == '.' {
            // Start of potential number
            let rest = lower.get(i..).unwrap_or_default();

            // Extract number
            let num_end = rest
                .find(|c: char| !c.is_ascii_digit() && c != '.')
                .unwrap_or(rest.len());

            if num_end == 0 {
                continue;
            }

            let num_str = rest.get(..num_end).unwrap_or_default();
            let num = num_str.parse::<f64>().ok()?;

            // Check for unit immediately after number
            let unit_str = rest.get(num_end..).unwrap_or_default();
            if unit_str.starts_with("gb") || unit_str.starts_with("gib") {
                return Some((num * 1024.0 * 1024.0 * 1024.0) as usize);
            } else if unit_str.starts_with("mb") || unit_str.starts_with("mib") {
                return Some((num * 1024.0 * 1024.0) as usize);
            } else if unit_str.starts_with("kb") || unit_str.starts_with("kib") {
                return Some((num * 1024.0) as usize);
            } else {
                // Unrecognized unit
            }
        }
    }

    None
}

/// Extract raw byte count (e.g., "allocate 1024 bytes")
fn extract_raw_bytes(msg: &str) -> Option<usize> {
    let lower = msg.to_lowercase();

    // Look for "allocate" or "requested" followed by number and "bytes"
    for pattern in &["allocate ", "requested "] {
        if let Some(idx) = lower.find(pattern) {
            let rest = lower.get(idx + pattern.len()..).unwrap_or_default();

            // Extract number
            let num_end = rest
                .find(|c: char| !c.is_ascii_digit())
                .unwrap_or(rest.len());

            if num_end > 0 {
                let num_str = rest.get(..num_end).unwrap_or_default();
                if let Ok(num) = num_str.parse::<usize>() {
                    // Verify "bytes" follows
                    let unit_str = rest.get(num_end..).unwrap_or_default().trim_start();
                    if unit_str.starts_with("bytes") || unit_str.starts_with("byte") {
                        return Some(num);
                    }
                }
            }
        }
    }

    None
}

/// Helper struct for constructing OOM-like errors in tests.
///
/// Replaces the previous `CandleError::Msg` usage -- any `Debug` type works.
#[cfg(test)]
#[derive(Debug)]
struct TestError(String);

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

    // ========================================================================
    // OOM Detection Tests
    // ========================================================================

    #[test]
    fn test_cuda_oom_detection() {
        // CUDA error patterns from test_gpu_oom_handling.rs
        let cuda_errors = vec![
            TestError("CUDA error 2: out of memory".to_owned()),
            TestError("cuda error 2".to_owned()),
            TestError("CUDA_ERROR_OUT_OF_MEMORY".to_owned()),
            TestError("cudaMalloc failed".to_owned()),
        ];

        for err in &cuda_errors {
            assert!(is_oom_error(err), "Failed to detect CUDA OOM: {:?}", err);
        }
    }

    #[test]
    fn test_cpu_oom_detection() {
        // CPU allocation failure patterns
        let cpu_errors = vec![
            TestError("failed to allocate 1024 MB".to_owned()),
            TestError("memory allocation failed".to_owned()),
            TestError("could not allocate tensor".to_owned()),
        ];

        for err in &cpu_errors {
            assert!(is_oom_error(err), "Failed to detect CPU OOM: {:?}", err);
        }
    }

    #[test]
    fn test_generic_oom_detection() {
        // Generic OOM patterns
        let oom_errors = vec![
            TestError("out of memory".to_owned()),
            TestError("OOM occurred".to_owned()),
            TestError("Out_of_memory error".to_owned()),
        ];

        for err in &oom_errors {
            assert!(
                is_oom_error(err),
                "Failed to detect generic OOM: {:?}",
                err
            );
        }
    }

    #[test]
    fn test_non_oom_errors() {
        // Non-OOM error patterns
        let non_oom_errors = vec![
            TestError("dimension mismatch".to_owned()),
            TestError("invalid shape".to_owned()),
            TestError("file not found".to_owned()),
            TestError("network error".to_owned()),
        ];

        for err in &non_oom_errors {
            assert!(
                !is_oom_error(err),
                "False positive OOM detection: {:?}",
                err
            );
        }
    }

    // ========================================================================
    // Memory Size Extraction Tests
    // ========================================================================

    #[test]
    fn test_extract_size_gb() {
        let err = TestError("tried to allocate 1.2GB".to_owned());
        let size = extract_oom_size(&err);
        assert!(size.is_some(), "Failed to extract GB size");

        let size_bytes = size.unwrap();
        let expected = (1.2 * 1024.0 * 1024.0 * 1024.0) as usize;

        // Allow 1% tolerance for floating point
        let diff = if size_bytes > expected {
            size_bytes - expected
        } else {
            expected - size_bytes
        };
        assert!(
            diff < expected / 100,
            "GB extraction mismatch: got {}, expected ~{}",
            size_bytes,
            expected
        );
    }

    #[test]
    fn test_extract_size_mb() {
        let err = TestError("failed to allocate 512MB".to_owned());
        let size = extract_oom_size(&err);
        assert!(size.is_some(), "Failed to extract MB size");
        assert_eq!(size.unwrap(), 512 * 1024 * 1024);
    }

    #[test]
    fn test_extract_size_kb() {
        let err = TestError("out of memory (requested 2048KB)".to_owned());
        let size = extract_oom_size(&err);
        assert!(size.is_some(), "Failed to extract KB size");
        assert_eq!(size.unwrap(), 2048 * 1024);
    }

    #[test]
    fn test_extract_size_bytes() {
        let err = TestError("allocate 1024 bytes failed".to_owned());
        let size = extract_oom_size(&err);
        assert!(size.is_some(), "Failed to extract byte size");
        assert_eq!(size.unwrap(), 1024);
    }

    #[test]
    fn test_extract_size_no_info() {
        // Error messages without size information
        let errors = vec![
            TestError("out of memory".to_owned()),
            TestError("cuda error 2".to_owned()),
            TestError("allocation failed".to_owned()),
        ];

        for err in &errors {
            assert!(
                extract_oom_size(err).is_none(),
                "Should return None for error without size: {:?}",
                err
            );
        }
    }

    #[test]
    fn test_extract_size_case_insensitive() {
        let errors = vec![
            TestError("allocate 1GB".to_owned()),
            TestError("allocate 1gb".to_owned()),
            TestError("allocate 1Gb".to_owned()),
            TestError("allocate 1gB".to_owned()),
        ];

        for err in &errors {
            let size = extract_oom_size(err);
            assert!(
                size.is_some(),
                "Failed case-insensitive extraction: {:?}",
                err
            );
            assert_eq!(size.unwrap(), 1024 * 1024 * 1024);
        }
    }

    // ========================================================================
    // Edge Cases
    // ========================================================================

    #[test]
    fn test_multiple_sizes_in_message() {
        // Should extract first size found
        let err = TestError("tried to allocate 2GB but only 1GB available".to_owned());
        let size = extract_oom_size(&err);
        assert!(size.is_some());
        assert_eq!(size.unwrap(), 2 * 1024 * 1024 * 1024);
    }

    #[test]
    fn test_decimal_sizes() {
        let err = TestError("allocate 1.5GB failed".to_owned());
        let size = extract_oom_size(&err);
        assert!(size.is_some());

        let expected = (1.5 * 1024.0 * 1024.0 * 1024.0) as usize;
        let actual = size.unwrap();
        let diff = if actual > expected {
            actual - expected
        } else {
            expected - actual
        };
        assert!(diff < expected / 100, "Decimal size mismatch");
    }

    #[test]
    fn test_gib_vs_gb() {
        // Both GiB and GB should work (treated as binary)
        let err1 = TestError("allocate 1GiB".to_owned());
        let err2 = TestError("allocate 1GB".to_owned());

        assert_eq!(
            extract_oom_size(&err1),
            extract_oom_size(&err2),
            "GiB and GB should extract same size"
        );
    }
}
