use crate::mem::addr;
use crate::mem::phys::BitmapFrameAllocator;
use crate::sync::IrqMutex;
use lancer_core::object_layout::KernelObject;
use x86_64::PhysAddr;

const SLOTS_PER_FRAME: usize = 64;
const MAX_FRAMES: usize = 256;
const SLOT_SIZE: usize = 64;
const INIT_FRAMES: usize = 192;

struct KernelObjectAllocatorInner {
    frame_phys: [u64; MAX_FRAMES],
    bitmap: [u64; MAX_FRAMES],
    frame_count: usize,
}

impl KernelObjectAllocatorInner {
    const fn new() -> Self {
        Self {
            frame_phys: [0; MAX_FRAMES],
            bitmap: [0; MAX_FRAMES],
            frame_count: 0,
        }
    }

    fn install_frames(&mut self, phys_addrs: &[u64]) {
        phys_addrs
            .iter()
            .for_each(|&phys| match self.frame_count < MAX_FRAMES {
                true => {
                    self.frame_phys[self.frame_count] = phys;
                    self.bitmap[self.frame_count] = 0;
                    self.frame_count += 1;
                }
                false => {}
            });
    }

    fn alloc_slot(&mut self) -> Option<u64> {
        (0..self.frame_count).find_map(|fi| {
            let bm = self.bitmap[fi];
            match bm == u64::MAX {
                true => None,
                false => {
                    let bit = (!bm).trailing_zeros() as usize;
                    self.bitmap[fi] |= 1u64 << bit;
                    Some(self.frame_phys[fi] + (bit * SLOT_SIZE) as u64)
                }
            }
        })
    }

    fn free_slot(&mut self, phys: u64) -> bool {
        (0..self.frame_count)
            .find(|&fi| {
                let base = self.frame_phys[fi];
                phys >= base && phys < base + 4096
            })
            .map(|fi| {
                let offset = (phys - self.frame_phys[fi]) as usize;
                let bit = offset / SLOT_SIZE;
                let mask = 1u64 << bit;
                let was_allocated = self.bitmap[fi] & mask != 0;
                debug_assert!(
                    was_allocated,
                    "kernel_objects: double-free at phys {phys:#x}"
                );
                self.bitmap[fi] &= !mask;
                was_allocated
            })
            .unwrap_or(false)
    }
}

static KERNEL_OBJ_ALLOC: IrqMutex<KernelObjectAllocatorInner, 3> =
    IrqMutex::new(KernelObjectAllocatorInner::new());

pub fn init(allocator: &mut BitmapFrameAllocator) {
    let mut frame_addrs = [0u64; INIT_FRAMES];
    let count = (0..INIT_FRAMES)
        .scan((), |_, _| allocator.allocate())
        .enumerate()
        .map(|(i, frame)| {
            let phys = frame.phys_addr();
            let virt = addr::phys_to_virt(phys);
            unsafe {
                core::ptr::write_bytes(virt.as_mut_ptr::<u8>(), 0, 4096);
            }
            frame_addrs[i] = phys.as_u64();
            core::mem::forget(frame);
            i + 1
        })
        .last()
        .unwrap_or(0);

    {
        let mut alloc = KERNEL_OBJ_ALLOC.lock();
        alloc.install_frames(&frame_addrs[..count]);
    }

    crate::kprintln!(
        "  KernelObjectAllocator: {} frames ({} slots)",
        count,
        count * SLOTS_PER_FRAME,
    );
}

pub fn alloc_slot() -> Option<u64> {
    KERNEL_OBJ_ALLOC.lock().alloc_slot()
}

pub fn write_at<T: KernelObject>(phys: u64, obj: T) {
    let virt = addr::phys_to_virt(PhysAddr::new(phys));
    unsafe {
        core::ptr::write(virt.as_mut_ptr::<T>(), obj);
    }
}

pub fn free_slot(phys: u64) {
    let _ = KERNEL_OBJ_ALLOC.lock().free_slot(phys);
}