//! Physical memory manager implementation
use core::sync::atomic::{AtomicUsize, Ordering};

use kernel_arch::KernelTableManagerImpl;
use libk_mm_interface::{address::PhysicalAddress, pointer::PhysicalRefMut};
use yggdrasil_abi::{error::Error, system::SystemMemoryStats};

use crate::L3_PAGE_SIZE;

#[cfg(any(target_pointer_width = "64", rust_analyzer))]
pub type BitmapWord = u64;
#[cfg(any(target_pointer_width = "32", rust_analyzer))]
pub type BitmapWord = u32;

pub(super) const BITMAP_WORD_SIZE: usize = BitmapWord::BITS as usize;

#[cfg(any(target_pointer_width = "64", rust_analyzer))]
pub(super) const BITMAP_PAGE_COUNT: usize = 512;

#[cfg(any(target_pointer_width = "32", rust_analyzer))]
pub(super) const BITMAP_PAGE_COUNT: usize = 32;

const HUGE_PAGE_WORD_COUNT: usize = 512 / BITMAP_WORD_SIZE;

pub(super) const TRACKED_PAGE_LIMIT: usize = (BITMAP_PAGE_COUNT * L3_PAGE_SIZE) * 8;

struct MemoryStats {
    available_pages: AtomicUsize,
    used_pages: AtomicUsize,
}

static STATS: MemoryStats = MemoryStats {
    available_pages: AtomicUsize::new(0),
    used_pages: AtomicUsize::new(0),
};

/// Physical memory management interface
pub struct PhysicalMemoryManager {
    bitmap: PhysicalRefMut<'static, [BitmapWord], KernelTableManagerImpl>,
    last_free_bit: usize,
    offset: usize,
    page_count: usize,
}

impl PhysicalMemoryManager {
    pub unsafe fn new(
        bitmap_phys_base: PhysicalAddress,
        offset: usize,
        page_count: usize,
    ) -> PhysicalMemoryManager {
        let bitmap_len = page_count.div_ceil(BITMAP_WORD_SIZE);
        let mut bitmap = PhysicalRefMut::<'static, _, KernelTableManagerImpl>::map_slice(
            bitmap_phys_base,
            bitmap_len,
        );

        bitmap.fill(BitmapWord::MAX);

        Self {
            bitmap,
            page_count,
            offset,
            last_free_bit: 0,
        }
    }

    #[inline]
    fn mark_alloc(&mut self, index: usize) {
        self.bitmap[index / BITMAP_WORD_SIZE] |= 1 << (index & (BITMAP_WORD_SIZE - 1));
    }

    #[inline]
    fn mark_free(&mut self, index: usize) {
        self.bitmap[index / BITMAP_WORD_SIZE] &= !(1 << (index & (BITMAP_WORD_SIZE - 1)));
    }

    #[inline(always)]
    fn is_alloc(&self, index: usize) -> bool {
        self.bitmap[index / BITMAP_WORD_SIZE] & (1 << (index & (BITMAP_WORD_SIZE - 1))) != 0
    }

    /// Allocates a single page, marking it as used with `usage`
    pub fn alloc_page(&mut self) -> Result<PhysicalAddress, Error> {
        for i in self.last_free_bit..self.page_count {
            if self.is_alloc(i) {
                continue;
            }

            self.last_free_bit = i + 1;
            self.mark_alloc(i);

            STATS.used_pages.fetch_add(1, Ordering::Relaxed);

            return Ok(PhysicalAddress::from_usize(i * L3_PAGE_SIZE + self.offset));
        }

        if self.last_free_bit != 0 {
            self.last_free_bit = 0;
            self.alloc_page()
        } else {
            Err(Error::OutOfMemory)
        }
    }

    pub fn alloc_2m_page(&mut self) -> Result<PhysicalAddress, Error> {
        let aligned_bit = self.last_free_bit & !511;

        'l0: for i in (aligned_bit..self.page_count).step_by(512) {
            for j in 0..HUGE_PAGE_WORD_COUNT {
                if self.bitmap[i / BITMAP_WORD_SIZE + j] != 0 {
                    continue 'l0;
                }
            }

            for j in 0..HUGE_PAGE_WORD_COUNT {
                self.bitmap[i / BITMAP_WORD_SIZE + j] = BitmapWord::MAX;
            }
            self.last_free_bit = i + 512;

            STATS.used_pages.fetch_add(512, Ordering::Relaxed);

            return Ok(PhysicalAddress::from_usize(i * L3_PAGE_SIZE + self.offset));
        }

        if self.last_free_bit != 0 {
            self.last_free_bit = 0;
            self.alloc_2m_page()
        } else {
            Err(Error::OutOfMemory)
        }
    }

    /// Allocates a contiguous range of physical pages, marking it as used with `usage`
    pub fn alloc_contiguous_pages(&mut self, count: usize) -> Result<PhysicalAddress, Error> {
        'l0: for i in self.last_free_bit..self.page_count {
            for j in 0..count {
                if self.is_alloc(i + j) {
                    continue 'l0;
                }
            }

            for j in 0..count {
                self.mark_alloc(i + j);
            }
            self.last_free_bit = i + count;

            STATS.used_pages.fetch_add(count, Ordering::Relaxed);

            return Ok(PhysicalAddress::from_usize(i * L3_PAGE_SIZE + self.offset));
        }

        if self.last_free_bit != 0 {
            self.last_free_bit = 0;
            self.alloc_contiguous_pages(count)
        } else {
            Err(Error::OutOfMemory)
        }
    }

    /// Deallocates a physical memory page.
    ///
    /// # Safety
    ///
    /// `addr` must be a page-aligned physical address previously allocated by this implementation.
    pub unsafe fn free_page(&mut self, page: PhysicalAddress) {
        let page = page.try_into_usize().unwrap();
        assert!(page >= self.offset);
        let index = (page - self.offset) / L3_PAGE_SIZE;

        STATS.used_pages.fetch_sub(1, Ordering::Relaxed);

        assert!(self.is_alloc(index));
        self.mark_free(index);
    }

    /// Marks a previously reserved page as available.
    ///
    /// # Panics
    ///
    /// Will panic if the address does not point to a valid, reserved (and unallocated) page.
    pub fn add_available_page(&mut self, page: PhysicalAddress) {
        let page = page.try_into_usize().unwrap();
        assert!(page >= self.offset);
        let index = (page - self.offset) / L3_PAGE_SIZE;

        STATS.available_pages.fetch_add(1, Ordering::Relaxed);

        assert!(self.is_alloc(index));
        self.mark_free(index);
    }

    /// Returns memory usage stats
    pub fn stats() -> SystemMemoryStats {
        let available = STATS.available_pages.load(Ordering::Relaxed);
        let used = STATS.used_pages.load(Ordering::Relaxed);
        let free = available - used;

        SystemMemoryStats {
            total_usable_pages: available,
            allocated_pages: used,
            free_pages: free,
            page_size: L3_PAGE_SIZE,
        }
    }
}
