use core::ops::Range;

use aarch64_cpu::registers::{TTBR0_EL1, TTBR1_EL1};
use kernel_arch_interface::{mem::DeviceMemoryAttributes, KERNEL_VIRT_OFFSET};
use libk_mm_interface::{
    address::PhysicalAddress,
    device::{DevicePageManager, DevicePageTableLevel},
    table::EntryLevel,
};

use crate::mem::{
    auto_lower_address,
    table::{PageAttributes, PageEntry, PageTable, L1, L2, L3},
    tlb_flush_range_va,
};

pub const IDENTITY_SIZE_L1: usize = 8;
pub const DEVICE_L1: usize = IDENTITY_SIZE_L1;
pub const DEVICE_MAPPING_L3_COUNT: usize = 32;
pub const DEVICE_MAPPING_OFFSET: usize = KERNEL_VIRT_OFFSET + (DEVICE_L1 << L1::SHIFT);

#[repr(transparent)]
pub struct L2DeviceMemory(pub PageTable<L2>);
#[repr(transparent)]
pub struct L3DeviceMemory(pub [PageTable<L3>; DEVICE_MAPPING_L3_COUNT]);

static mut KERNEL_L1: PageTable<L1> = PageTable::zeroed();
pub(super) static mut DEVICE_MEMORY: DevicePageManager<L3DeviceMemory, L2DeviceMemory> =
    DevicePageManager::new(
        L3DeviceMemory([PageTable::zeroed(); DEVICE_MAPPING_L3_COUNT]),
        L2DeviceMemory(PageTable::zeroed()),
    );

impl DevicePageTableLevel for L2DeviceMemory {
    type Level = L2;
    const VIRTUAL_BASE: usize = DEVICE_MAPPING_OFFSET;
    const INDEX_RANGE: Range<usize> = DEVICE_MAPPING_L3_COUNT..512;

    fn map_page(
        &mut self,
        index: usize,
        physical: PhysicalAddress,
        attrs: &DeviceMemoryAttributes,
    ) {
        let _ = attrs;
        self.0[index] = PageEntry::device_block(physical);
    }

    fn unmap_page(&mut self, index: usize) {
        self.0[index] = PageEntry::INVALID;
    }

    fn is_mapped(&self, index: usize) -> bool {
        self.0[index].is_present()
    }

    fn flush_range(range: Range<usize>) {
        let start = range.start * L2::SIZE + Self::VIRTUAL_BASE;
        let size = (range.end - range.start) * L2::SIZE;
        tlb_flush_range_va(start, size);
    }
}

impl DevicePageTableLevel for L3DeviceMemory {
    type Level = L3;
    const VIRTUAL_BASE: usize = DEVICE_MAPPING_OFFSET;
    const INDEX_RANGE: Range<usize> = 0..512 * DEVICE_MAPPING_L3_COUNT;

    fn map_page(
        &mut self,
        index: usize,
        physical: PhysicalAddress,
        attrs: &DeviceMemoryAttributes,
    ) {
        let _ = attrs;
        let l2i = index / 512;
        let l3i = index % 512;
        self.0[l2i][l3i] = PageEntry::device_page(physical);
    }

    fn unmap_page(&mut self, index: usize) {
        let l2i = index / 512;
        let l3i = index % 512;
        self.0[l2i][l3i] = PageEntry::INVALID;
    }

    fn is_mapped(&self, index: usize) -> bool {
        let l2i = index / 512;
        let l3i = index % 512;
        self.0[l2i][l3i].is_present()
    }

    fn flush_range(range: Range<usize>) {
        let start = range.start * L3::SIZE + Self::VIRTUAL_BASE;
        let size = (range.end - range.start) * L3::SIZE;
        tlb_flush_range_va(start, size);
    }
}

pub unsafe fn setup() {
    // 0..IDENTITY_SIZE_L1 -> lower RAM region
    for i in 0..IDENTITY_SIZE_L1 {
        let phys = PhysicalAddress::from_usize(i << L1::SHIFT);
        KERNEL_L1[i] = PageEntry::normal_block(phys, PageAttributes::empty());
    }

    // DEVICE_L1 -> Device L2 table
    //               0..DEVICE_MAPPING_L3_COUNT -> Device L3 tables -> Device L3 pages
    //               ..512                      -> Device L2 pages
    for i in 0..DEVICE_MAPPING_L3_COUNT {
        let phys =
            PhysicalAddress::from_usize(auto_lower_address(&raw const DEVICE_MEMORY.normal.0[i]));
        DEVICE_MEMORY.large.0[i] = PageEntry::table(phys, PageAttributes::empty());
    }
    let phys = PhysicalAddress::from_usize(auto_lower_address(&raw const DEVICE_MEMORY.large.0));
    KERNEL_L1[DEVICE_L1] = PageEntry::table(phys, PageAttributes::empty());
}

pub unsafe fn load() {
    let ttbr_physical = auto_lower_address(&raw const KERNEL_L1) as u64;

    TTBR0_EL1.set_baddr(ttbr_physical);
    TTBR1_EL1.set_baddr(ttbr_physical);
}
