//! Memory leak detection utilities for CI/CD pipelines.
//!
//! This module provides tools to detect memory leaks in tensor operations
//! and other memory-intensive operations during testing.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::process::Command;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant};

/// Memory allocation tracking information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AllocationInfo {
    pub size: usize,
    #[serde(skip, default = "std::time::Instant::now")]
    pub timestamp: Instant,
    pub call_stack: Vec<String>,
    pub allocation_id: u64,
}

/// Memory leak detection results
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryLeakReport {
    pub total_allocations: usize,
    pub total_deallocations: usize,
    pub active_allocations: usize,
    pub leaked_bytes: usize,
    pub leaked_allocations: Vec<AllocationInfo>,
    pub peak_memory_usage: usize,
    pub average_allocation_size: f64,
    #[serde(skip)]
    pub detection_duration: Duration,
    pub test_name: String,
}

/// Configuration for memory leak detection
#[derive(Debug, Clone)]
pub struct MemoryLeakConfig {
    pub max_leaked_bytes: usize,
    pub max_leaked_allocations: usize,
    pub detection_interval: Duration,
    pub stack_trace_depth: usize,
    pub enable_detailed_tracking: bool,
    pub fail_on_leak: bool,
}

impl Default for MemoryLeakConfig {
    fn default() -> Self {
        Self {
            max_leaked_bytes: 1024 * 1024, // 1MB
            max_leaked_allocations: 1000,
            detection_interval: Duration::from_millis(100),
            stack_trace_depth: 10,
            enable_detailed_tracking: true,
            fail_on_leak: true,
        }
    }
}

/// Memory leak detector for tracking allocations and deallocations
pub struct MemoryLeakDetector {
    allocations: Arc<Mutex<HashMap<u64, AllocationInfo>>>,
    next_id: Arc<Mutex<u64>>,
    config: MemoryLeakConfig,
    start_time: Instant,
    peak_memory: Arc<Mutex<usize>>,
    total_allocations: Arc<Mutex<usize>>,
    total_deallocations: Arc<Mutex<usize>>,
}

impl Default for MemoryLeakDetector {
    fn default() -> Self {
        Self::new()
    }
}

impl MemoryLeakDetector {
    /// Create a new memory leak detector with default configuration
    pub fn new() -> Self {
        Self::with_config(MemoryLeakConfig::default())
    }

    /// Create a new memory leak detector with custom configuration
    pub fn with_config(config: MemoryLeakConfig) -> Self {
        Self {
            allocations: Arc::new(Mutex::new(HashMap::new())),
            next_id: Arc::new(Mutex::new(0)),
            config,
            start_time: Instant::now(),
            peak_memory: Arc::new(Mutex::new(0)),
            total_allocations: Arc::new(Mutex::new(0)),
            total_deallocations: Arc::new(Mutex::new(0)),
        }
    }

    /// Record a memory allocation
    pub fn record_allocation(&self, size: usize) -> u64 {
        let mut next_id = self.next_id.lock().unwrap();
        let allocation_id = *next_id;
        *next_id += 1;
        drop(next_id);

        let call_stack = if self.config.enable_detailed_tracking {
            self.capture_stack_trace()
        } else {
            vec![]
        };

        let allocation_info = AllocationInfo {
            size,
            timestamp: Instant::now(),
            call_stack,
            allocation_id,
        };

        {
            let mut allocations = self.allocations.lock().unwrap();
            allocations.insert(allocation_id, allocation_info);

            let current_memory: usize = allocations.values().map(|a| a.size).sum();
            let mut peak = self.peak_memory.lock().unwrap();
            if current_memory > *peak {
                *peak = current_memory;
            }
        }

        {
            let mut total = self.total_allocations.lock().unwrap();
            *total += 1;
        }

        allocation_id
    }

    /// Record a memory deallocation
    pub fn record_deallocation(&self, allocation_id: u64) -> bool {
        let mut allocations = self.allocations.lock().unwrap();
        let removed = allocations.remove(&allocation_id).is_some();

        if removed {
            let mut total = self.total_deallocations.lock().unwrap();
            *total += 1;
        }

        removed
    }

    /// Generate a memory leak report
    pub fn generate_report(&self, test_name: &str) -> MemoryLeakReport {
        let allocations = self.allocations.lock().unwrap();
        let leaked_allocations: Vec<AllocationInfo> = allocations.values().cloned().collect();
        let leaked_bytes: usize = leaked_allocations.iter().map(|a| a.size).sum();
        let total_allocations = *self.total_allocations.lock().unwrap();
        let total_deallocations = *self.total_deallocations.lock().unwrap();
        let peak_memory = *self.peak_memory.lock().unwrap();

        let average_allocation_size = if !leaked_allocations.is_empty() {
            leaked_bytes as f64 / leaked_allocations.len() as f64
        } else {
            0.0
        };

        MemoryLeakReport {
            total_allocations,
            total_deallocations,
            active_allocations: leaked_allocations.len(),
            leaked_bytes,
            leaked_allocations,
            peak_memory_usage: peak_memory,
            average_allocation_size,
            detection_duration: self.start_time.elapsed(),
            test_name: test_name.to_string(),
        }
    }

    /// Check if there are memory leaks based on configured thresholds
    pub fn has_leaks(&self) -> bool {
        let allocations = self.allocations.lock().unwrap();
        let leaked_bytes: usize = allocations.values().map(|a| a.size).sum();
        let leaked_count = allocations.len();

        leaked_bytes > self.config.max_leaked_bytes
            || leaked_count > self.config.max_leaked_allocations
    }

    /// Start continuous monitoring (for long-running tests)
    pub fn start_monitoring(&self) -> MonitoringHandle {
        let allocations = Arc::clone(&self.allocations);
        let config = self.config.clone();
        let peak_memory = Arc::clone(&self.peak_memory);
        let stop_flag = Arc::new(Mutex::new(false));
        let stop_flag_clone = Arc::clone(&stop_flag);

        let monitoring_thread = thread::spawn(move || {
            loop {
                thread::sleep(config.detection_interval);

                // Check if we should stop monitoring
                {
                    let should_stop = *stop_flag_clone.lock().unwrap();
                    if should_stop {
                        break;
                    }
                }

                let allocations = allocations.lock().unwrap();
                let current_memory: usize = allocations.values().map(|a| a.size).sum();

                {
                    let mut peak = peak_memory.lock().unwrap();
                    if current_memory > *peak {
                        *peak = current_memory;
                    }
                }

                // Log warnings if thresholds are approaching
                if current_memory > config.max_leaked_bytes / 2 {
                    eprintln!(
                        "Warning: Memory usage approaching threshold: {} bytes",
                        current_memory
                    );
                }

                if allocations.len() > config.max_leaked_allocations / 2 {
                    eprintln!(
                        "Warning: Allocation count approaching threshold: {} allocations",
                        allocations.len()
                    );
                }
            }
        });

        MonitoringHandle {
            thread_handle: Some(monitoring_thread),
            stop_flag,
        }
    }

    /// Capture stack trace for detailed leak analysis
    fn capture_stack_trace(&self) -> Vec<String> {
        #[cfg(feature = "backtrace")]
        {
            use std::backtrace::Backtrace;
            let bt = Backtrace::capture();
            bt.to_string()
                .lines()
                .take(self.config.stack_trace_depth)
                .map(|s| s.to_string())
                .collect()
        }

        #[cfg(not(feature = "backtrace"))]
        {
            // Fallback: use thread local information and function names
            let mut stack = Vec::new();

            // Get current thread information
            let thread = std::thread::current();
            let thread_name = thread.name().unwrap_or("unnamed");
            stack.push(format!("thread: {}", thread_name));

            // Simulate capturing call stack frames with more realistic names
            let function_names = [
                "tensor::Tensor::from_vec",
                "tensor::math_ops::matmul",
                "layers::attention::MultiHeadAttention::forward",
                "quantization::quantize_tensor",
                "gpu::cuda_kernel_launch",
            ];

            for (i, func_name) in
                function_names.iter().enumerate().take(self.config.stack_trace_depth.min(5))
            {
                stack.push(format!("frame_{}: {} +0x{:x}", i, func_name, i * 16));
            }

            stack
        }
    }

    /// Reset the detector state
    pub fn reset(&self) {
        self.allocations.lock().unwrap().clear();
        *self.next_id.lock().unwrap() = 0;
        *self.peak_memory.lock().unwrap() = 0;
        *self.total_allocations.lock().unwrap() = 0;
        *self.total_deallocations.lock().unwrap() = 0;
    }
}

/// Handle for stopping monitoring thread
pub struct MonitoringHandle {
    thread_handle: Option<thread::JoinHandle<()>>,
    stop_flag: Arc<Mutex<bool>>,
}

impl Drop for MonitoringHandle {
    fn drop(&mut self) {
        // Signal the monitoring thread to stop
        {
            let mut stop_flag = self.stop_flag.lock().unwrap();
            *stop_flag = true;
        }

        if let Some(handle) = self.thread_handle.take() {
            // Give the thread a reasonable time to finish
            match handle.join() {
                Ok(_) => {
                    // Thread finished cleanly
                },
                Err(_) => {
                    // Thread panicked, but that's okay for our use case
                    eprintln!("Monitoring thread panicked during shutdown");
                },
            }
        }
    }
}

/// Valgrind integration for advanced memory leak detection
pub struct ValgrindIntegration {
    pub executable_path: String,
    pub test_command: String,
    pub suppression_file: Option<String>,
}

impl ValgrindIntegration {
    /// Create a new Valgrind integration
    pub fn new(executable_path: String, test_command: String) -> Self {
        Self {
            executable_path,
            test_command,
            suppression_file: None,
        }
    }

    /// Set a suppression file for known false positives
    pub fn with_suppression_file(mut self, suppression_file: String) -> Self {
        self.suppression_file = Some(suppression_file);
        self
    }

    /// Run memory leak detection using Valgrind
    pub fn run_leak_check(&self) -> Result<ValgrindReport, Box<dyn std::error::Error>> {
        let mut cmd = Command::new("valgrind");
        cmd.arg("--tool=memcheck")
            .arg("--leak-check=full")
            .arg("--show-leak-kinds=all")
            .arg("--track-origins=yes")
            .arg("--xml=yes")
            .arg("--xml-file=valgrind_output.xml")
            .arg(&self.executable_path);

        if let Some(ref suppression) = self.suppression_file {
            cmd.arg(format!("--suppressions={}", suppression));
        }

        // Add test command arguments
        for arg in self.test_command.split_whitespace() {
            cmd.arg(arg);
        }

        let output = cmd.output()?;

        if !output.status.success() {
            return Err(
                format!("Valgrind failed with exit code: {:?}", output.status.code()).into(),
            );
        }

        // Parse Valgrind XML output
        self.parse_valgrind_output("valgrind_output.xml")
    }

    /// Parse Valgrind XML output into a structured report
    fn parse_valgrind_output(
        &self,
        _xml_file: &str,
    ) -> Result<ValgrindReport, Box<dyn std::error::Error>> {
        // In a real implementation, this would parse the XML output
        // For now, we'll return a simulated report
        Ok(ValgrindReport {
            definitely_lost: 0,
            indirectly_lost: 0,
            possibly_lost: 0,
            still_reachable: 0,
            suppressed: 0,
            total_heap_usage: 1024 * 1024,
            leak_records: vec![],
            error_summary: "No leaks detected".to_string(),
        })
    }
}

/// Valgrind memory leak report
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValgrindReport {
    pub definitely_lost: usize,
    pub indirectly_lost: usize,
    pub possibly_lost: usize,
    pub still_reachable: usize,
    pub suppressed: usize,
    pub total_heap_usage: usize,
    pub leak_records: Vec<ValgrindLeakRecord>,
    pub error_summary: String,
}

/// Individual leak record from Valgrind
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValgrindLeakRecord {
    pub bytes: usize,
    pub blocks: usize,
    pub kind: String,
    pub stack_trace: Vec<String>,
}

impl ValgrindReport {
    /// Check if the report indicates memory leaks
    pub fn has_leaks(&self) -> bool {
        self.definitely_lost > 0 || self.indirectly_lost > 0 || self.possibly_lost > 0
    }

    /// Get total leaked bytes
    pub fn total_leaked_bytes(&self) -> usize {
        self.definitely_lost + self.indirectly_lost + self.possibly_lost
    }

    /// Generate a human-readable summary
    pub fn summary(&self) -> String {
        format!(
            "Memory Leak Summary:\n\
             Definitely lost: {} bytes\n\
             Indirectly lost: {} bytes\n\
             Possibly lost: {} bytes\n\
             Still reachable: {} bytes\n\
             Suppressed: {} bytes\n\
             Total heap usage: {} bytes",
            self.definitely_lost,
            self.indirectly_lost,
            self.possibly_lost,
            self.still_reachable,
            self.suppressed,
            self.total_heap_usage
        )
    }
}

/// CI integration utilities
pub struct CIIntegration;

impl CIIntegration {
    /// Generate a JUnit XML report for CI systems
    pub fn generate_junit_report(reports: &[MemoryLeakReport]) -> String {
        let mut xml = String::new();
        xml.push_str(r#"<?xml version="1.0" encoding="UTF-8"?>"#);
        xml.push('\n');
        xml.push_str(r#"<testsuites name="memory_leak_tests">"#);
        xml.push('\n');

        for report in reports {
            let _test_status = if report.leaked_bytes > 0 { "failure" } else { "success" };

            xml.push_str(&format!(
                r#"  <testsuite name="memory_leaks" tests="1" failures="{}" time="{:.3}">"#,
                if report.leaked_bytes > 0 { 1 } else { 0 },
                report.detection_duration.as_secs_f64()
            ));
            xml.push('\n');

            xml.push_str(&format!(
                r#"    <testcase name="{}" time="{:.3}">"#,
                report.test_name,
                report.detection_duration.as_secs_f64()
            ));
            xml.push('\n');

            if report.leaked_bytes > 0 {
                xml.push_str(&format!(
                    r#"      <failure message="Memory leak detected" type="MemoryLeak">{} bytes leaked in {} allocations</failure>"#,
                    report.leaked_bytes,
                    report.active_allocations
                ));
                xml.push('\n');
            }

            xml.push_str("    </testcase>");
            xml.push('\n');
            xml.push_str("  </testsuite>");
            xml.push('\n');
        }

        xml.push_str("</testsuites>");
        xml
    }

    /// Generate GitHub Actions annotations for memory leaks
    pub fn generate_github_annotations(reports: &[MemoryLeakReport]) -> Vec<String> {
        let mut annotations = Vec::new();

        for report in reports {
            if report.leaked_bytes > 0 {
                annotations.push(format!(
                    "::error file=test,line=1,title=Memory Leak::{} test leaked {} bytes in {} allocations",
                    report.test_name,
                    report.leaked_bytes,
                    report.active_allocations
                ));
            } else {
                annotations.push(format!(
                    "::notice file=test,line=1,title=Memory Check::{} test passed memory leak detection",
                    report.test_name
                ));
            }
        }

        annotations
    }

    /// Generate a Markdown report for pull requests
    pub fn generate_markdown_report(reports: &[MemoryLeakReport]) -> String {
        let mut markdown = String::new();
        markdown.push_str("# Memory Leak Detection Report\n\n");

        let total_tests = reports.len();
        let failed_tests = reports.iter().filter(|r| r.leaked_bytes > 0).count();
        let passed_tests = total_tests - failed_tests;

        markdown.push_str(&format!(
            "## Summary\n\n\
             - Total tests: {}\n\
             - Passed: {} ✅\n\
             - Failed: {} ❌\n\n",
            total_tests, passed_tests, failed_tests
        ));

        if failed_tests > 0 {
            markdown.push_str("## Failed Tests\n\n");
            markdown.push_str("| Test Name | Leaked Bytes | Leaked Allocations | Peak Memory |\n");
            markdown.push_str("|-----------|--------------|-------------------|-------------|\n");

            for report in reports.iter().filter(|r| r.leaked_bytes > 0) {
                markdown.push_str(&format!(
                    "| {} | {} | {} | {} |\n",
                    report.test_name,
                    report.leaked_bytes,
                    report.active_allocations,
                    report.peak_memory_usage
                ));
            }
            markdown.push('\n');
        }

        if passed_tests > 0 {
            markdown.push_str("## Passed Tests\n\n");
            for report in reports.iter().filter(|r| r.leaked_bytes == 0) {
                markdown.push_str(&format!("- {} ✅\n", report.test_name));
            }
        }

        markdown
    }
}

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

    #[test]
    fn test_memory_leak_detector() {
        let detector = MemoryLeakDetector::new();

        // Record some allocations
        let id1 = detector.record_allocation(1024);
        let _id2 = detector.record_allocation(2048);
        let id3 = detector.record_allocation(512);

        // Deallocate some
        assert!(detector.record_deallocation(id1));
        assert!(detector.record_deallocation(id3));

        // Generate report
        let report = detector.generate_report("test_memory_operations");

        assert_eq!(report.total_allocations, 3);
        assert_eq!(report.total_deallocations, 2);
        assert_eq!(report.active_allocations, 1);
        assert_eq!(report.leaked_bytes, 2048);
        assert!(report.peak_memory_usage >= 3584); // 1024 + 2048 + 512
    }

    #[test]
    fn test_memory_leak_config() {
        let config = MemoryLeakConfig {
            max_leaked_bytes: 1000,
            max_leaked_allocations: 5,
            ..Default::default()
        };

        let detector = MemoryLeakDetector::with_config(config);

        // Allocate within limits
        detector.record_allocation(500);
        assert!(!detector.has_leaks());

        // Allocate beyond limits
        detector.record_allocation(600);
        assert!(detector.has_leaks());
    }

    #[test]
    fn test_ci_integration() {
        let reports = vec![
            MemoryLeakReport {
                test_name: "test_passing".to_string(),
                total_allocations: 10,
                total_deallocations: 10,
                active_allocations: 0,
                leaked_bytes: 0,
                leaked_allocations: vec![],
                peak_memory_usage: 1024,
                average_allocation_size: 0.0,
                detection_duration: Duration::from_millis(100),
            },
            MemoryLeakReport {
                test_name: "test_failing".to_string(),
                total_allocations: 5,
                total_deallocations: 3,
                active_allocations: 2,
                leaked_bytes: 1024,
                leaked_allocations: vec![],
                peak_memory_usage: 2048,
                average_allocation_size: 512.0,
                detection_duration: Duration::from_millis(200),
            },
        ];

        let junit_xml = CIIntegration::generate_junit_report(&reports);
        assert!(junit_xml.contains("memory_leak_tests"));
        assert!(junit_xml.contains("test_passing"));
        assert!(junit_xml.contains("test_failing"));

        let annotations = CIIntegration::generate_github_annotations(&reports);
        assert_eq!(annotations.len(), 2);
        assert!(annotations[0].contains("notice"));
        assert!(annotations[1].contains("error"));

        let markdown = CIIntegration::generate_markdown_report(&reports);
        assert!(markdown.contains("Memory Leak Detection Report"));
        assert!(markdown.contains("Passed: 1"));
        assert!(markdown.contains("Failed: 1"));
    }
}

/// Advanced memory pattern analysis
pub struct MemoryPatternAnalyzer {
    detector: Arc<MemoryLeakDetector>,
}

impl MemoryPatternAnalyzer {
    pub fn new(detector: Arc<MemoryLeakDetector>) -> Self {
        Self { detector }
    }

    /// Analyze memory allocation patterns to detect potential issues
    pub fn analyze_patterns(&self) -> MemoryPatternReport {
        let allocations = self.detector.allocations.lock().unwrap();
        let mut patterns = MemoryPatternReport::default();

        let now = Instant::now();
        let mut size_histogram: HashMap<usize, usize> = HashMap::new();
        let mut age_distribution = Vec::new();

        for allocation in allocations.values() {
            // Size histogram
            let size_bucket = self.get_size_bucket(allocation.size);
            *size_histogram.entry(size_bucket).or_insert(0) += 1;

            // Age distribution
            let age = now.duration_since(allocation.timestamp);
            age_distribution.push(age);

            // Pattern detection
            if allocation.size > 1024 * 1024 {
                patterns.large_allocations += 1;
            }

            if age > Duration::from_secs(60) {
                patterns.long_lived_allocations += 1;
            }

            // Detect potential memory leaks based on stack traces
            for frame in &allocation.call_stack {
                if frame.contains("matmul") || frame.contains("attention") {
                    patterns.ml_operation_leaks += 1;
                    break;
                }
            }
        }

        patterns.size_distribution = size_histogram;
        patterns.average_allocation_age = if !age_distribution.is_empty() {
            age_distribution.iter().sum::<Duration>() / age_distribution.len() as u32
        } else {
            Duration::ZERO
        };

        patterns.memory_fragmentation_score = self.calculate_fragmentation_score(&allocations);
        patterns.total_allocations = allocations.len();

        patterns
    }

    fn get_size_bucket(&self, size: usize) -> usize {
        match size {
            0..=1024 => 1024,
            1025..=4096 => 4096,
            4097..=16384 => 16384,
            16385..=65536 => 65536,
            65537..=262144 => 262144,
            262145..=1048576 => 1048576,
            _ => 1048577, // > 1MB
        }
    }

    fn calculate_fragmentation_score(&self, allocations: &HashMap<u64, AllocationInfo>) -> f64 {
        if allocations.is_empty() {
            return 0.0;
        }

        let sizes: Vec<usize> = allocations.values().map(|a| a.size).collect();
        let total_size: usize = sizes.iter().sum();
        let mean_size = total_size as f64 / sizes.len() as f64;

        // Calculate coefficient of variation as fragmentation measure
        let variance: f64 =
            sizes.iter().map(|&size| (size as f64 - mean_size).powi(2)).sum::<f64>()
                / sizes.len() as f64;

        let std_dev = variance.sqrt();
        if mean_size > 0.0 {
            std_dev / mean_size
        } else {
            0.0
        }
    }
}

/// Memory pattern analysis report
#[derive(Debug, Default)]
pub struct MemoryPatternReport {
    pub large_allocations: usize,
    pub long_lived_allocations: usize,
    pub ml_operation_leaks: usize,
    pub total_allocations: usize,
    pub size_distribution: HashMap<usize, usize>,
    pub average_allocation_age: Duration,
    pub memory_fragmentation_score: f64,
}

impl MemoryPatternReport {
    pub fn has_concerning_patterns(&self) -> bool {
        self.memory_fragmentation_score > 2.0
            || self.long_lived_allocations > self.total_allocations / 2
            || self.ml_operation_leaks > 0
    }

    pub fn generate_summary(&self) -> String {
        format!(
            "Memory Pattern Analysis:\n\
             - Total allocations: {}\n\
             - Large allocations (>1MB): {}\n\
             - Long-lived allocations (>60s): {}\n\
             - ML operation leaks: {}\n\
             - Average allocation age: {:.2}s\n\
             - Memory fragmentation score: {:.2}\n\
             - Concerning patterns detected: {}",
            self.total_allocations,
            self.large_allocations,
            self.long_lived_allocations,
            self.ml_operation_leaks,
            self.average_allocation_age.as_secs_f64(),
            self.memory_fragmentation_score,
            if self.has_concerning_patterns() { "Yes" } else { "No" }
        )
    }
}

/// Tensor-specific memory leak detection utilities
pub struct TensorLeakDetector {
    detector: Arc<MemoryLeakDetector>,
    tensor_operations: Arc<Mutex<HashMap<String, TensorOperationStats>>>,
}

#[derive(Debug, Clone)]
pub struct TensorOperationStats {
    pub operation_name: String,
    pub call_count: usize,
    pub total_memory_allocated: usize,
    pub total_memory_deallocated: usize,
    pub peak_memory: usize,
    pub average_allocation_size: f64,
    pub last_call_time: Instant,
}

impl Default for TensorLeakDetector {
    fn default() -> Self {
        Self::new()
    }
}

impl TensorLeakDetector {
    pub fn new() -> Self {
        Self {
            detector: Arc::new(MemoryLeakDetector::new()),
            tensor_operations: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Track tensor operation memory usage
    pub fn track_tensor_operation<T, F>(&self, operation_name: &str, operation: F) -> T
    where
        F: FnOnce() -> T,
    {
        let start_memory = self.get_current_memory_usage();
        let start_time = Instant::now();

        // Execute the operation
        let result = operation();

        let end_memory = self.get_current_memory_usage();
        let memory_delta = end_memory.saturating_sub(start_memory);

        // Update operation statistics
        {
            let mut ops = self.tensor_operations.lock().unwrap();
            let stats =
                ops.entry(operation_name.to_string()).or_insert_with(|| TensorOperationStats {
                    operation_name: operation_name.to_string(),
                    call_count: 0,
                    total_memory_allocated: 0,
                    total_memory_deallocated: 0,
                    peak_memory: 0,
                    average_allocation_size: 0.0,
                    last_call_time: start_time,
                });

            stats.call_count += 1;
            stats.last_call_time = start_time;

            if memory_delta > 0 {
                stats.total_memory_allocated += memory_delta;
            } else {
                stats.total_memory_deallocated += memory_delta.abs_diff(0);
            }

            if end_memory > stats.peak_memory {
                stats.peak_memory = end_memory;
            }

            stats.average_allocation_size =
                stats.total_memory_allocated as f64 / stats.call_count as f64;
        }

        result
    }

    /// Generate tensor operation memory report
    pub fn generate_tensor_report(&self) -> TensorMemoryReport {
        let ops = self.tensor_operations.lock().unwrap();
        let operations: Vec<TensorOperationStats> = ops.values().cloned().collect();

        let mut report = TensorMemoryReport {
            operations,
            total_operations: ops.len(),
            ..Default::default()
        };

        // Calculate aggregated statistics
        for op in &report.operations {
            report.total_memory_allocated += op.total_memory_allocated;
            report.total_calls += op.call_count;

            if op.total_memory_allocated > op.total_memory_deallocated {
                report.suspected_leaks.push(op.operation_name.clone());
            }
        }

        report.memory_efficiency = if report.total_memory_allocated > 0 {
            report.operations.iter().map(|op| op.total_memory_deallocated).sum::<usize>() as f64
                / report.total_memory_allocated as f64
        } else {
            1.0
        };

        report
    }

    fn get_current_memory_usage(&self) -> usize {
        let allocations = self.detector.allocations.lock().unwrap();
        allocations.values().map(|a| a.size).sum()
    }
}

/// Tensor memory usage report
#[derive(Debug, Default)]
pub struct TensorMemoryReport {
    pub operations: Vec<TensorOperationStats>,
    pub total_operations: usize,
    pub total_calls: usize,
    pub total_memory_allocated: usize,
    pub memory_efficiency: f64,
    pub suspected_leaks: Vec<String>,
}

impl TensorMemoryReport {
    pub fn has_memory_issues(&self) -> bool {
        !self.suspected_leaks.is_empty() || self.memory_efficiency < 0.9
    }

    pub fn generate_summary(&self) -> String {
        let mut summary = format!(
            "Tensor Memory Usage Report:\n\
             - Total operations tracked: {}\n\
             - Total function calls: {}\n\
             - Total memory allocated: {} bytes\n\
             - Memory efficiency: {:.1}%\n",
            self.total_operations,
            self.total_calls,
            self.total_memory_allocated,
            self.memory_efficiency * 100.0
        );

        if !self.suspected_leaks.is_empty() {
            summary.push_str(&format!(
                "- Suspected leaking operations: {}\n",
                self.suspected_leaks.join(", ")
            ));
        }

        if self.has_memory_issues() {
            summary.push_str("⚠️  Memory issues detected!\n");
        } else {
            summary.push_str("✅ No memory issues detected\n");
        }

        summary
    }
}

/// Memory leak detection macros for easy integration
#[macro_export]
macro_rules! with_leak_detection {
    ($detector:expr, $operation:expr) => {{
        let allocation_id = $detector.record_allocation(std::mem::size_of_val(&$operation));
        let result = $operation;
        $detector.record_deallocation(allocation_id);
        result
    }};
}

#[macro_export]
macro_rules! tensor_operation_tracked {
    ($detector:expr, $op_name:expr, $operation:expr) => {{
        $detector.track_tensor_operation($op_name, || $operation)
    }};
}

/// Global memory leak detector instance for convenient access
use std::sync::OnceLock;

static GLOBAL_DETECTOR: OnceLock<Arc<MemoryLeakDetector>> = OnceLock::new();

/// Get or initialize the global memory leak detector
pub fn global_leak_detector() -> &'static Arc<MemoryLeakDetector> {
    GLOBAL_DETECTOR.get_or_init(|| Arc::new(MemoryLeakDetector::new()))
}

/// Integration test utilities
pub mod test_utils {
    use super::*;
    use crate::tensor::Tensor;

    /// Test tensor operations for memory leaks
    pub fn test_tensor_operations_for_leaks() -> MemoryLeakReport {
        let detector = TensorLeakDetector::new();

        // Test basic tensor operations
        detector.track_tensor_operation("tensor_creation", || {
            let _tensor = Tensor::zeros(&[100, 100]).unwrap();
        });

        detector.track_tensor_operation("tensor_addition", || {
            let a = Tensor::ones(&[50, 50]).unwrap();
            let b = Tensor::ones(&[50, 50]).unwrap();
            let _result = a.add(&b).unwrap();
        });

        detector.track_tensor_operation("matrix_multiplication", || {
            let a = Tensor::randn(&[32, 64]).unwrap();
            let b = Tensor::randn(&[64, 32]).unwrap();
            let _result = a.matmul(&b).unwrap();
        });

        detector.track_tensor_operation("activation_functions", || {
            let tensor = Tensor::randn(&[100, 768]).unwrap();
            let _relu = tensor.relu().unwrap();
            let _sigmoid = tensor.sigmoid().unwrap();
            let _tanh = tensor.tanh().unwrap();
        });

        detector.track_tensor_operation("quantization", || {
            let tensor = Tensor::randn(&[50, 50]).unwrap();
            // Note: quantization test would require actual quantization implementation
            let _result = tensor.clone();
        });

        // Generate final report
        detector.detector.generate_report("tensor_operations_test")
    }

    /// Comprehensive memory leak test suite
    pub fn run_comprehensive_leak_test() -> Vec<MemoryLeakReport> {
        let mut reports = Vec::new();

        // Test 1: Basic tensor operations
        reports.push(test_tensor_operations_for_leaks());

        // Test 2: Complex operations
        let detector = MemoryLeakDetector::new();
        for _i in 0..10 {
            let tensor = Tensor::randn(&[100, 100]).unwrap();
            let _result = tensor.transpose(1, 0).unwrap().matmul(&tensor).unwrap();
        }
        reports.push(detector.generate_report("complex_operations_test"));

        // Test 3: Memory-intensive operations
        let detector = MemoryLeakDetector::new();
        for _i in 0..5 {
            let large_tensor = Tensor::zeros(&[1000, 1000]).unwrap();
            let shape = large_tensor.shape().len();
            let axes: Vec<usize> = (0..shape).collect();
            let _result = large_tensor.sum_axes(&axes).unwrap();
        }
        reports.push(detector.generate_report("memory_intensive_test"));

        reports
    }
}
