use core::{
    fmt,
    marker::PhantomData,
    ops::{Deref, DerefMut},
};

use kernel_arch_interface::mem::KernelTableManager;

use crate::address::{AsPhysicalAddress, PhysicalAddress};

/// Wrapper for immutably accessing a value at a physical address
#[repr(transparent)]
pub struct PhysicalRef<'a, T: ?Sized, K: KernelTableManager> {
    value: &'a T,
    _pd: PhantomData<K>,
}

/// Wrapper for mutably accessing a value at a physical address
#[repr(transparent)]
pub struct PhysicalRefMut<'a, T: ?Sized, K: KernelTableManager> {
    value: &'a mut T,
    _pd: PhantomData<K>,
}

// PhysicalRefMut<T> wrapper for safe mutable access to physical addresses

impl<'a, T: Sized, K: KernelTableManager> PhysicalRefMut<'a, T, K> {
    /// Maps a physical address into the kernel space as &mut T, allowing mmutable access to it.
    ///
    /// # Safety
    ///
    /// The caller must ensure the correct origin of the physical address as well that it actually
    /// contains T. The caller must also take care of access synchronization and make sure no
    /// aliasing occurs.
    pub unsafe fn map(physical: PhysicalAddress) -> PhysicalRefMut<'a, T, K> {
        let value = virtualize_raw::<_, K>(physical);
        PhysicalRefMut {
            value,
            _pd: PhantomData,
        }
    }

    /// Maps a physical address into the kernel space as &mut [T], allowing mmutable access to it.
    ///
    /// # Safety
    ///
    /// The caller must ensure the correct origin of the physical address as well that it actually
    /// contains [T; len]. The caller must also take care of access synchronization and make
    /// sure no aliasing occurs.
    pub unsafe fn map_slice(physical: PhysicalAddress, len: usize) -> PhysicalRefMut<'a, [T], K> {
        let value = virtualize_slice_raw::<_, K>(physical, len);
        PhysicalRefMut {
            value,
            _pd: PhantomData,
        }
    }
}

impl<T: ?Sized, K: KernelTableManager> PhysicalRefMut<'_, T, K> {
    /// Returns the "address" part of the reference
    #[inline]
    pub fn as_address(&self) -> usize {
        (self.value as *const T).addr()
    }
}

impl<T: ?Sized, K: KernelTableManager> AsPhysicalAddress for PhysicalRefMut<'_, T, K> {
    unsafe fn as_physical_address(&self) -> PhysicalAddress {
        PhysicalAddress::raw_from_virtualized::<K>(PhysicalRefMut::<T, K>::as_address(self))
    }
}

impl<T: ?Sized, K: KernelTableManager> Deref for PhysicalRefMut<'_, T, K> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.value
    }
}

impl<T: ?Sized, K: KernelTableManager> DerefMut for PhysicalRefMut<'_, T, K> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.value
    }
}

impl<T: ?Sized, K: KernelTableManager> fmt::Pointer for PhysicalRefMut<'_, T, K> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Pointer::fmt(&self.value, f)
    }
}

impl<T: ?Sized, K: KernelTableManager> Drop for PhysicalRefMut<'_, T, K> {
    fn drop(&mut self) {
        unsafe { K::unmap_physical_address(self.as_address()) }
    }
}

// PhysicalRef<T>: same as PhysicalRefMut<T>, except immutable

impl<'a, T: Sized, K: KernelTableManager> PhysicalRef<'a, T, K> {
    /// Maps a physical address into the kernel space as &T, allowing immutable access to it.
    ///
    /// # Safety
    ///
    /// The caller must ensure the correct origin of the physical address as well that it actually
    /// contains T.
    pub unsafe fn map(physical: PhysicalAddress) -> PhysicalRef<'a, T, K> {
        let value = virtualize_raw::<_, K>(physical);
        PhysicalRef {
            value,
            _pd: PhantomData,
        }
    }

    /// Maps a physical address into the kernel space as &[T] of given len, allowing immutable
    /// access to it.
    ///
    /// # Safety
    ///
    /// The caller must ensure the correct origin of the physical address as well that it actually
    /// contains [T; len].
    pub unsafe fn map_slice(physical: PhysicalAddress, len: usize) -> PhysicalRef<'a, [T], K> {
        let value = virtualize_slice_raw::<_, K>(physical, len);
        PhysicalRef {
            value,
            _pd: PhantomData,
        }
    }
}

impl<T: ?Sized, K: KernelTableManager> PhysicalRef<'_, T, K> {
    /// Returns the "address" part of the reference
    #[inline]
    pub fn as_address(&self) -> usize {
        (self.value as *const T).addr()
    }
}

impl<T: ?Sized, K: KernelTableManager> AsPhysicalAddress for PhysicalRef<'_, T, K> {
    unsafe fn as_physical_address(&self) -> PhysicalAddress {
        PhysicalAddress::raw_from_virtualized::<K>(PhysicalRef::<T, K>::as_address(self))
    }
}

impl<T: ?Sized, K: KernelTableManager> Deref for PhysicalRef<'_, T, K> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.value
    }
}

impl<T: ?Sized, K: KernelTableManager> Drop for PhysicalRef<'_, T, K> {
    fn drop(&mut self) {
        unsafe { K::unmap_physical_address(self.as_address()) }
    }
}

unsafe fn virtualize_raw<'a, T: Sized, K: KernelTableManager>(
    physical: PhysicalAddress,
) -> &'a mut T {
    // TODO check align
    let address = physical.raw_virtualize::<K>();
    &mut *(address as *mut T)
}

unsafe fn virtualize_slice_raw<'a, T: Sized, K: KernelTableManager>(
    physical: PhysicalAddress,
    len: usize,
) -> &'a mut [T] {
    // TODO check align
    let address = physical.raw_virtualize::<K>();
    core::slice::from_raw_parts_mut(address as *mut T, len)
}
