use core::ops::{Deref, DerefMut, Range};

use bitflags::bitflags;
use yggdrasil_abi::error::Error;

use super::address::PhysicalAddress;

#[cfg(any(not(target_arch = "x86"), rust_analyzer))]
const TABLE_SIZE_MASK: usize = 0x1FF;
#[cfg(any(target_arch = "x86", rust_analyzer))]
const TABLE_SIZE_MASK: usize = 0x3FF;

/// Interface for a single level of address translation
pub trait EntryLevel: Copy {
    /// The right shift needed to obtain an index of an entry at this level from an address
    const SHIFT: usize;
    /// The size of a page at this entry level
    const SIZE: usize = 1 << Self::SHIFT;
}

pub trait TableAllocator {
    fn allocate_page_table() -> Result<PhysicalAddress, Error>;
    /// # Safety
    ///
    /// `address` must point to a real translation table allocated by this trait impl.
    unsafe fn free_page_table(address: PhysicalAddress);
}

// TODO EXECUTABLE
bitflags! {
    /// Describes how a page translation mapping should behave
    #[derive(Clone, Copy)]
    pub struct MapAttributes: u64 {
        /// The data mapped can be read by the user process
        const USER_READ = 1 << 0;
        /// The data mapped can be written to by the user process
        const USER_WRITE = 1 << 1;
        /// The mapping is not global across the address spaces
        const NON_GLOBAL = 1 << 2;
        /// The mapping is marked dirty by the OS
        const DIRTY = 1 << 3;
    }
}

pub const fn page_index<T: EntryLevel>(address: usize) -> usize {
    address >> T::SHIFT & TABLE_SIZE_MASK
}

pub const fn page_offset<T: EntryLevel>(address: usize) -> usize {
    address & (T::SIZE - 1)
}

pub const fn page_count<T: EntryLevel>(address: usize) -> usize {
    address.div_ceil(T::SIZE)
}

pub const fn page_align_up<T: EntryLevel>(address: usize) -> usize {
    (address + T::SIZE - 1) & !(T::SIZE - 1)
}

pub const fn page_align_down<T: EntryLevel>(address: usize) -> usize {
    address & !(T::SIZE - 1)
}

pub const fn is_page_aligned_for<T: EntryLevel>(address: usize) -> bool {
    page_offset::<T>(address) == 0
}

pub trait EntryLevelExt: Sized {
    fn page_index<T: EntryLevel>(&self) -> usize;
    fn page_offset<T: EntryLevel>(&self) -> usize;
    fn page_count<T: EntryLevel>(&self) -> usize;
    fn page_align_up<T: EntryLevel>(&self) -> Self;
    fn page_align_down<T: EntryLevel>(&self) -> Self;
    fn is_page_aligned_for<T: EntryLevel>(&self) -> bool;
}

trait AddressLike: Sized + Copy {
    fn into_usize(self) -> usize;
    fn from_usize(v: usize) -> Self;
}

/// Interface for destroying memory translation tables
pub trait EntryLevelDrop {
    /// Range covering the whole table
    const FULL_RANGE: Range<usize>;

    /// Recursively destroys the specified range within the table.
    ///
    /// # Safety
    ///
    /// Caller must ensure the unmapped range is not in use by some other thread and will not be
    /// referred to from this point.
    unsafe fn drop_range<TA: TableAllocator>(&mut self, range: Range<usize>);

    /// Recursively destroys all the entries in within the table
    ///
    /// # Safety
    ///
    /// Caller must ensure the unmapped range is not in use by some other thread and will not be
    /// referred to from this point.
    unsafe fn drop_all<TA: TableAllocator>(&mut self) {
        self.drop_range::<TA>(Self::FULL_RANGE)
    }
}

/// Interface for non-terminal tables to retrieve the next level of address translation tables
pub trait NextPageTable {
    /// Type for the next-level page table
    type NextLevel;
    /// Type for an immutable reference to the next-level page table
    type TableRef: Deref<Target = Self::NextLevel>;
    /// Type for a mutable reference to the next-level page table
    type TableRefMut: DerefMut<Target = Self::NextLevel>;

    /// Tries looking up a next-level table at given index, allocating and mapping one if it is not
    /// present there
    fn get_mut_or_alloc<TA: TableAllocator>(
        &mut self,
        index: usize,
    ) -> Result<Self::TableRefMut, Error>;
    /// Returns a mutable reference to a next-level table at `index`, if present
    fn get_mut(&mut self, index: usize) -> Option<Self::TableRefMut>;
    /// Returns an immutable reference to a next-level table at `index`, if present
    fn get(&self, index: usize) -> Option<Self::TableRef>;
}

/// Tag trait to mark that the page table level may point to a next-level table
pub trait NonTerminalEntryLevel: EntryLevel {
    /// Tag type of the level this entry level may point to
    type NextLevel: EntryLevel;
}

impl AddressLike for usize {
    #[inline(always)]
    fn into_usize(self) -> usize {
        self
    }

    #[inline(always)]
    fn from_usize(v: usize) -> Self {
        v
    }
}

impl AddressLike for PhysicalAddress {
    fn from_usize(v: usize) -> Self {
        Self(v as _)
    }

    fn into_usize(self) -> usize {
        self.0 as _
    }
}

impl<T: AddressLike> EntryLevelExt for T {
    #[inline(always)]
    fn page_index<L: EntryLevel>(&self) -> usize {
        page_index::<L>(self.into_usize())
    }

    #[inline(always)]
    fn page_offset<L: EntryLevel>(&self) -> usize {
        page_offset::<L>(self.into_usize())
    }

    #[inline(always)]
    fn page_count<L: EntryLevel>(&self) -> usize {
        page_count::<L>(self.into_usize())
    }

    #[inline(always)]
    fn page_align_up<L: EntryLevel>(&self) -> Self {
        Self::from_usize(page_align_up::<L>(self.into_usize()))
    }

    #[inline(always)]
    fn page_align_down<L: EntryLevel>(&self) -> Self {
        Self::from_usize(page_align_down::<L>(self.into_usize()))
    }

    #[inline(always)]
    fn is_page_aligned_for<L: EntryLevel>(&self) -> bool {
        self.page_offset::<L>() == 0
    }
}
