// SPDX-License-Identifier: GPL-2.0

//! Kernel memory allocation (kmalloc)

use alloc::collections::BTreeMap;
use alloc::vec::Vec;
use core::ptr::NonNull;

use crate::error::{Error, Result};
use crate::memory::allocator::{alloc_pages, free_pages, GfpFlags, PageFrameNumber};
use crate::sync::Spinlock;

/// Kmalloc size classes (powers of 2)
const KMALLOC_SIZES: &[usize] = &[8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096];
const MAX_KMALLOC_SIZE: usize = 4096;

/// Slab allocator for small kernel allocations
/// Uses indices instead of raw pointers for thread safety
struct SlabAllocator {
	size_classes: BTreeMap<usize, Vec<usize>>, // Store offsets instead of pointers
	allocated_blocks: BTreeMap<usize, usize>,  // Maps offsets to size classes
	base_addr: usize,                          // Base address for calculations
}

impl SlabAllocator {
	const fn new() -> Self {
		Self {
			size_classes: BTreeMap::new(),
			allocated_blocks: BTreeMap::new(),
			base_addr: 0,
		}
	}

	fn init(&mut self, base_addr: usize) {
		self.base_addr = base_addr;
	}

	fn allocate(&mut self, size: usize) -> Result<*mut u8> {
		// Find appropriate size class
		let size_class = KMALLOC_SIZES
			.iter()
			.find(|&&s| s >= size)
			.copied()
			.unwrap_or(MAX_KMALLOC_SIZE);

		if size_class > MAX_KMALLOC_SIZE {
			return Err(Error::OutOfMemory);
		}

		// Try to get from free list
		if let Some(free_list) = self.size_classes.get_mut(&size_class) {
			if let Some(offset) = free_list.pop() {
				self.allocated_blocks.insert(offset, size_class);
				return Ok((self.base_addr + offset) as *mut u8);
			}
		}

		// Allocate new page and split it
		self.allocate_new_slab(size_class)
	}

	fn allocate_new_slab(&mut self, size_class: usize) -> Result<*mut u8> {
		// Allocate a page using buddy allocator
		let pfn = alloc_pages(0, GfpFlags::KERNEL)?;
		let page_addr = pfn.to_phys_addr().as_usize();
		let offset = page_addr - self.base_addr;

		// Split page into blocks of size_class
		let blocks_per_page = 4096 / size_class;
		let free_list = self.size_classes.entry(size_class).or_insert_with(Vec::new);

		for i in 1..blocks_per_page {
			let block_offset = offset + (i * size_class);
			free_list.push(block_offset);
		}

		// Return the first block
		self.allocated_blocks.insert(offset, size_class);
		Ok(page_addr as *mut u8)
	}

	fn deallocate(&mut self, ptr: *mut u8) -> Result<()> {
		let offset = (ptr as usize).saturating_sub(self.base_addr);
		if let Some(size_class) = self.allocated_blocks.remove(&offset) {
			let free_list =
				self.size_classes.entry(size_class).or_insert_with(Vec::new);
			free_list.push(offset);
			Ok(())
		} else {
			Err(Error::InvalidArgument)
		}
	}
}

static SLAB_ALLOCATOR: Spinlock<SlabAllocator> = Spinlock::new(SlabAllocator::new());

/// Allocate kernel memory
pub fn kmalloc(size: usize) -> Result<*mut u8> {
	// Increment performance counter
	crate::perf::counters::inc_memory_allocs();

	if size == 0 {
		return Err(Error::InvalidArgument);
	}

	if size <= MAX_KMALLOC_SIZE {
		// Use slab allocator for small allocations
		let mut allocator = SLAB_ALLOCATOR.lock();
		allocator.allocate(size)
	} else {
		// Use buddy allocator for large allocations
		let pages_needed = (size + 4095) / 4096;
		let order = pages_needed.next_power_of_two().trailing_zeros() as usize;
		let pfn = alloc_pages(order, GfpFlags::KERNEL)?;
		Ok(pfn.to_phys_addr().as_usize() as *mut u8)
	}
}

/// Free kernel memory
pub fn kfree(ptr: *mut u8) {
	if ptr.is_null() {
		return;
	}

	// Try slab allocator first
	if let Ok(()) = SLAB_ALLOCATOR.lock().deallocate(ptr) {
		return;
	}

	// Try buddy allocator for large allocations
	// TODO: Keep track of large allocations to know how many pages to free
	// For now, assume single page
	if let Some(_page) = NonNull::new(ptr as *mut crate::memory::Page) {
		let pfn =
			PageFrameNumber::from_phys_addr(crate::types::PhysAddr::new(ptr as usize));
		free_pages(pfn, 0);
	}
}

/// Allocate zeroed kernel memory
pub fn kzalloc(size: usize) -> Result<*mut u8> {
	let ptr = kmalloc(size)?;
	unsafe {
		core::ptr::write_bytes(ptr, 0, size);
	}
	Ok(ptr)
}

/// Reallocate kernel memory
pub fn krealloc(ptr: *mut u8, old_size: usize, new_size: usize) -> Result<*mut u8> {
	if ptr.is_null() {
		return kmalloc(new_size);
	}

	if new_size == 0 {
		kfree(ptr);
		return Ok(core::ptr::null_mut());
	}

	let new_ptr = kmalloc(new_size)?;
	let copy_size = core::cmp::min(old_size, new_size);

	unsafe {
		core::ptr::copy_nonoverlapping(ptr, new_ptr, copy_size);
	}

	kfree(ptr);
	Ok(new_ptr)
}

/// Initialize the slab allocator
pub fn init() -> Result<()> {
	let mut allocator = SLAB_ALLOCATOR.lock();
	// Use a reasonable base address for offset calculations
	allocator.init(0x_4000_0000_0000);
	Ok(())
}
