use core::{
    alloc::Layout,
    mem::size_of,
    ops::{Deref, DerefMut},
};

use alloc::sync::Arc;
use kernel_arch::KernelTableManagerImpl;
use libk_mm_interface::address::{AsPhysicalAddress, PhysicalAddress};
use yggdrasil_abi::error::Error;

pub use kernel_arch::mem::{DeviceMemoryAttributes, DeviceMemoryCaching};

pub type RawDeviceMemoryMapping = kernel_arch::mem::RawDeviceMemoryMapping<KernelTableManagerImpl>;

/// Describes a single untyped device memory mapping
#[derive(Clone, Debug)]
pub struct DeviceMemoryMapping {
    #[allow(unused)]
    inner: Arc<RawDeviceMemoryMapping>,
    address: usize,
}

/// Describes a single typed device memory mapping
#[derive(Clone, Debug)]
pub struct DeviceMemoryIo<'a, T: ?Sized> {
    #[allow(unused)]
    inner: Arc<RawDeviceMemoryMapping>,
    value: &'a T,
}

/// Describes a single typed and mutable device memory mapping
#[derive(Debug)]
pub struct DeviceMemoryIoMut<'a, T: ?Sized> {
    #[allow(unused)]
    inner: RawDeviceMemoryMapping,
    value: &'a mut T,
}

impl DeviceMemoryMapping {
    /// Maps a region of physical memory as device memory of given size.
    ///
    /// See [RawDeviceMemoryMapping::map].
    ///
    /// # Safety
    ///
    /// The caller must ensure proper access synchronization, as well as the address' origin.
    pub unsafe fn map(
        base: PhysicalAddress,
        size: usize,
        attrs: DeviceMemoryAttributes,
    ) -> Result<Self, Error> {
        let inner = RawDeviceMemoryMapping::map(base.into_u64(), size, attrs)?;
        let address = inner.address;
        Ok(Self {
            inner: Arc::new(inner),
            address,
        })
    }

    /// Returns the address to which the object is mapped
    pub fn address(&self) -> usize {
        self.address
    }
}

impl<'a, T: Sized> DeviceMemoryIo<'a, T> {
    /// Interprets a raw device memory mapping as pointing to a value of `T`.
    ///
    /// # Safety
    ///
    /// The caller must ensure the mapping actually contains the value of `T`, as well as proper
    /// access synchronization.
    pub unsafe fn from_raw(
        inner: Arc<RawDeviceMemoryMapping>,
    ) -> Result<DeviceMemoryIo<'a, T>, Error> {
        if size_of::<T>() > inner.page_size * inner.page_count {
            todo!();
        }
        // TODO check align
        let value = &*(inner.address as *const T);
        Ok(DeviceMemoryIo { inner, value })
    }

    /// Maps a physical address as device memory of type `[T]`.
    ///
    /// # Safety
    ///
    /// The caller must ensure the address actually points to a value of type `T`, as well as
    /// proper access synchronization. The caller must also ensure the `len` is valid.
    pub unsafe fn map_slice(
        base: PhysicalAddress,
        count: usize,
        attrs: DeviceMemoryAttributes,
    ) -> Result<DeviceMemoryIo<'a, [T]>, Error> {
        let layout = Layout::array::<T>(count).unwrap();
        let inner = RawDeviceMemoryMapping::map(base.into_u64(), layout.size(), attrs)?;
        let value = core::slice::from_raw_parts(inner.address as *mut T, count);

        Ok(DeviceMemoryIo {
            inner: Arc::new(inner),
            value,
        })
    }

    /// Maps a physical address as device memory of type `T`.
    ///
    /// # Safety
    ///
    /// The caller must ensure the address actually points to a value of type `T`, as well as
    /// proper access synchronization.
    pub unsafe fn map(
        base: PhysicalAddress,
        attrs: DeviceMemoryAttributes,
    ) -> Result<DeviceMemoryIo<'a, T>, Error> {
        let inner = RawDeviceMemoryMapping::map(base.into_u64(), size_of::<T>(), attrs)?;
        let value = &*(inner.address as *const T);

        Ok(DeviceMemoryIo {
            inner: Arc::new(inner),
            value,
        })
    }
}

impl<'a, T: ?Sized> DeviceMemoryIo<'a, T> {
    /// Extracts an inner reference to `U` from within `T`.
    ///
    /// # Safety
    ///
    /// To use this safely, the caller must guarantee the "extracted" reference will be the only
    /// valid reference and that `&U` won't be usable again through its parent `T`. The caller must
    /// also guarantee that `&U` is, in fact, contained within `T`.
    pub unsafe fn extract<'b, U: ?Sized + 'b, F: FnOnce(&'a T) -> *const U>(
        &'a self,
        f: F,
    ) -> DeviceMemoryIo<'b, U> {
        let value = f(self.value);

        DeviceMemoryIo {
            inner: self.inner.clone(),
            value: &*value,
        }
    }
}

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

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

impl<T: ?Sized> AsPhysicalAddress for DeviceMemoryIo<'_, T> {
    unsafe fn as_physical_address(&self) -> PhysicalAddress {
        PhysicalAddress::from_usize(self.inner.base_address)
    }
}

unsafe impl<T: ?Sized + Send> Send for DeviceMemoryIo<'_, T> {}
unsafe impl<T: ?Sized + Sync> Sync for DeviceMemoryIo<'_, T> {}
//impl<T: ?Sized> !Sync for DeviceMemoryIo<'_, T> {}

impl<'a, T: Sized> DeviceMemoryIoMut<'a, T> {
    /// Maps a physical address as device memory to a slice `[T; len]`
    ///
    /// # Safety
    ///
    /// The caller must ensure the address actually points to a value of type `T`, as well as
    /// proper access synchronization. The caller must also ensure the `len` is valid.
    pub unsafe fn map_slice(
        base: PhysicalAddress,
        len: usize,
        attrs: DeviceMemoryAttributes,
    ) -> Result<DeviceMemoryIoMut<'a, [T]>, Error> {
        let layout = Layout::array::<T>(len).unwrap();
        let inner = RawDeviceMemoryMapping::map(base.into_u64(), layout.size(), attrs)?;
        let value = core::slice::from_raw_parts_mut(inner.address as *mut T, len);

        Ok(DeviceMemoryIoMut { inner, value })
    }
}

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

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

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

impl<T: ?Sized> !Send for DeviceMemoryIoMut<'_, T> {}
impl<T: ?Sized> !Sync for DeviceMemoryIoMut<'_, T> {}
