//! x86-64-specific memory translation table management interfaces and functions
use core::{
    marker::PhantomData,
    ops::{Index, IndexMut, Range},
};

use bitflags::bitflags;
use libk_mm_interface::{
    address::{AsPhysicalAddress, PhysicalAddress},
    pointer::{PhysicalRef, PhysicalRefMut},
    process::PageAttributeUpdate,
    table::{
        EntryLevel, EntryLevelDrop, MapAttributes, NextPageTable, NonTerminalEntryLevel,
        TableAllocator,
    },
};
use yggdrasil_abi::error::Error;

use crate::KernelTableManagerImpl;

bitflags! {
    /// Describes how each page table entry is mapped
    pub struct PageAttributes: u64 {
        /// When set, the mapping is considered valid and pointing somewhere
        const PRESENT = 1 << 0;
        /// For tables, allows writes to further translation levels, for pages/blocks, allows
        /// writes to the region covered by the entry
        const WRITABLE = 1 << 1;
        /// When set for L2 entries, the mapping specifies a 2MiB page instead of a page table
        /// reference
        const BLOCK = 1 << 7;
        /// For tables, allows user access to further translation levels, for pages/blocks, allows
        /// user access to the region covered by the entry
        const USER = 1 << 2;
        /// If set, the page has been written to
        const DIRTY = 1 << 6;
    }
}

/// Represents a single virtual address space mapping depending on its translation level
#[derive(Clone, Copy, Debug)]
#[repr(transparent)]
pub struct PageEntry<L: EntryLevel>(u64, PhantomData<L>);

/// Table describing a single level of address translation
#[derive(Clone, Copy)]
#[repr(C, align(0x1000))]
pub struct PageTable<L: EntryLevel> {
    data: [PageEntry<L>; 512],
}

/// Translation level 0 (PML4): Entry is 512GiB table
#[derive(Clone, Copy, Debug)]
pub struct L0;
/// Translation level 1 (PDPT): Entry is 1GiB table
#[derive(Clone, Copy, Debug)]
pub struct L1;
/// Translation level 2 (Page directory): Entry is 2MiB block/table
#[derive(Clone, Copy, Debug)]
pub struct L2;
/// Translation level 3 (Page table): Entry is 4KiB page
#[derive(Clone, Copy, Debug)]
pub struct L3;

impl NonTerminalEntryLevel for L0 {
    type NextLevel = L1;
}
impl NonTerminalEntryLevel for L1 {
    type NextLevel = L2;
}
impl NonTerminalEntryLevel for L2 {
    type NextLevel = L3;
}

impl EntryLevel for L0 {
    const SHIFT: usize = 39;
}

impl EntryLevel for L1 {
    const SHIFT: usize = 30;
}

impl EntryLevel for L2 {
    const SHIFT: usize = 21;
}

impl EntryLevel for L3 {
    const SHIFT: usize = 12;
}

impl PageEntry<L3> {
    /// Constructs a mapping which points to a 4KiB page
    pub fn page(phys: PhysicalAddress, attrs: PageAttributes) -> Self {
        Self(
            u64::from(phys) | (attrs | PageAttributes::PRESENT | PageAttributes::USER).bits(),
            PhantomData,
        )
    }

    /// Returns the physical address of the page this entry refers to, returning None if it does
    /// not
    pub fn as_page(self) -> Option<PhysicalAddress> {
        if self.0 & PageAttributes::PRESENT.bits() != 0 {
            Some(PhysicalAddress::from_u64(self.0 & !Self::ATTR_MASK))
        } else {
            None
        }
    }

    pub fn is_dirty(&self) -> bool {
        self.0 & PageAttributes::DIRTY.bits() != 0
    }
}

impl PageEntry<L2> {
    /// Constructs a mapping which points to a 2MiB block
    pub fn block(phys: PhysicalAddress, attrs: PageAttributes) -> Self {
        Self(
            u64::from(phys) | (attrs | PageAttributes::PRESENT | PageAttributes::BLOCK).bits(),
            PhantomData,
        )
    }
}

impl PageEntry<L1> {
    /// Constructs a mapping which points to a 1GiB block
    pub fn block(phys: PhysicalAddress, attrs: PageAttributes) -> Self {
        Self(
            u64::from(phys) | (attrs | PageAttributes::PRESENT | PageAttributes::BLOCK).bits(),
            PhantomData,
        )
    }
}

impl<L: NonTerminalEntryLevel> PageEntry<L> {
    /// Constructs a mapping which points to a next-level table
    pub fn table(phys: PhysicalAddress, attrs: PageAttributes) -> Self {
        Self(
            u64::from(phys)
                | (attrs
                    | PageAttributes::PRESENT
                    | PageAttributes::WRITABLE
                    | PageAttributes::USER)
                    .bits(),
            PhantomData,
        )
    }

    /// Returns the physical address of the table this entry refers to, returning None if it
    /// does not
    pub fn as_table(self) -> Option<PhysicalAddress> {
        if self.0 & PageAttributes::PRESENT.bits() != 0
            && self.0 & PageAttributes::BLOCK.bits() == 0
        {
            Some(PhysicalAddress::from_u64(self.0 & !Self::ATTR_MASK))
        } else {
            None
        }
    }

    /// Returns `true` if the mapping represents a "page"/"block" and not a table
    pub fn is_block(self) -> bool {
        self.0 & PageAttributes::BLOCK.bits() != 0
    }
}

impl<L: EntryLevel> PageEntry<L> {
    const ATTR_MASK: u64 = 0xFFF | (1 << 63);

    /// An entry that is not mapped
    pub const INVALID: Self = Self(0, PhantomData);

    /// Reinterprets raw [u64] as a [PageEntry].
    ///
    /// # Safety
    ///
    /// Unsafe: the caller must ensure the value is a valid page translation entry.
    pub const unsafe fn from_raw(raw: u64) -> Self {
        Self(raw, PhantomData)
    }

    /// Returns the translation attributes of the entry
    pub fn attributes(&self) -> PageAttributes {
        PageAttributes::from_bits_retain(self.0)
    }

    /// Returns `true` if the entry contains a valid mapping to either a table or to a page/block
    pub fn is_present(&self) -> bool {
        self.0 & PageAttributes::PRESENT.bits() != 0
    }

    pub fn update(&mut self, update: &PageAttributeUpdate) -> Result<(), Error> {
        let mut attrs = PageAttributes::from_bits_retain(self.0);
        if let Some(write) = update.user_write {
            if write {
                attrs |= PageAttributes::WRITABLE;
            } else {
                attrs &= !PageAttributes::WRITABLE;
            }
        }
        // Dirty is ignored, it's hardware-managed
        self.0 &= !Self::ATTR_MASK;
        self.0 |= attrs.bits() & Self::ATTR_MASK;
        Ok(())
    }
}

impl<L: EntryLevel> PageTable<L> {
    /// Constructs a page table filled with invalid (non-present) entries
    pub const fn zeroed() -> Self {
        Self {
            data: [PageEntry::INVALID; 512],
        }
    }

    /// Reinterprets given [PageEntry] slice as a reference to [PageTable].
    ///
    /// # Safety
    ///
    /// Unsafe: the caller must ensure the provided reference is properly aligned and contains sane
    /// data.
    pub unsafe fn from_raw_slice_mut(data: &mut [PageEntry<L>; 512]) -> &mut Self {
        core::mem::transmute(data)
    }

    /// Allocates a new page table, filling it with non-preset entries
    pub fn new_zeroed<'a, TA: TableAllocator>(
    ) -> Result<PhysicalRefMut<'a, Self, KernelTableManagerImpl>, Error> {
        let physical = TA::allocate_page_table()?;
        let mut table =
            unsafe { PhysicalRefMut::<'a, Self, KernelTableManagerImpl>::map(physical) };

        for i in 0..512 {
            table[i] = PageEntry::INVALID;
        }

        Ok(table)
    }

    /// Recursively clears and deallocates the translation table.
    ///
    /// # Safety
    ///
    /// The caller must ensure the table is no longer in use and is not referenced anymore.
    pub unsafe fn free<TA: TableAllocator>(this: PhysicalRefMut<Self, KernelTableManagerImpl>) {
        let physical = this.as_physical_address();
        TA::free_page_table(physical);
    }

    // /// Returns the physical address of this table
    // pub fn physical_address(&self) -> usize {
    //     unsafe { (self.data.as_ptr() as usize).physicalize() }
    // }
}

impl<L: NonTerminalEntryLevel + 'static> NextPageTable for PageTable<L> {
    type NextLevel = PageTable<L::NextLevel>;
    type TableRef = PhysicalRef<'static, Self::NextLevel, KernelTableManagerImpl>;
    type TableRefMut = PhysicalRefMut<'static, Self::NextLevel, KernelTableManagerImpl>;

    fn get(&self, index: usize) -> Option<Self::TableRef> {
        self[index]
            .as_table()
            .map(|addr| unsafe { PhysicalRef::map(addr) })
    }

    fn get_mut(&mut self, index: usize) -> Option<Self::TableRefMut> {
        self[index]
            .as_table()
            .map(|addr| unsafe { PhysicalRefMut::map(addr) })
    }

    fn get_mut_or_alloc<TA: TableAllocator>(
        &mut self,
        index: usize,
    ) -> Result<Self::TableRefMut, Error> {
        let entry = self[index];

        if let Some(table) = entry.as_table() {
            Ok(unsafe { PhysicalRefMut::map(table) })
        } else {
            let table = PageTable::new_zeroed::<TA>()?;
            self[index] = PageEntry::<L>::table(
                unsafe { table.as_physical_address() },
                PageAttributes::WRITABLE | PageAttributes::USER,
            );
            Ok(table)
        }
    }
}

impl EntryLevelDrop for PageTable<L3> {
    const FULL_RANGE: Range<usize> = 0..512;

    // Do nothing
    unsafe fn drop_range<TA: TableAllocator>(&mut self, _range: Range<usize>) {}
}

impl<L: NonTerminalEntryLevel + 'static> EntryLevelDrop for PageTable<L>
where
    PageTable<L::NextLevel>: EntryLevelDrop,
{
    const FULL_RANGE: Range<usize> = 0..512;

    unsafe fn drop_range<TA: TableAllocator>(&mut self, range: Range<usize>) {
        for index in range {
            let entry = self[index];

            if let Some(table) = entry.as_table() {
                let mut table_ref: PhysicalRefMut<PageTable<L::NextLevel>, KernelTableManagerImpl> =
                    PhysicalRefMut::map(table);

                table_ref.drop_all::<TA>();

                TA::free_page_table(table);
            } else if entry.is_present() {
                // Memory must've been cleared beforehand, so no non-table entries must be present
                panic!(
                    "Expected a table containing only tables, got table[{}] = {:#x?}",
                    index, entry.0
                );
            }

            self[index] = PageEntry::INVALID;
        }
    }
}

impl<L: EntryLevel> Index<usize> for PageTable<L> {
    type Output = PageEntry<L>;

    fn index(&self, index: usize) -> &Self::Output {
        &self.data[index]
    }
}

impl<L: EntryLevel> IndexMut<usize> for PageTable<L> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        &mut self.data[index]
    }
}

impl From<MapAttributes> for PageAttributes {
    fn from(value: MapAttributes) -> Self {
        let mut res = PageAttributes::WRITABLE;
        if value.intersects(MapAttributes::USER_READ | MapAttributes::USER_WRITE) {
            res |= PageAttributes::USER;
        }
        res
    }
}

impl From<PageAttributes> for MapAttributes {
    fn from(value: PageAttributes) -> Self {
        let mut res = MapAttributes::empty();
        if value.contains(PageAttributes::USER) {
            res |= MapAttributes::USER_READ;
            if value.contains(PageAttributes::WRITABLE) {
                res |= MapAttributes::USER_WRITE;
            }
        }
        // TODO ???
        res |= MapAttributes::NON_GLOBAL;
        res
    }
}
