use kernel_arch_interface::mem::{
    DeviceMemoryAttributes, KernelTableManager, RawDeviceMemoryMapping,
};
use libk_mm_interface::{
    address::PhysicalAddress,
    table::{page_index, EntryLevel, EntryLevelExt},
};
use tock_registers::interfaces::Writeable;
use yggdrasil_abi::error::Error;

use crate::{
    mem::table::{PageTable, L1, L3},
    registers::SATP,
};

pub use intrinsics::*;

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

pub const KERNEL_VIRT_OFFSET: usize = kernel_arch_interface::KERNEL_VIRT_OFFSET;

/// Any VAs above this one are sign-extended
pub const USER_BOUNDARY: usize = 0x40_00000000;

#[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 _ = attrs;
        let _lock = fixed::LOCK.lock();
        let base = PhysicalAddress::from_u64(base);
        let l3_aligned_base = base.page_align_down::<L3>();
        let l3_aligned_end = base.add(count).page_align_up::<L3>();
        let l3_offset = base - l3_aligned_base;
        let l3_page_count = (l3_aligned_end - l3_aligned_base).page_count::<L3>();
        let l3_aligned_virt = l3_aligned_base.add(KERNEL_VIRT_OFFSET).into_usize();

        Ok(unsafe {
            RawDeviceMemoryMapping::from_raw_parts(
                l3_aligned_base.into_u64(),
                l3_aligned_virt + l3_offset,
                l3_aligned_virt,
                l3_page_count,
                L3::SIZE,
            )
        })
    }

    unsafe fn unmap_device_pages(mapping: &RawDeviceMemoryMapping<Self>) {
        let _ = mapping;
    }
}

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

/// Enables the memory translation.
///
/// # Safety
///
/// Only meant to be called once per each HART during their early init.
pub unsafe fn enable_mmu() {
    let l1_phys = auto_lower_address(&raw const fixed::KERNEL_L1) as u64;
    tlb_flush_full();
    SATP.write(SATP::PPN.val(l1_phys >> 12) + SATP::MODE::Sv39);
}

pub fn clone_kernel_tables(dst: &mut PageTable<L1>) {
    let _lock = fixed::LOCK.lock();
    for l1i in page_index::<L1>(USER_BOUNDARY)..512 {
        dst[l1i] = unsafe { fixed::KERNEL_L1[l1i] };
    }
}
