use kernel_arch_interface::mem::{
    DeviceMemoryAttributes, KernelTableManager, RawDeviceMemoryMapping,
};
use libk_mm_interface::{address::PhysicalAddress, table::EntryLevel};
use yggdrasil_abi::error::Error;

use crate::KERNEL_VIRT_OFFSET;

use self::table::{PageTable, L0, L1};

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

#[derive(Debug)]
pub struct KernelTableManagerImpl;

impl KernelTableManager for KernelTableManagerImpl {
    fn virtualize(address: u64) -> usize {
        let address = address as usize;
        if address < fixed::IDENTITY_SIZE_L1 * L1::SIZE {
            address + KERNEL_VIRT_OFFSET
        } else {
            panic!("Invalid physical address: {address:#x}");
        }
    }

    fn physicalize(address: usize) -> u64 {
        if address < KERNEL_VIRT_OFFSET
            || address - KERNEL_VIRT_OFFSET >= fixed::IDENTITY_SIZE_L1 * L1::SIZE
        {
            panic!("Invalid virtualized address: {address:#x}");
        }
        (address - KERNEL_VIRT_OFFSET) as u64
    }

    unsafe fn map_device_pages(
        base: u64,
        count: usize,
        attrs: DeviceMemoryAttributes,
    ) -> Result<RawDeviceMemoryMapping<Self>, Error> {
        let _lock = fixed::LOCK.lock();
        #[allow(static_mut_refs)]
        fixed::DEVICE_MEMORY.map_device_pages(PhysicalAddress::from_u64(base), count, attrs)
    }

    unsafe fn unmap_device_pages(mapping: &RawDeviceMemoryMapping<Self>) {
        let _lock = fixed::LOCK.lock();
        #[allow(static_mut_refs)]
        fixed::DEVICE_MEMORY.unmap_device_pages(mapping);
    }
}

pub fn clone_kernel_tables(dst: &mut PageTable<L0>) {
    unsafe {
        dst[fixed::KERNEL_L0I] = fixed::KERNEL_PML4[fixed::KERNEL_L0I];
    }
}

pub fn auto_lower_address<T>(pointer: *const T) -> usize {
    let address = pointer.addr();
    if address < KERNEL_VIRT_OFFSET {
        address
    } else {
        address - KERNEL_VIRT_OFFSET
    }
}

/// Sets up the following memory map:
/// ...: KERNEL_TABLES.l0:
///     * 0xFFFFFF0000000000 .. 0xFFFFFFFF8000000000        : RAM_MAPPING_L1
///     * 0xFFFFFF8000000000 .. ...                         : KERNEL_TABLES.kernel_l1:
///         * 0xFFFFFF8000000000 .. 0xFFFFFF8040000000      : KERNEL_TABLES.kernel_l2
///             * 0xFFFFFF8000000000 .. 0xFFFFFF8000200000  : ---
///             * 0xFFFFFF8000200000 .. 0xFFFFFF8000400000  : EARLY_MAPPING_L3
///             * 0xFFFFFF8000400000 .. ...                 : KERNEL_TABLES.kernel_l3s
///         * 0xFFFFFF8040000000 .. 0xFFFFFF8080000000      : ---
///         * 0xFFFFFF8080000000 .. 0xFFFFFF8100000000      : DEVICE_MAPPING_L2
///             * 0xFFFFFF8080000000 .. 0xFFFFFF8080800000  : DEVICE_MAPPING_L3S
///             * 0xFFFFFF8080800000 .. 0xFFFFFF8100000000  : ...
///
/// # Safety
///
/// Unsafe, must only be called by BSP during its early init, must already be in "higher-half"
#[inline(never)]
pub unsafe fn init_fixed_tables(have_1gib_pages: bool, bsp: bool) {
    fixed::setup(have_1gib_pages);
    if bsp {
        fixed::load();
    }
}

/// # 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));
}
