use limine::memory_map::Entry;
use x86_64::PhysAddr;

use super::RawSlice;
use super::addr;
use super::phys::BitmapFrameAllocator;
use crate::sync::IrqMutex;

const PAGE_SIZE: u64 = 4096;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RefcountError {
    OutOfRange(PhysAddr),
    Overflow(PhysAddr),
    Underflow(PhysAddr),
}

static REFCOUNTS: IrqMutex<RawSlice<u16>, 4> = IrqMutex::new(RawSlice::empty());

pub fn init(memory_map: &[&Entry]) {
    let max_addr = memory_map
        .iter()
        .filter(|entry| entry.entry_type == limine::memory_map::EntryType::USABLE)
        .map(|entry| entry.base + entry.length)
        .fold(0u64, u64::max);

    let total_frames = (max_addr / PAGE_SIZE) as usize;
    let bytes_needed = total_frames * 2;
    let frames_needed = bytes_needed.div_ceil(PAGE_SIZE as usize);

    let allocator = BitmapFrameAllocator;
    let phys_base = allocator
        .allocate_contiguous(frames_needed)
        .expect("[refcount] failed to allocate contiguous frames for refcount table");

    let virt = addr::phys_to_virt(phys_base);
    let ptr = virt.as_u64() as *mut u16;

    unsafe {
        core::ptr::write_bytes(ptr, 0, total_frames);
    }

    let mut rc = REFCOUNTS.lock();
    rc.init(ptr, total_frames);
}

pub fn increment(phys: PhysAddr) -> Result<(), RefcountError> {
    let idx = (phys.as_u64() / PAGE_SIZE) as usize;
    let mut rc = REFCOUNTS.lock();
    let slot = rc
        .as_slice_mut()
        .get_mut(idx)
        .ok_or(RefcountError::OutOfRange(phys))?;
    let new = slot.checked_add(1).ok_or(RefcountError::Overflow(phys))?;
    *slot = new;
    Ok(())
}

pub fn decrement(phys: PhysAddr) -> Result<u16, RefcountError> {
    let idx = (phys.as_u64() / PAGE_SIZE) as usize;
    let mut rc = REFCOUNTS.lock();
    let slot = rc
        .as_slice_mut()
        .get_mut(idx)
        .ok_or(RefcountError::OutOfRange(phys))?;
    let new = slot.checked_sub(1).ok_or(RefcountError::Underflow(phys))?;
    *slot = new;
    Ok(new)
}

#[allow(dead_code)]
pub fn get(phys: PhysAddr) -> u16 {
    let idx = (phys.as_u64() / PAGE_SIZE) as usize;
    let rc = REFCOUNTS.lock();
    if idx < rc.len() {
        rc.as_slice()[idx]
    } else {
        0
    }
}