use core::{
    marker::PhantomData,
    ops::{Deref, DerefMut, Index, IndexMut},
    ptr::NonNull,
};

use alloc::{boxed::Box, slice, vec::Vec};

use crate::mem::MemoryRegion;

pub struct PhysicalBox<T: ?Sized> {
    data: NonNull<T>,
    #[allow(dead_code)]
    region: MemoryRegion,
    physical_address: usize,
    _marker: PhantomData<T>,
}

impl<T> PhysicalBox<T> {
    pub fn new(data: T) -> Self {
        let (memory_region, paddr) =
            MemoryRegion::contiguous_physical(size_of::<T>() as u64).expect("Failed to allocate");
        // UNWRAP: We know this isn't null.
        let ptr = NonNull::new(memory_region.mut_ptr_at_offset(0)).unwrap();
        unsafe { ptr.write(data) };
        Self {
            data: ptr,
            region: memory_region,
            physical_address: paddr as usize,
            _marker: PhantomData,
        }
    }
}

impl<T: ?Sized> PhysicalBox<T> {
    pub fn physical_address(&self) -> usize {
        self.physical_address
    }
}

impl<T: ?Sized> Deref for PhysicalBox<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        // SAFETY:
        //   - Alignment: This is page aligned.
        //   - Dereferenceable: Guaranteed in same allocation.
        //   - Aliasing: The borrow rules ensure this
        unsafe { self.data.as_ref() }
    }
}

impl<T: ?Sized> DerefMut for PhysicalBox<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        // SAFETY:
        //   - Alignment: This is page aligned.
        //   - Dereferenceable: Guaranteed in same allocation.
        //   - Aliasing: The borrow rules ensure this
        unsafe { self.data.as_mut() }
    }
}

impl<T> PhysicalBox<[T]> {
    pub fn default_with_count(default: T, len: usize) -> Self
    where
        T: Clone,
    {
        let layout = core::alloc::Layout::array::<T>(len).expect("Layout overflow");

        // TODO: Implement a function like alloc that takes a layout.
        let (memory_region, paddr) =
            MemoryRegion::contiguous_physical(layout.size() as u64).expect("Failed to allocate");

        let ptr: *mut T = memory_region.mut_ptr_at_offset(0);
        for i in 0..len {
            unsafe {
                ptr.add(i).write(default.clone());
            }
        }

        let slice_ptr = core::ptr::slice_from_raw_parts_mut(ptr, len);

        Self {
            // UNWRAP: We know this isn't null.
            data: NonNull::new(slice_ptr).unwrap(),
            region: memory_region,
            physical_address: paddr as usize,
            _marker: PhantomData,
        }
    }

    pub fn from_vec(mut vec: Vec<T>) -> Self {
        let len = vec.len();
        let layout = core::alloc::Layout::array::<T>(len).expect("Layout overflow");

        // TODO: Implement a function like alloc that takes a layout.
        let (memory_region, paddr) =
            MemoryRegion::contiguous_physical(layout.size() as u64).expect("Failed to allocate");

        let ptr: *mut T = memory_region.mut_ptr_at_offset(0);
        for (i, item) in vec.into_iter().enumerate() {
            unsafe {
                ptr.add(i).write(item);
            }
        }

        let slice_ptr = core::ptr::slice_from_raw_parts_mut(ptr, len);

        Self {
            // UNWRAP: We know this isn't null.
            data: NonNull::new(slice_ptr).unwrap(),
            region: memory_region,
            physical_address: paddr as usize,
            _marker: PhantomData,
        }
    }
    pub fn len(&self) -> usize {
        (**self).len()
    }
}

impl<I, T> Index<I> for PhysicalBox<[T]>
where
    I: slice::SliceIndex<[T]>,
{
    type Output = I::Output;

    fn index(&self, index: I) -> &Self::Output {
        &(**self)[index]
    }
}

impl<I, T> IndexMut<I> for PhysicalBox<[T]>
where
    I: slice::SliceIndex<[T]>,
{
    fn index_mut(&mut self, index: I) -> &mut Self::Output {
        &mut (**self)[index]
    }
}

/// SAFETY: We are the only owner of this pointer.
unsafe impl<T: ?Sized> Send for PhysicalBox<T> where Box<T>: Send {}

/// SAFETY: You must have a mutable reference to this
/// type to modify the data at the pointer.
unsafe impl<T: ?Sized> Sync for PhysicalBox<T> where Box<T>: Sync {}

impl<T: ?Sized> Drop for PhysicalBox<T> {
    fn drop(&mut self) {
        // SAFETY:
        //   - We own this data.
        unsafe { core::ptr::drop_in_place(self.data.as_ptr()) }
    }
}