use fixed::FixedTables;
use kernel_arch_interface::{
    mem::{DeviceMemoryAttributes, KernelTableManager, RawDeviceMemoryMapping},
    KERNEL_VIRT_OFFSET,
};
use libk_mm_interface::{address::PhysicalAddress, KernelImageObject};
use table::{PageAttributes, PageEntry};
use yggdrasil_abi::error::Error;

pub mod fixed;
pub mod process;
pub mod table;

pub use process::ProcessAddressSpaceImpl;

#[derive(Debug)]
pub struct KernelTableManagerImpl;

#[link_section = ".data.tables"]
pub static mut KERNEL_TABLES: KernelImageObject<FixedTables> =
    unsafe { KernelImageObject::new(FixedTables::zeroed()) };

impl KernelTableManager for KernelTableManagerImpl {
    unsafe fn map_device_pages(
        base: u64,
        count: usize,
        _attrs: DeviceMemoryAttributes,
    ) -> Result<RawDeviceMemoryMapping<Self>, Error> {
        // TODO page align up
        let end = base + count as u64;
        if end < fixed::MAX_FIXED_PHYSICAL.into_u64() {
            // 1:1
            let address = Self::virtualize(base);
            Ok(RawDeviceMemoryMapping::from_raw_parts(
                address, address, 0, 0,
            ))
        } else {
            loop {}
        }
    }

    unsafe fn unmap_device_pages(mapping: &RawDeviceMemoryMapping<Self>) {
        loop {}
    }

    fn virtualize(phys: u64) -> usize {
        unsafe { KERNEL_TABLES.virtualize(PhysicalAddress::from_u64(phys)) }
    }

    fn physicalize(virt: usize) -> u64 {
        unsafe { KERNEL_TABLES.physicalize(virt) }
            .expect("Invalid virtual address")
            .into_u64()
    }
}

pub unsafe fn init_fixed_tables() {
    // Unmap lower stuff
    for (i, entry) in KERNEL_TABLES.l0.lower.iter_mut().enumerate() {
        *entry = PageEntry::INVALID;
        flush_tlb_entry(i << 22);
    }

    // Map the rest of fixed translation
    for (i, entry) in KERNEL_TABLES.l0.kernel.iter_mut().enumerate() {
        let virt = KERNEL_VIRT_OFFSET + i << 22;
        let phys = (i << 22) as u32;
        *entry = PageEntry::block(PhysicalAddress::from_u32(phys), PageAttributes::WRITABLE);
        flush_tlb_entry(virt);
    }
}

/// # Safety
///
/// `address` must be page-aligned.
#[inline]
pub unsafe fn flush_tlb_entry(address: usize) {
    core::arch::asm!("invlpg ({0})", in(reg) address, options(att_syntax));
}