#![allow(clippy::missing_safety_doc)]
use aarch64_cpu::{
    asm::barrier,
    registers::{MAIR_EL1, SCTLR_EL1, TCR_EL1},
};
use kernel_arch_interface::{
    mem::{DeviceMemoryAttributes, KernelTableManager, RawDeviceMemoryMapping},
    sync::IrqSafeSpinlock,
    KERNEL_VIRT_OFFSET,
};
use libk_mm_interface::{address::PhysicalAddress, table::EntryLevel};
use tock_registers::interfaces::{ReadWriteable, Writeable};
use yggdrasil_abi::error::Error;

pub use intrinsics::*;

use crate::{mem::table::L1, ArchitectureImpl};

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

#[derive(Debug)]
pub struct KernelTableManagerImpl;

static KERNEL_MEMORY_LOCK: IrqSafeSpinlock<ArchitectureImpl, ()> = IrqSafeSpinlock::new(());

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 virtual (-> physical) 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 = KERNEL_MEMORY_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 = KERNEL_MEMORY_LOCK.lock();
        #[allow(static_mut_refs)]
        fixed::DEVICE_MEMORY.unmap_device_pages(mapping);
    }
}

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

#[inline]
pub fn auto_upper_address<T>(ptr: *const T) -> usize {
    let address = ptr.addr();
    if address < KERNEL_VIRT_OFFSET {
        address + KERNEL_VIRT_OFFSET
    } else {
        address
    }
}

fn setup_memory_attributes() {
    // TODO: Figure out why WriteBack_NonTransient_ReadWriteAlloc doesn't work on Pi 4B

    MAIR_EL1.write(
        //// Attribute 0 -- normal memory
        MAIR_EL1::Attr0_Normal_Inner::WriteBack_NonTransient +
         MAIR_EL1::Attr0_Normal_Outer::WriteBack_NonTransient +
         //// Attribute 1 -- normal non-cacheable memory
         MAIR_EL1::Attr0_Normal_Inner::NonCacheable +
         MAIR_EL1::Attr0_Normal_Outer::NonCacheable +
         //// Attribute 2 -- device memory
         MAIR_EL1::Attr1_Device::nonGathering_nonReordering_EarlyWriteAck,
    );
}

unsafe fn enable_mmu() {
    TCR_EL1.write(
        TCR_EL1::AS::ASID8Bits +
         TCR_EL1::A1::TTBR0 +
         TCR_EL1::HD::CLEAR +
         // General
         TCR_EL1::IPS::Bits_48 +
         // TTBR0
         TCR_EL1::TG0::KiB_4 + TCR_EL1::T0SZ.val(25) + TCR_EL1::SH0::Inner +
         // TTBR1
         TCR_EL1::TG1::KiB_4 + TCR_EL1::T1SZ.val(25) + TCR_EL1::SH1::Inner,
    );

    barrier::dsb(barrier::ISHST);
    barrier::isb(barrier::SY);

    SCTLR_EL1.modify(
        SCTLR_EL1::E0E::LittleEndian
            + SCTLR_EL1::EE::LittleEndian
            + SCTLR_EL1::WXN::Disable
            + SCTLR_EL1::SA0::Enable
            + SCTLR_EL1::SA::Enable
            + SCTLR_EL1::A::Enable
            + SCTLR_EL1::I::NonCacheable
            + SCTLR_EL1::C::NonCacheable,
    );

    barrier::dsb(barrier::ISH);
    barrier::isb(barrier::SY);

    // Enable translation
    SCTLR_EL1.modify(SCTLR_EL1::M::Enable);

    // Enable caches
    enable_icache();
    enable_dcache();
}

pub unsafe fn init_lower(bsp: bool) {
    setup_memory_attributes();
    if bsp {
        fixed::setup();
    }
    fixed::load();
    enable_mmu();
}
