use core::ops::Range;

use kernel_arch_interface::mem::{
    DeviceMemoryAttributes, KernelTableManager, RawDeviceMemoryMapping,
};
use yggdrasil_abi::error::Error;

use crate::{
    address::PhysicalAddress,
    table::{EntryLevel, EntryLevelExt},
};

pub trait DevicePageTableLevel {
    type Level: EntryLevel;
    const INDEX_RANGE: Range<usize>;
    const VIRTUAL_BASE: usize;

    fn map_page(&mut self, index: usize, physical: PhysicalAddress, attrs: &DeviceMemoryAttributes);
    fn unmap_page(&mut self, index: usize);
    fn is_mapped(&self, index: usize) -> bool;
    fn flush_range(range: Range<usize>);

    fn allocate_mapping(
        &mut self,
        base: PhysicalAddress,
        page_count: usize,
        attrs: &DeviceMemoryAttributes,
    ) -> Option<usize> {
        'l0: for i in Self::INDEX_RANGE {
            for j in 0..page_count {
                if self.is_mapped(i + j) {
                    continue 'l0;
                }
            }

            for j in 0..page_count {
                self.map_page(i + j, base.add(j * Self::Level::SIZE), attrs);
            }

            let start = Self::VIRTUAL_BASE + i * Self::Level::SIZE;
            Self::flush_range(i..i + page_count);
            return Some(start);
        }
        None
    }

    /// # Safety
    ///
    /// Removes validity of memory regions, unsafe.
    unsafe fn remove_mapping(&mut self, base: usize, page_count: usize) {
        let base_index = (base - Self::VIRTUAL_BASE) / Self::Level::SIZE;
        for i in 0..page_count {
            self.unmap_page(base_index + i);
        }
    }
}

pub struct DevicePageManager<N: DevicePageTableLevel, L: DevicePageTableLevel> {
    pub normal: N,
    pub large: L,
}

impl<N: DevicePageTableLevel, L: DevicePageTableLevel> DevicePageManager<N, L> {
    pub const fn new(normal: N, large: L) -> Self {
        Self { normal, large }
    }

    /// # Safety
    ///
    /// Provides ability to map arbitrary ranges of memory, unsafe.
    pub unsafe fn map_device_pages<A: KernelTableManager>(
        &mut self,
        base: PhysicalAddress,
        size: usize,
        attrs: DeviceMemoryAttributes,
    ) -> Result<RawDeviceMemoryMapping<A>, Error> {
        let small_aligned_base = base.page_align_down::<N::Level>();
        let small_aligned_end = base.add(size).page_align_up::<N::Level>();
        let small_offset = base - small_aligned_base;
        let small_page_count = (small_aligned_end - small_aligned_base).page_count::<N::Level>();

        if small_page_count > 128 {
            // Allocate from large page pool
            let large_aligned_base = base.page_align_down::<L::Level>();
            let large_aligned_end = base.add(size).page_align_up::<L::Level>();
            let large_offset = base - large_aligned_base;
            let large_page_count =
                (large_aligned_end - large_aligned_base).page_count::<L::Level>();

            let mapped_base = self
                .large
                .allocate_mapping(large_aligned_base, large_page_count, &attrs)
                .ok_or(Error::OutOfMemory)?;
            let mapped_address = mapped_base + large_offset;

            Ok(RawDeviceMemoryMapping::from_raw_parts(
                large_aligned_base.into_u64(),
                mapped_address,
                mapped_base,
                large_page_count,
                L::Level::SIZE,
            ))
        } else {
            // Allocate from small page pool
            let mapped_base = self
                .normal
                .allocate_mapping(small_aligned_base, small_page_count, &attrs)
                .ok_or(Error::OutOfMemory)?;
            let mapped_address = mapped_base + small_offset;

            Ok(RawDeviceMemoryMapping::from_raw_parts(
                small_aligned_base.into_u64(),
                mapped_address,
                mapped_base,
                small_page_count,
                N::Level::SIZE,
            ))
        }
    }

    /// # Safety
    ///
    /// Removes validity of memory regions, unsafe.
    pub unsafe fn unmap_device_pages<A: KernelTableManager>(
        &mut self,
        mapping: &RawDeviceMemoryMapping<A>,
    ) {
        if mapping.page_size == N::Level::SIZE {
            self.normal
                .remove_mapping(mapping.base_address, mapping.page_count);
        } else if mapping.page_size == L::Level::SIZE {
            self.large
                .remove_mapping(mapping.base_address, mapping.page_count);
        } else {
            unreachable!(
                "Invalid device memory mapping with page size {:#x}",
                mapping.page_size
            )
        }
    }
}
