//! x86-64-specific process address space management functions
use core::marker::PhantomData;

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

use crate::KernelTableManagerImpl;

use super::{
    clone_kernel_tables, flush_tlb_entry,
    table::{PageEntry, PageTable, L0, L1, L2, L3},
};

/// Represents a process or kernel address space. Because x86-64 does not have cool stuff like
/// TTBR0 and TTBR1, all address spaces are initially cloned from the kernel space.
#[repr(C)]
pub struct ProcessAddressSpaceImpl<TA: TableAllocator> {
    l0: PhysicalRefMut<'static, PageTable<L0>, KernelTableManagerImpl>,
    _alloc: PhantomData<TA>,
}

impl<TA: TableAllocator> ProcessAddressSpaceManager<TA> for ProcessAddressSpaceImpl<TA> {
    // Start with 8GiB
    const LOWER_LIMIT_PFN: usize = (8 << 30) / L3::SIZE;
    // 16GiB VM limit
    const UPPER_LIMIT_PFN: usize = (16 << 30) / L3::SIZE;

    fn new() -> Result<Self, Error> {
        let mut l0 = unsafe {
            PhysicalRefMut::<'static, PageTable<L0>, KernelTableManagerImpl>::map(
                TA::allocate_page_table()?,
            )
        };

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

        clone_kernel_tables(&mut l0);

        Ok(Self {
            l0,
            _alloc: PhantomData,
        })
    }

    #[inline]
    unsafe fn map_page(
        &mut self,
        address: usize,
        physical: PhysicalAddress,
        flags: MapAttributes,
    ) -> Result<(), Error> {
        self.write_l3_entry(address, PageEntry::page(physical, flags.into()), false)
    }

    unsafe fn update_page_attributes(
        &mut self,
        address: usize,
        update: &PageAttributeUpdate,
    ) -> Result<(), Error> {
        self.update_l3_entry(address, |entry| entry.update(update))
    }

    unsafe fn unmap_page(&mut self, address: usize) -> Result<(PhysicalAddress, bool), Error> {
        self.pop_l3_entry(address)
    }

    #[inline]
    fn translate(&self, address: usize) -> Result<(PhysicalAddress, MapAttributes), Error> {
        self.read_l3_entry(address)
            .ok_or(Error::InvalidMemoryOperation)
    }

    fn as_address_with_asid(&self) -> (u64, u64) {
        // TODO x86-64 PCID/ASID?
        (unsafe { self.l0.as_physical_address().into_u64() }, 0)
    }

    unsafe fn clear(&mut self) {
        self.l0
            .drop_range::<TA>(0..((Self::UPPER_LIMIT_PFN * L3::SIZE).page_index::<L1>()));
    }
}

impl<TA: TableAllocator> ProcessAddressSpaceImpl<TA> {
    // Write a single 4KiB entry
    fn write_l3_entry(
        &mut self,
        virt: usize,
        entry: PageEntry<L3>,
        overwrite: bool,
    ) -> Result<(), Error> {
        let l0i = virt.page_index::<L0>();
        let l1i = virt.page_index::<L1>();
        let l2i = virt.page_index::<L2>();
        let l3i = virt.page_index::<L3>();

        let mut l1 = self.l0.get_mut_or_alloc::<TA>(l0i)?;
        let mut l2 = l1.get_mut_or_alloc::<TA>(l1i)?;
        let mut l3 = l2.get_mut_or_alloc::<TA>(l2i)?;

        if l3[l3i].is_present() && !overwrite {
            todo!();
        }

        l3[l3i] = entry;
        unsafe {
            flush_tlb_entry(virt);
        }

        Ok(())
    }

    fn update_l3_entry<F: FnOnce(&mut PageEntry<L3>) -> Result<(), Error>>(
        &mut self,
        virt: usize,
        mapper: F,
    ) -> Result<(), Error> {
        let l0i = virt.page_index::<L0>();
        let l1i = virt.page_index::<L1>();
        let l2i = virt.page_index::<L2>();
        let l3i = virt.page_index::<L3>();

        let mut l1 = self.l0.get_mut(l0i).ok_or(Error::DoesNotExist)?;
        let mut l2 = l1.get_mut(l1i).ok_or(Error::DoesNotExist)?;
        let mut l3 = l2.get_mut(l2i).ok_or(Error::DoesNotExist)?;

        let entry = &mut l3[l3i];
        if !entry.is_present() {
            return Err(Error::DoesNotExist);
        }
        mapper(entry)?;
        unsafe {
            flush_tlb_entry(virt);
        }

        Ok(())
    }

    fn pop_l3_entry(&mut self, virt: usize) -> Result<(PhysicalAddress, bool), Error> {
        let l0i = virt.page_index::<L0>();
        let l1i = virt.page_index::<L1>();
        let l2i = virt.page_index::<L2>();
        let l3i = virt.page_index::<L3>();

        // TODO somehow drop tables if they're known to be empty?
        let mut l1 = self.l0.get_mut(l0i).ok_or(Error::DoesNotExist)?;
        let mut l2 = l1.get_mut(l1i).ok_or(Error::DoesNotExist)?;
        let mut l3 = l2.get_mut(l2i).ok_or(Error::DoesNotExist)?;

        let page = l3[l3i].as_page().ok_or(Error::DoesNotExist)?;
        let dirty = l3[l3i].is_dirty();

        l3[l3i] = PageEntry::INVALID;
        unsafe {
            flush_tlb_entry(virt);
        }

        Ok((page, dirty))
    }

    fn read_l3_entry(&self, virt: usize) -> Option<(PhysicalAddress, MapAttributes)> {
        let l0i = virt.page_index::<L0>();
        let l1i = virt.page_index::<L1>();
        let l2i = virt.page_index::<L2>();
        let l3i = virt.page_index::<L3>();

        let l1 = self.l0.get(l0i)?;
        let l2 = l1.get(l1i)?;
        let l3 = l2.get(l2i)?;

        let page = l3[l3i].as_page()?;

        Some((page.add(virt & 0xFFF), l3[l3i].attributes().into()))
    }
}

impl<TA: TableAllocator> Drop for ProcessAddressSpaceImpl<TA> {
    fn drop(&mut self) {
        // SAFETY: with safe usage of the ProcessAddressSpaceImpl, clearing and dropping
        // is safe, no one refers to the memory
        unsafe {
            self.clear();
            let l0_phys = self.l0.as_physical_address();
            TA::free_page_table(l0_phys);
        }
    }
}
