//! AArch64-specific process address space management
use core::{
    marker::PhantomData,
    sync::atomic::{AtomicU8, Ordering},
};

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::{mem::table::PageEntry, KernelTableManagerImpl};

use super::{
    dc_cvac, ic_iallu,
    table::{PageAttributes, PageTable, L1, L2, L3},
    tlb_flush_asid, tlb_flush_vaae1,
};

/// AArch64 implementation of a process address space table
#[repr(C)]
pub struct ProcessAddressSpaceImpl<TA: TableAllocator> {
    l1: PhysicalRefMut<'static, PageTable<L1>, KernelTableManagerImpl>,
    asid: u8,
    _alloc: PhantomData<TA>,
}

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

    fn new() -> Result<Self, Error> {
        static LAST_ASID: AtomicU8 = AtomicU8::new(1);

        let asid = LAST_ASID.fetch_add(1, Ordering::AcqRel);

        let mut l1 = unsafe {
            PhysicalRefMut::<'static, PageTable<L1>, KernelTableManagerImpl>::map(
                TA::allocate_page_table()?,
            )
        };

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

        tlb_flush_asid(asid);

        Ok(Self {
            l1,
            asid,
            _alloc: PhantomData,
        })
    }

    fn translate(&self, address: usize) -> Result<(PhysicalAddress, MapAttributes), Error> {
        self.read_l3_entry(address).ok_or(Error::DoesNotExist)
    }

    unsafe fn map_page(
        &mut self,
        address: usize,
        physical: PhysicalAddress,
        flags: MapAttributes,
    ) -> Result<(), Error> {
        self.write_l3_entry(
            address,
            PageEntry::normal_page(
                physical,
                PageAttributes::from(flags) | PageAttributes::NON_GLOBAL,
            ),
            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)
    }

    fn as_address_with_asid(&self) -> (u64, u64) {
        let physical = unsafe { u64::from(self.l1.as_physical_address()) };
        (physical, self.asid as u64)
    }

    unsafe fn clear(&mut self) {
        self.l1
            .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 l1i = virt.page_index::<L1>();
        let l2i = virt.page_index::<L2>();
        let l3i = virt.page_index::<L3>();

        let mut l2 = self.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;
        dc_cvac((&raw const l3[l3i]).addr());
        tlb_flush_vaae1(virt);

        Ok(())
    }

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

        let mut l2 = self.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)?;
        ic_iallu();
        dc_cvac((&raw const l3[l3i]).addr());
        tlb_flush_vaae1(virt);

        Ok(())
    }

    fn pop_l3_entry(&mut self, virt: usize) -> Result<(PhysicalAddress, bool), Error> {
        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 l2 = self.l1.get_mut(l1i).ok_or(Error::DoesNotExist)?;
        let mut l3 = l2.get_mut(l2i).ok_or(Error::DoesNotExist)?;

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

        l3[l3i] = PageEntry::INVALID;
        ic_iallu();
        dc_cvac((&raw const l3[l3i]).addr());
        tlb_flush_vaae1(virt);

        Ok((page, dirty))
    }

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

        let l2 = self.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 l1_phys = self.l1.as_physical_address();
            TA::free_page_table(l1_phys);
        }
    }
}
