/// Memory allocator stuff
/// Compared to the original source, this will be some what "lean n' mean"
use std::cell::RefCell;
use std::collections::HashMap;

#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
/// The PurgeTags (PUs)
/// Tags < 100 are not overwritten untill freed
pub enum PurgeTag {
    /// Static the entire execution time
    Static = 1,
    /// Static while playing
    Sound = 2,
    /// static while playing
    Music = 3,
    /// Original source comment: anything else dave wants static
    Dave = 4,
    /// static until level exit
    Level = 50,
    /// static for special thinkers in a level
    LevSpec = 51,
    // Tags >= 100 are "purgable" whenever needed
    PurgeLevel = 100,
    Cache = 101,
}


/// Provides metadata on 
struct AllocMeta {
    /// The size of the allocation
    size: usize,
    /// The PU tag associated with the allocation
    tag: i32,
    /// A pointer to the caller's tracking pointer, cleared automatically upon freeing
    user: *mut *mut u8,
}


thread_local! {
    // We promise not to be naughty and access this from another thread :)
    /// Mapping from pointers (as usize) to the correspond AllocationMeta
    static REGISTRY: RefCell<HashMap<usize, AllocMeta>> = RefCell::new(HashMap::new());
}

// Helper method that we can use to access the registry
fn with_registry<F, R>(f: F) -> R 
where
    F: FnOnce(&mut HashMap<usize, AllocMeta>) -> R
{
    REGISTRY.with(|registry| f(&mut *registry.borrow_mut()))
}

#[allow(non_snake_case)]
/// Initalized the Zone memory allocator sub-system
pub fn Z_Init() {
    // do nothing, we get lazy init of the registry.
}

#[allow(non_snake_case)]
/// Allocate a contiguous chunk of zeroed memory from the zone allocator
/// # Safety
/// This function returns a raw pointer that bypasses Rust's automatic memory 
/// tracking. The returned memory must be manually reclaimed using `Z_Free`.
pub unsafe fn Z_Malloc(size: usize, tag: i32, user: *mut *mut u8) -> *mut u8 {
    let mut buffer = vec![0u8; size];
    let ptr = buffer.as_mut_ptr();

    // ensures Rust does not drop the buffer at the end of this block
    std::mem::forget(buffer);

    // Update the engine's tracking pointer if one was provided
    if !user.is_null() {
        unsafe {
            *user = ptr;
        }
    }

    // Save allocation details using the pointer address as the key
    with_registry(|map| {
        map.insert(ptr as usize, AllocMeta { size, tag, user });
    });

    ptr
}

#[allow(non_snake_case)]
/// Reclaims and frees a allocation from Z_Malloc
/// # Safety
/// The provided pointer must point to a valid allocation tracked by the registry.
/// Passing an invalid or already freed pointer can result in undefined behavior
pub unsafe fn Z_Free(ptr: *mut u8) {
    if ptr.is_null() {
        return;
    }
    let removed = with_registry(|map| map.remove(&(ptr as usize)));
    
    if let Some(meta) = removed {
        // Clear the caller's tracking pointer to prevent use-after-free
        if !meta.user.is_null() {
            unsafe {
                *meta.user = std::ptr::null_mut();
            }
        }
        // re-build the vector so Rust can drop it
        unsafe {
            let _dropped_vec = Vec::from_raw_parts(ptr, meta.size, meta.size);
        }
    }
}

#[allow(non_snake_case)]
/// Frees all memory allocations between the specific lifetime PU tags.
/// # Safety
/// All pointers tracked within this tag boundary are invalidated and deallocate
pub unsafe fn Z_FreeTags(low_tag: i32, high_tag: i32) {

    // get all of the allocations to free
    let tags_to_free: Vec<usize> = with_registry(|map| {
        map.iter()
            .filter(|(_, meta)| meta.tag >= low_tag && meta.tag <= high_tag)
            .map(|(ptr_addr,_)| *ptr_addr)
            .collect()
    });

    // remove each pointer returned from the map and drop the memory

    for ptr_addr in tags_to_free {
        let removed_ptr = with_registry(|map| map.remove(&ptr_addr));

        if let Some(meta) = removed_ptr {
            let p = ptr_addr as *mut u8;
            unsafe {
                // clear caller tracking pointer to preven use after free
                if !meta.user.is_null() {
                    *meta.user = std::ptr::null_mut();
                }

                let _dropped_vec = Vec::from_raw_parts(p, meta.size, meta.size);
            }
        }
    }
}

// Internal ChangeTag function
#[allow(non_snake_case)]
fn Z_ChangeTagInternal(ptr: *mut u8, tag: i32) {
    if ptr.is_null() {
        return;
    }

    with_registry(|map| {
        if let Some(meta) = map.get_mut(&(ptr as usize)) {
            meta.tag = tag;
        }
    })
}
/// Modifies the lifetime tag of an active allocation.
/// 
/// # Safety
/// The provided pointer must point to a valid allocation tracked by the registry
#[allow(non_snake_case)]
pub unsafe fn Z_ChangeTag(ptr: *mut u8, tag: i32) {
    // removes the original helper macro from the original C, we don't need withcraft here
    Z_ChangeTagInternal(ptr, tag);
}
