use core::{
	alloc::{GlobalAlloc, Layout},
	sync::atomic::{AtomicUsize, Ordering::Relaxed},
};

pub struct TalcAllocator(
	talc::cell::TalcSyncCell<talc::wasm::WasmGrowAndExtend, talc::wasm::WasmBinning>,
);

impl TalcAllocator {
	pub const fn new() -> Self {
		Self(talc::cell::TalcSyncCell::new_wasm(
			talc::wasm::WasmGrowAndExtend::new(),
		))
	}
}

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

unsafe impl GlobalAlloc for TalcAllocator {
	unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
		unsafe { self.0.alloc(layout) }
	}

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

	unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
		unsafe { self.0.alloc_zeroed(layout) }
	}

	unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
		unsafe { self.0.realloc(ptr, layout, new_size) }
	}
}

const CHUNK_SIZE: usize = 64;
const PAGE_SIZE: usize = 64 * 1024;

unsafe extern "C" {
	safe static __heap_base: usize;
}

pub struct BumpAllocator {
	offset: AtomicUsize,
	capacity: AtomicUsize,
	free_8: AtomicUsize,
	free_16: AtomicUsize,
	free_32: AtomicUsize,
	free_64: AtomicUsize,
	free_128: AtomicUsize,
	free_256: AtomicUsize,
	free_512: AtomicUsize,
	free_1024: AtomicUsize,
	free_2048: AtomicUsize,
	free_4096: AtomicUsize,
	free_8192: AtomicUsize,
	free_16384: AtomicUsize,
	free_32768: AtomicUsize,
}

fn space_for(layout: Layout) -> usize {
	usize::div_ceil(layout.size(), CHUNK_SIZE) * CHUNK_SIZE
}

fn pop_free_list(list: &AtomicUsize) -> Option<*mut u8> {
	let offset = list.swap(0, Relaxed);
	if offset != 0 {
		// safety: offset is a valid pointer for these operations when non-zero
		unsafe {
			let ptr = offset as *mut u8;
			let next = *ptr.cast::<usize>();
			list.swap(next, Relaxed);
			Some(ptr)
		}
	} else {
		None
	}
}

unsafe fn push_free_list(list: &AtomicUsize, ptr: *mut u8) {
	// safety: all of the pointers here are valid
	unsafe {
		let new_offset = ptr.addr();
		let current_head = list.swap(new_offset, Relaxed);
		*ptr.cast::<usize>() = current_head;
	}
}

impl BumpAllocator {
	pub const fn new() -> Self {
		Self {
			offset: AtomicUsize::new(8),
			capacity: AtomicUsize::new(PAGE_SIZE - 8),
			free_8: AtomicUsize::new(0),
			free_16: AtomicUsize::new(0),
			free_32: AtomicUsize::new(0),
			free_64: AtomicUsize::new(0),
			free_128: AtomicUsize::new(0),
			free_256: AtomicUsize::new(0),
			free_512: AtomicUsize::new(0),
			free_1024: AtomicUsize::new(0),
			free_2048: AtomicUsize::new(0),
			free_4096: AtomicUsize::new(0),
			free_8192: AtomicUsize::new(0),
			free_16384: AtomicUsize::new(0),
			free_32768: AtomicUsize::new(0),
		}
	}
}

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

unsafe impl GlobalAlloc for BumpAllocator {
	unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
		self.offset.fetch_max(__heap_base, Relaxed);
		let space_needed = space_for(layout);
		let ptr = match space_needed {
			8 => pop_free_list(&self.free_8),
			16 => pop_free_list(&self.free_16),
			32 => pop_free_list(&self.free_32),
			64 => pop_free_list(&self.free_64),
			128 => pop_free_list(&self.free_128),
			256 => pop_free_list(&self.free_256),
			512 => pop_free_list(&self.free_512),
			1024 => pop_free_list(&self.free_1024),
			2048 => pop_free_list(&self.free_2048),
			4096 => pop_free_list(&self.free_4096),
			8192 => pop_free_list(&self.free_8192),
			16384 => pop_free_list(&self.free_16384),
			32768 => pop_free_list(&self.free_32768),
			_ => None,
		};
		if let Some(ptr) = ptr {
			return ptr;
		}

		let offset = self.offset.fetch_add(space_needed, Relaxed);
		let pointer = offset as *mut u8;
		let available_space = self.capacity.load(Relaxed) - offset;
		if available_space < layout.size() {
			let pages = layout.size().div_ceil(PAGE_SIZE);
			core::arch::wasm32::memory_grow(0, pages);
			self.capacity.fetch_add(pages * PAGE_SIZE, Relaxed);
		}

		pointer
	}

	unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
		let space_used = space_for(layout);
		let ptr_offset = ptr.addr();
		if self
			.offset
			.compare_exchange(ptr_offset + space_used, ptr_offset, Relaxed, Relaxed)
			.is_ok()
		{
			return;
		}

		// safety: ptr must be valid and on the heap
		unsafe {
			match space_used {
				8 => push_free_list(&self.free_8, ptr),
				16 => push_free_list(&self.free_16, ptr),
				32 => push_free_list(&self.free_32, ptr),
				64 => push_free_list(&self.free_64, ptr),
				128 => push_free_list(&self.free_128, ptr),
				256 => push_free_list(&self.free_256, ptr),
				512 => push_free_list(&self.free_512, ptr),
				1024 => push_free_list(&self.free_1024, ptr),
				2048 => push_free_list(&self.free_2048, ptr),
				4096 => push_free_list(&self.free_4096, ptr),
				8192 => push_free_list(&self.free_8192, ptr),
				16384 => push_free_list(&self.free_16384, ptr),
				32768 => push_free_list(&self.free_32768, ptr),
				_ => {}
			};
		}
	}

	unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
		let space_used = space_for(layout);
		if new_size <= space_used {
			return ptr;
		}

		unsafe {
			self.dealloc(ptr, layout);
			self.alloc(Layout::from_size_align(new_size, layout.align()).unwrap_unchecked())
		}
	}
}