use core::{
    fmt,
    marker::PhantomData,
    ops::{Index, IndexMut, Range},
};

use bitflags::bitflags;
use kernel_arch_interface::KERNEL_VIRT_OFFSET;
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;

use super::dc_cvac;

bitflags! {
    #[derive(Clone, Copy, PartialEq, Eq)]
    pub struct PageAttributes: u64 {
        const PRESENT = 1 << 0;

        const TABLE = 1 << 1;
        const PAGE = 1 << 1;
        const BLOCK = 0 << 1;

        const ACCESS = 1 << 10;

        const AP_KERNEL_READWRITE = 0 << 6;
        const AP_BOTH_READWRITE = 1 << 6;
        const AP_KERNEL_READONLY = 2 << 6;
        const AP_BOTH_READONLY = 3 << 6;
        const AP_ACCESS_MASK = 3 << 6;

        const SH_OUTER = 2 << 8;
        const SH_INNER = 3 << 8;

        const PAGE_ATTR_NORMAL = 0 << 2;
        const PAGE_ATTR_NORMAL_NC = 1 << 2;
        const PAGE_ATTR_DEVICE = 2 << 2;

        const NON_GLOBAL = 1 << 11;

        const DIRTY = 1 << 51;
        const PXN = 1 << 53;
        const UXN = 1 << 54;
    }
}

#[derive(Clone, Copy)]
#[repr(C, align(0x1000))]
pub struct PageTable<L: EntryLevel> {
    entries: [PageEntry<L>; 512],
}

#[derive(Clone, Copy)]
pub struct PageEntry<L: EntryLevel>(u64, PhantomData<L>);

#[derive(Clone, Copy)]
pub struct L1;
#[derive(Clone, Copy)]
pub struct L2;
#[derive(Clone, Copy)]
pub struct L3;

#[derive(Debug, Clone, Copy)]
pub enum EntryType {
    Table(PhysicalAddress),
    Page(PhysicalAddress),
    Invalid,
}

impl NonTerminalEntryLevel for L1 {
    type NextLevel = L2;
}

impl NonTerminalEntryLevel for L2 {
    type NextLevel = L3;
}

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<L: EntryLevel> PageTable<L> {
    pub const fn zeroed() -> Self {
        Self {
            entries: [PageEntry::INVALID; 512],
        }
    }

    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)
    }

    /// Creates a reference to [PageTable] from a physical address.
    ///
    /// # Safety
    ///
    /// The function takes in a raw physical address.
    pub unsafe fn from_physical(
        physical: PhysicalAddress,
    ) -> Option<PhysicalRefMut<'static, Self, KernelTableManagerImpl>> {
        if physical.into_usize() >= KERNEL_VIRT_OFFSET {
            // Looks fishy
            return None;
        }
        if !physical.is_aligned_for::<L3>() {
            return None;
        }

        let inner = PhysicalRefMut::map(physical);
        Some(inner)
    }
}

impl<L: EntryLevel> PageEntry<L> {
    const ATTR_MASK: u64 = 0xFFF | (0xFFFF << 48);
    pub const INVALID: Self = Self(0, PhantomData);

    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 {
            // Make writeable/non-writeable
            if write {
                attrs &= !PageAttributes::AP_ACCESS_MASK;
                attrs |= PageAttributes::AP_BOTH_READWRITE;
            } else {
                todo!();
            }
        }
        if let Some(dirty) = update.dirty {
            if dirty {
                attrs |= PageAttributes::DIRTY;
            } else {
                attrs &= !PageAttributes::DIRTY;
            }
        }
        self.0 &= !Self::ATTR_MASK;
        self.0 |= attrs.bits() & Self::ATTR_MASK;
        Ok(())
    }

    pub const fn is_present(self) -> bool {
        self.0 & PageAttributes::PRESENT.bits() != 0
    }

    pub fn attributes(self) -> PageAttributes {
        PageAttributes::from_bits_retain(self.0)
    }
}

impl<L: NonTerminalEntryLevel> PageTable<L> {
    pub fn walk(&self, index: usize) -> EntryType {
        self[index].classify()
    }
}

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

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

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

    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::empty(),
            );
            dc_cvac((&raw const self[index]).addr());
            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;
            dc_cvac((&raw const self[index]).addr());
        }
    }
}

impl<L: NonTerminalEntryLevel> PageEntry<L> {
    pub fn table(phys: PhysicalAddress, attrs: PageAttributes) -> Self {
        Self(
            phys.into_u64() | (PageAttributes::TABLE | PageAttributes::PRESENT | attrs).bits(),
            PhantomData,
        )
    }

    pub fn normal_block(phys: PhysicalAddress, attrs: PageAttributes) -> Self {
        Self(
            phys.into_u64()
                | (PageAttributes::BLOCK
                    | PageAttributes::PRESENT
                    | PageAttributes::ACCESS
                    | PageAttributes::SH_OUTER
                    | PageAttributes::PAGE_ATTR_NORMAL
                    | attrs)
                    .bits(),
            PhantomData,
        )
    }

    pub fn device_block(phys: PhysicalAddress) -> Self {
        Self(
            phys.into_u64()
                | (PageAttributes::BLOCK
                    | PageAttributes::PRESENT
                    | PageAttributes::ACCESS
                    | PageAttributes::SH_OUTER
                    | PageAttributes::PAGE_ATTR_DEVICE
                    | PageAttributes::UXN
                    | PageAttributes::PXN)
                    .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
        }
    }

    pub fn classify(self) -> EntryType {
        if !self.is_present() {
            EntryType::Invalid
        } else if let Some(table) = self.as_table() {
            EntryType::Table(table)
        } else {
            EntryType::Page(PhysicalAddress::from_u64(self.0 & !Self::ATTR_MASK))
        }
    }
}

impl PageEntry<L3> {
    pub fn normal_page(phys: PhysicalAddress, attrs: PageAttributes) -> Self {
        Self(
            phys.into_u64()
                | (PageAttributes::PAGE
                    | PageAttributes::PRESENT
                    | PageAttributes::ACCESS
                    | PageAttributes::SH_INNER
                    | PageAttributes::PAGE_ATTR_NORMAL
                    | attrs)
                    .bits(),
            PhantomData,
        )
    }

    pub fn device_page(phys: PhysicalAddress) -> Self {
        Self(
            phys.into_u64()
                | (PageAttributes::PAGE
                    | PageAttributes::PRESENT
                    | PageAttributes::ACCESS
                    | PageAttributes::SH_OUTER
                    | PageAttributes::PAGE_ATTR_DEVICE)
                    .bits(),
            PhantomData,
        )
    }

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

    pub fn as_page(&self) -> Option<PhysicalAddress> {
        let mask = (PageAttributes::PRESENT | PageAttributes::PAGE).bits();
        if self.0 & mask == mask {
            Some(PhysicalAddress::from_u64(self.0 & !Self::ATTR_MASK))
        } else {
            None
        }
    }
}

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

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

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

impl From<MapAttributes> for PageAttributes {
    fn from(value: MapAttributes) -> Self {
        let mut out = PageAttributes::empty();
        // TODO kernel cannot write el0 readonly pages
        if value.contains(MapAttributes::USER_WRITE) {
            // Read/write
            out |= PageAttributes::AP_BOTH_READWRITE;
        } else if value.contains(MapAttributes::USER_READ) {
            // Read only
            out |= PageAttributes::AP_BOTH_READONLY;
        } else {
            // No read/write
            out |= PageAttributes::AP_KERNEL_READONLY;
        }

        if value.contains(MapAttributes::DIRTY) {
            out |= PageAttributes::DIRTY;
        }

        if value.contains(MapAttributes::NON_GLOBAL) {
            out |= PageAttributes::NON_GLOBAL;
        }

        out
    }
}

impl From<PageAttributes> for MapAttributes {
    fn from(value: PageAttributes) -> Self {
        let mut out = MapAttributes::empty();

        out |= match value.intersection(PageAttributes::AP_ACCESS_MASK) {
            PageAttributes::AP_BOTH_READWRITE => {
                MapAttributes::USER_WRITE | MapAttributes::USER_READ
            }
            PageAttributes::AP_BOTH_READONLY => MapAttributes::USER_READ,
            PageAttributes::AP_KERNEL_READONLY => MapAttributes::empty(),
            PageAttributes::AP_KERNEL_READWRITE => panic!("This variant cannot be constructed"),
            _ => unreachable!(),
        };

        if value.contains(PageAttributes::DIRTY) {
            out |= MapAttributes::DIRTY;
        }

        if value.contains(PageAttributes::NON_GLOBAL) {
            out |= MapAttributes::NON_GLOBAL;
        }

        out
    }
}

impl fmt::Display for EntryType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Table(address) => write!(f, "table @ {address:#x}"),
            Self::Page(address) => write!(f, "page @ {address:#x}"),
            Self::Invalid => f.write_str("<invalid>"),
        }
    }
}
