thag_profiler 1.0.1

#![allow(
    clippy::branches_sharing_code,
    clippy::if_same_then_else,
    clippy::uninlined_format_args,
    unused_variables
)]
#![deny(unsafe_op_in_unsafe_fn)]
//! Task-aware memory allocator for profiling.
//!
//! This module provides a memory allocator that tracks allocations by logical tasks
//! rather than threads, making it suitable for async code profiling. It also contains
//! the custom memory allocator implementation that enables memory profiling.
use crate::{
    debug_log, file_stem_from_path, find_profile, flush_debug_log, fn_name,
    get_global_profile_type, get_root_module, internal_doc, is_detailed_memory,
    mem_attribution::{DetailedAddressRegistry, ProfileReg},
    profiling::{
        build_stack, clean_function_name, extract_detailed_alloc_callstack,
        get_memory_detail_dealloc_path, get_memory_detail_path, get_memory_path,
        is_profiling_state_enabled, MemoryDetailDeallocFile, MemoryDetailFile, MemoryProfileFile,
    },
    safe_alloc, warn_once, Profile, ProfileRef, ProfileType,
};
use backtrace::{resolve_frame, trace};
use parking_lot::Mutex;
use regex::Regex;
use std::{
    alloc::{GlobalAlloc, Layout, System},
    collections::{HashMap, HashSet},
    env, fmt,
    io::{self, Write},
    sync::{
        atomic::{AtomicUsize, Ordering},
        LazyLock,
    },
    time::Instant,
};
use thag_common::{lazy_static_var, re};

use std::{cell::Cell, thread_local};

/// Regular expression pattern to identify allocation start points in backtraces
pub static ALLOC_START_PATTERN: LazyLock<&'static Regex> =
    LazyLock::new(|| re!("thag_profiler::mem_tracking.+Dispatcher"));

// Thread-local storage for better async/threading isolation
// Each thread maintains its own flag, preventing cross-thread interference
thread_local! {
    static USING_SYSTEM_ALLOCATOR: Cell<bool> = const { Cell::new(false) };
}

/// Get the current state of the system allocator flag
///
/// Returns `true` if the system allocator is currently being used,
/// `false` if the tracking allocator is being used.
#[internal_doc]
#[inline]
#[must_use]
pub fn get_using_system() -> bool {
    USING_SYSTEM_ALLOCATOR.with(Cell::get)
}

/// Set the current state of the system allocator flag
///
/// # Arguments
///
/// * `value` - `true` to use the system allocator, `false` to use the tracking allocator
#[internal_doc]
#[inline]
pub fn set_using_system(value: bool) {
    USING_SYSTEM_ALLOCATOR.with(|cell| cell.set(value));
}

/// Try swapping the boolean TLS or global `USING_SYSTEM_ALLOCATOR` value and return the outcome as a Result.
///
/// # Errors
///
/// This function will return an error if `USING_SYSTEM_ALLOCATOR` was already set to the desired value.
/// We expect to handle this error in normal operation.
#[internal_doc]
#[inline]
pub fn compare_exchange_using_system(current: bool, new: bool) -> Result<bool, bool> {
    USING_SYSTEM_ALLOCATOR.with(|cell| {
        let actual = cell.get();
        if actual == current {
            cell.set(new);
            Ok(actual)
        } else {
            Err(actual)
        }
    })
}

/// Reset allocator state using the unified approach
#[internal_doc]
pub fn reset_allocator_state() {
    USING_SYSTEM_ALLOCATOR.with(|flag| flag.set(false));
}

// Maximum safe allocation size - 1 GB, anything larger is suspicious
const MAX_SAFE_ALLOCATION: usize = 1024 * 1024 * 1024;

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
/// Memory allocator types for the profiling system
///
/// This enum defines the different allocator backends that can be used
/// during profiling operations. The dispatcher switches between these
/// based on the current profiling state.
pub enum Allocator {
    /// Task-aware allocator that tracks which task allocated memory
    Tracking,
    /// System allocator for profiling operations
    System,
}

impl fmt::Display for Allocator {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Tracking => write!(f, "Tracking"),
            Self::System => write!(f, "System"),
        }
    }
}

/// Get the current allocator based on the configured approach
#[internal_doc]
#[must_use]
pub fn current_allocator() -> Allocator {
    let using_system = USING_SYSTEM_ALLOCATOR.with(Cell::get) || !crate::is_profiling_enabled();
    if using_system {
        Allocator::System
    } else {
        Allocator::Tracking
    }
}

/// Dispatcher allocator that routes allocation requests to the appropriate allocator
pub struct Dispatcher {
    /// Task-aware allocator that tracks allocations by logical tasks
    pub tracking: TrackingAllocator,
    /// Standard system allocator for fallback operations
    pub system: std::alloc::System,
}

impl Dispatcher {
    /// Creates a new dispatcher with default tracking and system allocators.
    #[must_use]
    pub const fn new() -> Self {
        Self {
            tracking: TrackingAllocator,
            system: std::alloc::System,
        }
    }
}

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

unsafe impl GlobalAlloc for Dispatcher {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let current = current_allocator();

        match current {
            Allocator::System => unsafe { self.system.alloc(layout) },
            Allocator::Tracking => {
                // // Use a recursive guard here to prevent infinite loops
                // let recursion_depth = RECURSION_DEPTH.load(Ordering::Relaxed);
                // if recursion_depth > 10 {
                //     // Emergency fallback to system allocator
                //     unsafe { self.system.alloc(layout) }
                // } else {
                //     RECURSION_DEPTH.store(recursion_depth + 1, Ordering::SeqCst);
                //     let ptr = unsafe { self.tracking.alloc(layout) };
                //     let recursion_depth = RECURSION_DEPTH.load(Ordering::Relaxed);
                //     if recursion_depth > 0 {
                //         RECURSION_DEPTH.store(recursion_depth - 1, Ordering::SeqCst);
                //     }
                //     ptr
                // }
                unsafe { self.tracking.alloc(layout) }
            }
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        if ptr.is_null() {
            return;
        }

        // Safety check for unreasonably large deallocations
        if layout.size() > MAX_SAFE_ALLOCATION {
            safe_alloc! {
                eprintln!(
                    "WARNING: Extremely large deallocation request of {} bytes",
                    layout.size()
                )
            }
            // Still need to deallocate it to avoid memory leaks
        }

        match current_allocator() {
            Allocator::System => unsafe { self.system.dealloc(ptr, layout) },
            Allocator::Tracking => {
                // // Use a recursive guard here to prevent infinite loops
                // let recursion_depth = RECURSION_DEPTH.load(Ordering::Relaxed);
                // if recursion_depth > 10 {
                //     // Emergency fallback to system allocator
                //     unsafe { self.system.dealloc(ptr, layout) }
                // } else {
                //     RECURSION_DEPTH.store(recursion_depth + 1, Ordering::SeqCst);
                //     unsafe { self.tracking.dealloc(ptr, layout) };
                //     let recursion_depth = RECURSION_DEPTH.load(Ordering::Relaxed);
                //     if recursion_depth > 0 {
                //         RECURSION_DEPTH.store(recursion_depth - 1, Ordering::SeqCst);
                //     }
                // }
                unsafe { self.tracking.dealloc(ptr, layout) }
            }
        }
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        if ptr.is_null() {
            return unsafe {
                self.alloc(Layout::from_size_align_unchecked(new_size, layout.align()))
            };
        }

        // Safety check for unreasonably large reallocations
        // if new_size > MAX_SAFE_ALLOCATION {
        //     safe_alloc! {
        //         eprintln!(
        //             "WARNING: Extremely large reallocation request of {} bytes",
        //             layout.size()
        //         )
        //     };
        //     return std::ptr::null_mut();
        // }

        match current_allocator() {
            Allocator::System => unsafe { self.system.realloc(ptr, layout, new_size) },
            Allocator::Tracking => {
                // // Use a recursive guard here to prevent infinite loops
                // let recursion_depth = RECURSION_DEPTH.load(Ordering::Relaxed);
                // if recursion_depth > 10 {
                //     // Emergency fallback to system allocator
                //     unsafe { self.system.realloc(ptr, layout, new_size) }
                // } else {
                //     RECURSION_DEPTH.store(recursion_depth + 1, Ordering::SeqCst);
                //     let ptr = unsafe { self.tracking.realloc(ptr, layout, new_size) };
                //     let recursion_depth = RECURSION_DEPTH.load(Ordering::Relaxed);
                //     if recursion_depth > 0 {
                //         RECURSION_DEPTH.store(recursion_depth - 1, Ordering::SeqCst);
                //     }
                //     ptr
                // }
                unsafe { self.tracking.realloc(ptr, layout, new_size) }
            }
        }
    }
}

/// Task-aware allocator that tracks memory allocations
pub struct TrackingAllocator;

// Static instance for global access
static TRACKING_ALLOCATOR: TrackingAllocator = TrackingAllocator;

/// Helper to get the allocator instance
#[must_use]
pub fn get_allocator() -> &'static TrackingAllocator {
    &TRACKING_ALLOCATOR
}

#[allow(clippy::unused_self)]
impl TrackingAllocator {
    /// Creates a new task context for tracking memory
    #[internal_doc]
    pub fn create_task_context(&'static self) -> TaskMemoryContext {
        let task_id = TASK_STATE.next_task_id.fetch_add(1, Ordering::SeqCst);

        // Initialize in profile registry
        activate_task(task_id);

        TaskMemoryContext { task_id }
    }
}

unsafe impl GlobalAlloc for TrackingAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let ptr = unsafe { System.alloc(layout) };

        safe_alloc! {
            if !ptr.is_null() && is_profiling_state_enabled() {
                let size = layout.size();
                // Potentially skip small allocations
                if size > *SIZE_TRACKING_THRESHOLD {
                    let address = ptr as usize;
                    record_alloc(address, size);
                }
            }
            // See ya later allocator
        };
        ptr
    }

    #[allow(clippy::too_many_lines)]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        safe_alloc! {
            if !ptr.is_null()
                && is_profiling_state_enabled()
                // Only record detailed deallocations to -memory_detail_dealloc.folded if requested
                && lazy_static_var!(bool, deref, is_detailed_memory())
            {
                // Potentially skip small allocations
                let size = layout.size();
                if size > *SIZE_TRACKING_THRESHOLD {
                    let address = ptr as usize;
                    record_dealloc(address, size);
                }
            }
        };

        // Forward to system allocator for deallocation
        unsafe { System.dealloc(ptr, layout) };
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        safe_alloc! {
            if !ptr.is_null()
                && is_profiling_state_enabled()
                // Only record detailed deallocations to -memory_detail_dealloc.folded if requested
                && lazy_static_var!(bool, deref, is_detailed_memory())
            {
                // Potentially skip small allocations
                let dealloc_size = layout.size();
                if dealloc_size > *SIZE_TRACKING_THRESHOLD {
                    let address = ptr as usize;
                    record_dealloc(address, dealloc_size);
                }
            }

            // Potentially skip small allocations
            if new_size > *SIZE_TRACKING_THRESHOLD {
                let address = ptr as usize;
                record_alloc(address, new_size);
            }
        };

        unsafe { System.realloc(ptr, layout, new_size) }
    }
}

#[allow(
    clippy::cognitive_complexity,
    clippy::too_many_lines,
    unreachable_code,
    unused_variables
)]
fn record_alloc(address: usize, size: usize) {
    // Simple recursion prevention without using TLS with destructors
    static mut IN_TRACKING: bool = false;
    struct Guard;
    impl Drop for Guard {
        fn drop(&mut self) {
            unsafe {
                IN_TRACKING = false;
            }
        }
    }

    safe_alloc! {
        if size == 0 {
            debug_log!("Zero-sized allocation found");
            return;
        }

        let profile_type = lazy_static_var!(ProfileType, deref, get_global_profile_type());

        if profile_type != ProfileType::Memory && profile_type != ProfileType::Both {
            // debug_log!(
            //     "Skipping allocation recording because profile_type={:?}",
            //     profile_type
            // );
            return;
        }

        // Flag if we're already tracking in case it causes an infinite recursion
        let in_tracking = unsafe { IN_TRACKING };

        // Assertion disabled because not 100%
        // #[cfg(debug_assertions)]
        // assert!(!in_tracking);

        if in_tracking {
            debug_log!("*** Caution: already tracking: proceeding for allocation of {size} B");
            // return ptr;
        }

        // Set tracking flag and create guard for cleanup
        unsafe {
            IN_TRACKING = true;
        }
        let _guard = Guard;

        // Get backtrace without recursion
        // debug_log!("Attempting backtrace");
        let start_ident = Instant::now();

        let file_names = {
            safe_alloc! {
                ProfileReg::get()
                    // .lock()
                    .get_file_names()
            }
        };
        debug_log!("file_names={file_names:#?}");

        // let Some((filename, lineno, frame, fn_name, profile_ref)) = Backtrace::frames(&current_backtrace)
        let Some(frames) =
            extract_callstack_with_recursion_check(&file_names)
        else {
            debug_log!("Recursion detected");
            return;
        };

        safe_alloc! {
            if frames.is_empty() {
                debug_log!("No eligible profile found");
                return;
            }
            // debug_log!("func_and_ancestors={func_and_ancestors:#?}");

            let in_profile_code = frames.iter().any(|(_, _, frame, _, _)| {
                frame.contains("Profile::new")
            });

            if in_profile_code {
                debug_log!("Ignoring allocation request of size {size} for profiler code");
                return;
            }

            let (filename, lineno, frame, fn_name, profile_ref) = &frames[0];
            let detailed_memory = lazy_static_var!(bool, deref, is_detailed_memory());

            // if size == 40 {
            //     debug_log!("frames: {frames:#?} for size {size}");
            // }
            debug_log!("Found filename (file_name)={filename}, lineno={lineno}, fn_name: {fn_name:?}, frame: {frame:?} for size {size}");

            // Still record detailed allocations to -memory_detail.folded if requested
            if detailed_memory {
                record_detailed_alloc(
                    address,
                    size,
                    &ALLOC_START_PATTERN,
                    true,
                );
            }

            // Try to record the allocation in the new profile registry
            if !filename.is_empty()
                && *lineno > 0
                && record_allocation(filename, fn_name, *lineno, size)
            {
                debug_log!("Recorded allocation of {size} bytes in {filename}::{fn_name}:{lineno} to a profile");

                debug_log!(
                    "size={size}, time to assign = {}ms",
                    start_ident.elapsed().as_millis()
                );
            }
        };
    };
}

type FrameSummary = (String, u32, String, String, ProfileRef);

#[fn_name]
pub fn extract_callstack_with_recursion_check(file_names: &[String]) -> Option<Vec<FrameSummary>> {
    safe_alloc! {
        // Pre-allocate with fixed capacity to avoid reallocations
        let capacity = 100;
        let mut frames: Vec<(String, u32, String, String, ProfileRef)> = Vec::with_capacity(capacity); // Fixed size, no growing
        let mut found_recursion = false;
        let mut fin = false;
        let mut i = 0;

        trace(|frame| {
            let mut suppress = false;

            resolve_frame(frame, |symbol| {

                'process_symbol: {
                    let Some(name) = symbol.name() else {
                        suppress = true;
                        break 'process_symbol;
                    };
                    let name = name.to_string();
                    if name.contains("__rust_begin_short_backtrace") {
                        fin = true;
                        suppress = true;
                    }
                    if name.starts_with("backtrace::backtrace::") || name.starts_with('<') {
                        suppress = true;
                    }

                    if suppress { break 'process_symbol; }

                    // Check for our own functions (recursion detection)
                    if i > 0 && name.contains(fn_name) {
                        found_recursion = true;
                        break 'process_symbol;
                    }

                    let maybe_filename = symbol.filename();
                    let maybe_lineno = symbol.lineno();

                    // Apply the first filter
                    if maybe_filename.is_none()
                        || maybe_lineno.is_none()
                    {
                        suppress = true;
                        break 'process_symbol;
                    }
                    // Safe to unwrap now
                    let filename = safe_alloc! { file_stem_from_path(maybe_filename.unwrap()) };
                    let lineno = safe_alloc! { maybe_lineno.unwrap() };

                    if !file_names.contains(&filename) {
                        suppress = true;
                        break 'process_symbol;
                    }

                    // Apply second filter
                    let fn_name = clean_function_name(&mut name.clone());
                    let maybe_profile_ref = find_profile(&filename, &fn_name, lineno);
                    if let Some(profile_ref) = maybe_profile_ref {
                        // Safe to add this frame
                        frames.push((filename, lineno, name, fn_name, profile_ref));
                        i += 1;
                        if i >= capacity {
                            safe_alloc! {
                                 println!("frames={frames:#?}");
                             };
                             panic!("Max limit of {capacity} frames exceeded");
                        }
                    } else {
                        debug_log!("No profile found for {filename}, {fn_name}, {lineno}");
                    }
                }
            });
            !found_recursion && !fin
        });
        if found_recursion {
            None // Signal to skip tracking
        } else {
            Some(frames)
        }
    }
}

/// Record an allocation with the profile registry based on module path and line number
#[internal_doc]
#[must_use]
pub fn record_allocation(file_name: &str, fn_name: &str, line: u32, size: usize) -> bool {
    safe_alloc! {
        // First log (acquires debug log mutex)
        debug_log!(
            "Looking for profile to record allocation: module={file_name}, fn={fn_name}, line={line}, size={size}"
        );

        // Flush to release the debug log mutex
        flush_debug_log();

        // Print list of registered modules to help diagnose issues
        {
            let modules = ProfileReg::get()
                .get_file_names();
            debug_log!("Available modules in registry: {modules:?}");
            flush_debug_log();
        }

        // Now acquire the PROFILE_REGISTRY mutex
        let result;
        {
            debug_log!("About to call record_allocation on registry");
            // result = crate::mem_attribution::ProfileReg::get()
            result = ProfileReg::get().record_allocation(
                file_name,
                fn_name,
                line,
                size,
            );
            debug_log!("record_allocation on registry returned {result}");
        }

        // Log after releasing the mutex
        if result {
            debug_log!(
                "Successfully recorded allocation of {size} bytes in module {file_name}::{fn_name} at line {line}"
            );
        } else {
            debug_log!("No matching profile found to record allocation of {size} bytes in module {file_name}::{fn_name} at line {line}");
        }
        // flush_debug_log();

        result
    }
}

/// Register a detailed allocation with its stack trace and size for later deallocation tracking
///
/// This function stores allocation details in the detailed address registry, which is used
/// to track stack traces for deallocations when detailed memory profiling is enabled.
///
/// # Arguments
///
/// * `address` - The memory address of the allocation
/// * `size` - The size in bytes of the allocation
/// * `stack` - The call stack trace at the time of allocation
pub fn register_detailed_allocation(address: usize, size: usize, stack: Vec<String>) {
    safe_alloc! {
        if is_detailed_memory() {
            DetailedAddressRegistry::get().insert(address, (stack, size));
        }
    }
}

/// Records a detailed memory allocation with its call stack for profiling purposes.
///
/// This function captures the call stack at the allocation point, writes it to the appropriate
/// detail file if requested, and registers the allocation for later deallocation tracking.
///
/// # Arguments
///
/// * `address` - The memory address of the allocation
/// * `size` - The size in bytes of the allocation
/// * `start_pattern` - Regular expression pattern to identify allocation start points in backtraces
/// * `write_to_detail_file` - Whether to write the allocation details to the detail file
pub fn record_detailed_alloc(
    address: usize,
    size: usize,
    start_pattern: &Regex,
    write_to_detail_file: bool,
) {
    let detailed_stack = extract_detailed_alloc_callstack(start_pattern);
    write_detailed_stack_alloc(size, write_to_detail_file, &detailed_stack);
    register_detailed_allocation(address, size, detailed_stack);
}

#[allow(
    clippy::ptr_arg,
    clippy::missing_panics_doc,
    reason = "debug_assertions"
)]
/// Write detailed stack allocation information to the appropriate memory profile file.
///
/// This function formats and writes allocation details including the call stack and size
/// to either the detailed memory file or the standard memory profile file based on the
/// `write_to_detail_file` parameter.
///
/// # Arguments
///
/// * `size` - The size in bytes of the allocation
/// * `write_to_detail_file` - If true, writes to the detailed memory file; otherwise writes to the standard memory profile file
/// * `detailed_stack` - The call stack trace as a vector of strings
pub fn write_detailed_stack_alloc(
    size: usize,
    write_to_detail_file: bool,
    detailed_stack: &Vec<String>,
) {
    safe_alloc! {
        let root_module = lazy_static_var!(
            String,
            get_root_module()
                .as_ref()
                .map_or("root module", |v| v)
                .to_string()
        );

        let entry = if detailed_stack.is_empty() {
            format!("[Out of `{root_module}` scope] {size}")
        } else {
            let descr_stack = build_stack(detailed_stack, None, ";");

            debug_log!("descr_stack={descr_stack}");
            format!("{descr_stack} {size}")
        };

        let (memory_path, file) = if write_to_detail_file {
            (get_memory_detail_path().unwrap(), MemoryDetailFile::get())
        } else {
            (get_memory_path().unwrap(), MemoryProfileFile::get())
        };
        let _ = Profile::write_profile_event(memory_path, file, &entry);
    }
}

#[allow(
    clippy::too_many_lines,
    clippy::missing_panics_doc,
    reason = "debug_assertions"
)]
/// Records a memory deallocation for detailed profiling purposes.
///
/// This function tracks memory deallocations by capturing call stack information
/// and writing it to the detailed memory deallocation profile file when detailed
/// memory profiling is enabled. It includes recursion prevention and filtering
/// to avoid tracking profiler-internal allocations.
///
/// # Arguments
///
/// * `address` - The memory address being deallocated
/// * `size` - The size in bytes of the deallocation
pub fn record_dealloc(address: usize, size: usize) {
    // Simple recursion prevention without using TLS with destructors
    static mut IN_TRACKING: bool = false;
    struct Guard;
    impl Drop for Guard {
        fn drop(&mut self) {
            unsafe {
                IN_TRACKING = false;
            }
        }
    }

    // Assertion disabled because not 100%
    // #[cfg(debug_assertions)]
    // assert_eq!(current_allocator(), Allocator::System);

    let root_module = lazy_static_var!(
        String,
        get_root_module()
            .as_ref()
            .map_or("root module", |v| v)
            .to_string()
    );

    let profile_type = lazy_static_var!(ProfileType, deref, get_global_profile_type());
    let is_mem_prof = lazy_static_var!(bool, {
        profile_type == ProfileType::Memory || profile_type == ProfileType::Both
    });

    // Use the warn_once! macro for clean, optimized warning suppression
    warn_once!(
        !is_mem_prof,
        || {
            debug_log!("Skipping deallocation recording because profile_type={profile_type:?}");
        },
        return
    );

    // Flag if we're already tracking in case it causes an infinite recursion
    let in_tracking = unsafe { IN_TRACKING };

    // Assertion disabled because not 100%
    // #[cfg(debug_assertions)]
    // assert!(!in_tracking);

    if in_tracking {
        debug_log!("*** Caution: already tracking: proceeding for deallocation of {size} B");
        // return ptr;
    }

    // Set tracking flag and create guard for cleanup
    unsafe {
        IN_TRACKING = true;
    }
    let _guard = Guard;

    // Get backtrace without recursion
    // debug_log!("Attempting backtrace");
    // let start_ident = Instant::now();
    // let mut task_id = 0;
    // Now we can safely use backtrace without recursion!
    let start_pattern: &Regex = re!("thag_profiler::mem_tracking.+Dispatcher");

    let detailed_memory = lazy_static_var!(bool, deref, is_detailed_memory());
    if size > 0 && detailed_memory {
        let detailed_stack = extract_detailed_alloc_callstack(start_pattern);

        let in_profile_code = detailed_stack
            .iter()
            .any(|frame| frame.contains("::profiling::Profile"));

        if in_profile_code {
            debug_log!(
                "Detailed memory tracking ignoring detailed deallocation request of size {size} for profiler code: frame={:?}",
                detailed_stack
                    .iter()
                    .find(|frame| frame.contains("::profiling::Profile"))
            );
            // debug_log!("...current backtrace: {:#?}", current_backtrace);
            return;
        }

        let entry = if detailed_stack.is_empty() {
            let stack_and_size = {
                DetailedAddressRegistry::get()
                    .remove(&address)
                    .unwrap_or((0, (Vec::new(), size)))
            };

            let (stack, _) = stack_and_size.1;

            let legend = if stack.is_empty() {
                // debug_log!("Empty cleaned_stack and stack for backtrace={current_backtrace:#?}");
                format!("[Dealloc out of `{root_module}` scope]")
            } else {
                stack.join(";")
            };
            format!("{legend} {size}")
        } else {
            format!("{} {size}", detailed_stack.join(";"))
        };

        let memory_detail_dealloc_path = get_memory_detail_dealloc_path().unwrap();
        let _ = Profile::write_profile_event(
            memory_detail_dealloc_path,
            MemoryDetailDeallocFile::get(),
            &entry,
        );
    }
}

// // Create a direct static instance
#[global_allocator]
static ALLOCATOR: Dispatcher = Dispatcher::new();

// ========== ALLOCATION TRACKING DEFINITIONS ==========

/// Threshold for size-based memory allocation tracking.
///
/// This static value determines the minimum allocation size that will be tracked
/// during memory profiling. Allocations smaller than this threshold are ignored
/// to reduce profiling overhead. The value can be configured via the
/// `SIZE_TRACKING_THRESHOLD` environment variable, defaulting to 0 if not set.
///
/// When set to 0, all allocations are tracked regardless of size.
/// When set to a positive value, only allocations exceeding that threshold are tracked.
pub static SIZE_TRACKING_THRESHOLD: LazyLock<usize> = LazyLock::new(|| {
    let threshold = env::var("SIZE_TRACKING_THRESHOLD")
        .or_else(|_| Ok::<String, &str>(String::from("0")))
        .ok()
        .and_then(|val| val.parse::<usize>().ok())
        .expect("Value specified for SIZE_TRACKING_THRESHOLD must be a valid integer");
    if threshold == 0 {
        debug_log!("*** The SIZE_TRACKING_THRESHOLD environment variable is set or defaulted to 0, so all memory allocations and deallocations will be tracked.");
    } else {
        debug_log!("*** Only memory allocations and deallocations exceeding the specified threshold of {threshold} bytes will be tracked.");
    }
    threshold
});

// ========== PUBLIC REGISTRY API ==========

/// Add a task to active profiles
pub fn activate_task(task_id: usize) {
    safe_alloc! {
        ProfileReg::get().activate_task(task_id);
    };
}

/// Remove a task from active profiles
#[allow(dead_code)]
pub fn deactivate_task(task_id: usize) {
    safe_alloc! {
        ProfileReg::get().deactivate_task(task_id);
    };
}

// /// Get active tasks
// #[must_use]
// pub fn get_active_tasks() -> Vec<usize> {
//     safe_alloc! { ProfileReg::get().get_active_tasks() }
// }

/// Get the last active task
#[internal_doc]
#[must_use]
pub fn get_last_active_task() -> Option<usize> {
    safe_alloc! { ProfileReg::get().get_last_active_task() }
}

// ========== TASK CONTEXT DEFINITIONS ==========

/// Task context for tracking allocations
#[internal_doc]
#[derive(Debug, Clone)]
pub struct TaskMemoryContext {
    /// Unique identifier for this task context
    pub task_id: usize,
}

impl TaskMemoryContext {
    /// Gets the unique ID for this task
    #[must_use]
    pub const fn id(&self) -> usize {
        self.task_id
    }
}

/// Task context for tracking allocations
#[cfg(not(feature = "full_profiling"))]
#[derive(Debug, Default, Clone, Copy)]
pub struct TaskMemoryContext;

/// Creates a new task context for memory tracking.
#[internal_doc]
#[must_use]
pub fn create_memory_task() -> TaskMemoryContext {
    let allocator = get_allocator();
    allocator.create_task_context()
}

// ========== TASK STATE MANAGEMENT ==========

/// Task state management for memory profiling
///
/// This struct maintains global state for tracking memory allocations
/// across different tasks in the profiling system.
pub struct TaskState {
    /// Counter for generating unique task IDs
    ///
    /// This atomic counter ensures each task gets a unique identifier
    /// for tracking memory allocations in multi-threaded environments.
    pub next_task_id: AtomicUsize,
}

/// Global task state for memory profiling
///
/// This static instance manages the global state for tracking memory allocations
/// across different tasks in the profiling system. It provides thread-safe
/// access to task ID generation and other task-related state.
pub static TASK_STATE: LazyLock<TaskState> = LazyLock::new(|| TaskState {
    next_task_id: AtomicUsize::new(1),
});

/// Task guard that automatically deactivates a task when dropped
///
/// This guard ensures that tasks are properly cleaned up when they go out of scope,
/// preventing memory leaks in the task tracking system.
#[internal_doc]
#[derive(Clone, Debug)]
pub struct TaskGuard {
    task_id: usize,
}

impl TaskGuard {
    /// Creates a new task guard for the given task ID
    ///
    /// The guard will automatically deactivate the task when dropped.
    ///
    /// # Arguments
    ///
    /// * `task_id` - The unique identifier for the task to guard
    #[must_use]
    pub const fn new(task_id: usize) -> Self {
        Self { task_id }
    }
}

/// Task guard that automatically deactivates a task when dropped
///
/// This guard ensures that tasks are properly cleaned up when they go out of scope,
/// preventing memory leaks in the task tracking system.
#[cfg(not(feature = "full_profiling"))]
#[derive(Debug, Default, Clone, Copy)]
pub struct TaskGuard;

impl Drop for TaskGuard {
    fn drop(&mut self) {
        // Run these operations with System allocator
        safe_alloc! {
            // Remove from active profiles
            ProfileReg::get().deactivate_task(self.task_id);
            debug_log!("Deactivated task {}", self.task_id);

            // Flush logs directly
            if let Some(logger) = crate::DebugLogger::get() {
                let _ = logger.lock().flush();
            }
        };
    }
}

// ========== TASK PATH MANAGEMENT ==========

/// Registry mapping task IDs to their execution paths for flamegraph generation.
///
/// Each task ID maps to a vector of strings representing the call stack path.
/// This ensures all paths are written to the .folded file, even with zero allocations,
/// to provide complete call hierarchy information for flamegraph construction.
pub static TASK_PATH_REGISTRY: LazyLock<Mutex<HashMap<usize, Vec<String>>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

// ========== MEMORY PROFILING LIFECYCLE ==========

/// Initialize memory profiling.
/// This is called by the main `init_profiling` function.
#[allow(clippy::missing_panics_doc)]
pub fn initialize_memory_profiling() {
    // Set up allocator state with Tracking as the default using unified approach
    reset_allocator_state();

    // Use system allocator just for logging
    safe_alloc! {
        debug_log!("Memory profiling initialized");
        flush_debug_log();
    };
    #[cfg(debug_assertions)]
    assert_eq!(current_allocator(), Allocator::Tracking);
}

/// Finalize memory profiling and write out data.
/// This is called by the main `finalize_profiling` function.
pub fn finalize_memory_profiling() {
    write_final_memory_profile_data();
    // write_memory_dealloc_data();
    flush_debug_log();
}

/// Write final memory profile data to a file for completeness
fn write_final_memory_profile_data() {
    use std::{collections::HashMap, fs::File, path::Path};

    safe_alloc! {
        // Retrieve registries to get task allocations and names
        let memory_path = get_memory_path().unwrap_or("memory.folded");

        // Check if the file exists first
        let file_exists = Path::new(memory_path).exists();

        // If the file already exists, write the summary information to the existing file
        // Otherwise, create a new file with the appropriate headers
        let file_result = if file_exists {
            debug_log!("Opening existing file in append mode");
            File::options().append(true).open(memory_path)
        } else {
            debug_log!("Creating new file");
            match File::create(memory_path) {
                Ok(file) => {

                    Ok(file)
                }
                Err(e) => {
                    debug_log!("Error creating file: {e}");
                    Err(e)
                }
            }
        };

        if let Ok(file) = file_result {
            let mut writer = io::BufWriter::new(file);

            // Get the task path registry mapping for easier lookup
            let task_paths_map: HashMap<usize, Vec<String>> = {
                let binding = TASK_PATH_REGISTRY.lock();

                // Dump all entries for debugging
                // for (id, path) in binding.iter() {
                //     debug_log!("Registry entry: task {id}: path: {:?}", path);
                // }

                // Get all entries from the registry
                binding
                    .iter()
                    .map(|(task_id, path)| (*task_id, path.clone()))
                    .collect()
            };

            let mut already_written = HashSet::new();

            // Now write all tasks from registry that might not have allocations
            // This helps with keeping the full call hierarchy in the output
            for (task_id, path) in &task_paths_map {
                let task_id = *task_id;

                // let path_str = path.join(";");
                let path_str = build_stack(path, None, ";");
                if already_written.contains(&path_str) {
                    continue;
                }

                debug_log!("Writing for task {task_id} from registry: '{path_str}' with 0 bytes");

                // Write line with zero bytes to maintain call hierarchy
                write_alloc(task_id, 0, &mut writer, &mut already_written, &path_str);
            }

            // Make sure to flush the writer
            if let Err(e) = writer.flush() {
                debug_log!("Error flushing writer: {e}");
            }
        }
    };
}

fn write_alloc(
    task_id: usize,
    allocation: usize,
    writer: &mut io::BufWriter<std::fs::File>,
    already_written: &mut HashSet<String>,
    path_str: &str,
) {
    match writeln!(writer, "{} {}", path_str, allocation) {
        Ok(()) => {
            already_written.insert(path_str.to_string());
        }
        Err(e) => {
            debug_log!("Error writing line for task {task_id}: {e}");
        }
    }
}