use core::{
    alloc::{GlobalAlloc, Layout},
    ffi::c_void,
    ptr::{self, null_mut, NonNull},
};

use libyalloc::{allocator::BucketAllocator, sys::PageProvider};
use yggdrasil_rt::{
    mem::{MappingFlags, MappingSource},
    sync::mutex::Mutex,
    sys as syscall,
};

use crate::{
    error::{EResult, OptionExt},
    headers::{
        errno::Errno,
        string::mem::{memcpy, memset},
    },
};

struct Allocator;
struct PageProviderImpl;

unsafe impl GlobalAlloc for Allocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let pointer = YALLOC.lock().allocate(layout);
        match pointer {
            Some(ptr) => ptr.as_ptr(),
            None => null_mut(),
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let pointer = NonNull::new(ptr).expect("NULL pointer passed to dealloc()");
        YALLOC.lock().free(pointer, layout);
    }
}

impl PageProvider for PageProviderImpl {
    fn map_pages(count: usize) -> Option<NonNull<u8>> {
        match unsafe {
            syscall::map_memory(
                None,
                count * 0x1000,
                MappingFlags::WRITE,
                &MappingSource::Anonymous,
            )
        } {
            Ok(address) => {
                let pointer = ptr::with_exposed_provenance_mut(address);
                NonNull::new(pointer)
            }
            Err(_error) => None,
        }
    }

    fn unmap_pages(address: NonNull<u8>, count: usize) {
        let address = address.as_ptr().addr();
        unsafe { syscall::unmap_memory(address, count * 0x1000) }
            .expect("Could not unmap memory pages");
    }
}

unsafe fn get_allocation(ptr: NonNull<u8>) -> (NonNull<u8>, Layout) {
    assert!(usize::from(ptr.addr()) >= 16);
    let real_ptr = ptr.sub(16);
    let size = real_ptr.cast::<usize>().read();
    let align = real_ptr.cast::<usize>().add(1).read();
    assert_eq!(align, 16);
    let layout = Layout::from_size_align(size, align).unwrap();
    (real_ptr, layout)
}

pub fn c_alloc(size: usize, align: usize, zero: bool) -> EResult<NonNull<c_void>> {
    // allocation layout:
    // [ header ] [ ... ]
    // header: 16 bytes
    let align = match align {
        0..=16 => 16,
        _ => panic!("TODO: handle larger aligns: align={}", align),
    };
    let size = size + 16;
    let layout = Layout::from_size_align(size, align).expect("Couldn't setup malloc() layout");
    let ptr = YALLOC.lock().allocate(layout).e_ok_or(Errno::ENOMEM)?;

    if zero {
        unsafe {
            memset(ptr.cast().as_ptr(), 0, size);
        }
    }

    unsafe {
        ptr.cast::<usize>().write(size);
        ptr.cast::<usize>().add(1).write(align);
    }

    EResult::Ok(unsafe { ptr.add(16).cast() })
}

pub unsafe fn c_realloc(old_ptr: Option<NonNull<c_void>>, size: usize) -> EResult<NonNull<c_void>> {
    match old_ptr {
        Some(old_ptr) => {
            let (real_old_ptr, old_layout) = get_allocation(old_ptr.cast());
            assert!(old_layout.size() >= 16);
            let old_size = old_layout.size() - 16;

            if size == old_size {
                return EResult::Ok(old_ptr);
            }

            let new_ptr = c_alloc(size, old_layout.align(), false)?;

            // Copy no more memory than was allocated and will be allocated
            let memcpy_size = core::cmp::min(size, old_size);
            memcpy(new_ptr.as_ptr(), old_ptr.as_ptr(), memcpy_size);

            YALLOC.lock().free(real_old_ptr, old_layout);

            EResult::Ok(new_ptr)
        }
        None => c_alloc(size, 16, false),
    }
}

pub unsafe fn c_free(ptr: NonNull<c_void>) {
    let (real_ptr, layout) = get_allocation(ptr.cast());
    YALLOC.lock().free(real_ptr, layout);
}

#[global_allocator]
static ALLOCATOR: Allocator = Allocator;

static YALLOC: Mutex<BucketAllocator<PageProviderImpl>> = Mutex::new(BucketAllocator::new());
