use crate::script_implementations::wasm::script_function_cache::WasmPtr;
use std::mem::size_of;
use std::sync::atomic::{AtomicUsize, Ordering};

/// A TempWasmAllocator is a bump allocator for use with very short lived allocations within WASM.
/// This is for example used to allocate the values by reference in our script functions. These
/// live only during the call, and then get deallocated again. Once everything is deallocated, we
/// reset the allocator to the start again.
pub(super) struct TempWasmAllocator {
    /// The raw pointer in host memory where the allocation lives.
    data: *mut u8,
    /// The pointer in client memory where the allocation lives.
    wasm_pointer: u32,
    /// The total amount we've currently allocated. This fluctuates up and down, when allocating
    /// and de-allocating. When this hits 0, the allocation gets reset.
    currently_allocated: AtomicUsize,
    /// The bump position for our allocations. we increment this when we allocate, but don't decrement
    /// when we de-allocate. We reset this to 0 when currently_allocated hits 0.
    offset_high: AtomicUsize,
}

impl TempWasmAllocator {
    /// Creates a new allocator, with a given pointer to memory, and the associated position
    /// within WASM memory.
    pub(super) fn new(data: *mut u8, wasm_pointer: u32) -> Self {
        Self {
            data,
            wasm_pointer,
            currently_allocated: AtomicUsize::new(0),
            offset_high: AtomicUsize::new(0),
        }
    }

    /// Allocates a new object with a certain type.
    pub fn alloc<T>(&self, value: T) -> AllocatedObject<T> {
        self.currently_allocated.fetch_add(size_of::<T>(), Ordering::SeqCst);
        let ptr_offset = self.offset_high.fetch_add(size_of::<T>(), Ordering::SeqCst);
        let ptr = unsafe { self.data.add(ptr_offset) } as *mut T;
        unsafe {
            *ptr = value;
        }
        AllocatedObject::<T> {
            ptr,
            wasm_ptr: (self.wasm_pointer + ptr_offset as u32).into(),
            allocator: self as *const TempWasmAllocator,
        }
    }

    /// Drops an earlier allocated type.
    pub fn drop<T>(&self) {
        self.currently_allocated.fetch_sub(size_of::<T>(), Ordering::SeqCst);
        // As soon as we've no longer allocated anything, we reset our allocating back to the start.
        if self.currently_allocated.load(Ordering::SeqCst) == 0 {
            self.offset_high.store(0, Ordering::SeqCst);
        }
    }
}

/// A value allocated within WASM memory. Once this goes out of scope, it gets deallocated,
pub(super) struct AllocatedObject<T> {
    /// The pointer in host memory to where the object lives.
    pub ptr: *mut T,
    /// The pointer in client memory to where the object lives.
    pub wasm_ptr: WasmPtr<T>,
    /// A raw pointer to the allocator we use, so we know where to deallocate to.
    allocator: *const TempWasmAllocator,
}

impl<T> AllocatedObject<T> {
    /// Gets the currently underlying value from the pointer.
    pub fn value(&self) -> &T {
        unsafe { &*self.ptr }
    }
}

impl<T> Drop for AllocatedObject<T> {
    fn drop(&mut self) {
        unsafe {
            self.allocator.as_ref().unwrap().drop::<T>();
        }
    }
}