use core::marker::PhantomData;

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

use crate::{mem::flush_tlb_entry, KernelTableManagerImpl};

use super::{
    fixed::clone_kernel_tables,
    table::{PageEntry, PageTable, L0, L3},
};

#[repr(C)]
pub struct ProcessAddressSpaceImpl<TA: TableAllocator> {
    l0: PhysicalRefMut<'static, PageTable<L0>, KernelTableManagerImpl>,
    _alloc: PhantomData<TA>,
}

impl<TA: TableAllocator> ProcessAddressSpaceManager<TA> for ProcessAddressSpaceImpl<TA> {
    const UPPER_LIMIT_PFN: usize = KERNEL_VIRT_OFFSET >> L3::SHIFT;
    const LOWER_LIMIT_PFN: usize = 32;

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

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

        clone_kernel_tables(&mut l0);

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

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

    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::page(physical, flags.into()), false)
    }

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

    fn as_address_with_asid(&self) -> u64 {
        unsafe { self.l0.as_physical_address().into_u64() }
    }
}

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 l3i = virt.page_index::<L3>();

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

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

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

        Ok(())
    }

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

        let mut l3 = self.l0.get_mut(l0i).ok_or(Error::DoesNotExist)?;
        let page = l3[l3i].as_page().ok_or(Error::DoesNotExist)?;

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

        Ok(page)
    }

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

        let l3 = self.l0.get(l0i)?;
        let page = l3[l3i].as_page()?;

        Some((page, l3[l3i].attributes().into()))
    }
}