use core::ops::Range;

use kernel_arch_interface::{mem::DeviceMemoryAttributes, sync::IrqSafeSpinlock, Architecture};
use kernel_arch_x86::registers::CR3;
use libk_mm_interface::{
    address::PhysicalAddress,
    device::{DevicePageManager, DevicePageTableLevel},
    table::{page_index, EntryLevel},
};

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

pub const IDENTITY_SIZE_L1: usize = 64;
pub const KERNEL_L0I: usize = page_index::<L0>(KERNEL_VIRT_OFFSET);

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);

pub static LOCK: IrqSafeSpinlock<ArchitectureImpl, ()> = IrqSafeSpinlock::new(());
pub static mut KERNEL_PDPT: PageTable<L1> = PageTable::zeroed();
pub static mut KERNEL_PML4: PageTable<L0> = PageTable::zeroed();

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

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

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::<L2>::block(physical, PageAttributes::WRITABLE);
    }

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

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

    fn flush_range(range: Range<usize>) {
        let _ = range;
    }
}

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::page(physical, PageAttributes::WRITABLE);
    }

    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 _ = range;
    }
}

pub(super) unsafe fn setup(have_1gib_pages: bool) {
    let phys = PhysicalAddress::from_usize(auto_lower_address(&raw const KERNEL_PDPT));
    KERNEL_PML4[KERNEL_L0I] = PageEntry::table(phys, PageAttributes::WRITABLE);

    if have_1gib_pages {
        for i in 0..IDENTITY_SIZE_L1 {
            let phys = PhysicalAddress::from_usize(i * L1::SIZE);
            KERNEL_PDPT[i] = PageEntry::<L1>::block(phys, PageAttributes::WRITABLE);
        }
    } else {
        // TODO
        ArchitectureImpl::halt();
    }

    // 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::WRITABLE);
    }
    let phys = PhysicalAddress::from_usize(auto_lower_address(&raw const DEVICE_MEMORY.large.0));
    KERNEL_PDPT[DEVICE_L1] = PageEntry::table(phys, PageAttributes::WRITABLE);
}

pub(super) unsafe fn load() {
    CR3.set_address(auto_lower_address(&raw const KERNEL_PML4));
}
