use super::object::ObjectData;
use super::object_table::{DecRefOutcome, ObjectTable};
use crate::error::KernelError;
use crate::mem::addr;
use crate::mem::phys::BitmapFrameAllocator;
use crate::sync::IrqMutex;
use crate::types::{Generation, ObjectId};
use x86_64::PhysAddr;

const SLAB_FRAMES: usize = 128;

pub struct ObjectPool {
    table: ObjectTable,
    kernel_region_phys: u64,
    kernel_region_capacity: u32,
}

impl ObjectPool {
    const fn new() -> Self {
        Self {
            table: ObjectTable::new(),
            kernel_region_phys: 0,
            kernel_region_capacity: 0,
        }
    }

    pub fn init(&mut self, allocator: &mut BitmapFrameAllocator) {
        let base_frame = allocator
            .allocate_contiguous(SLAB_FRAMES)
            .expect("ObjectPool kernel region allocation failed");

        self.kernel_region_phys = base_frame.as_u64();

        let base_virt = addr::phys_to_virt(base_frame);
        unsafe {
            core::ptr::write_bytes(base_virt.as_mut_ptr::<u8>(), 0, SLAB_FRAMES * 4096);
        }

        let obj_size = core::mem::size_of::<ObjectData>();
        let total_bytes = SLAB_FRAMES * 4096;
        self.kernel_region_capacity = (total_bytes / obj_size) as u32;

        self.table.init();

        crate::kprintln!(
            "  ObjectPool: {} kernel region slots ({} frames, entry={}B)",
            self.kernel_region_capacity,
            SLAB_FRAMES,
            obj_size,
        );
    }

    pub fn capacity(&self) -> usize {
        self.kernel_region_capacity as usize
    }

    pub fn kernel_region(&self) -> (u64, usize) {
        (self.kernel_region_phys, SLAB_FRAMES * 4096)
    }

    fn slot_phys(&self, id: u16) -> u64 {
        let obj_size = core::mem::size_of::<ObjectData>() as u64;
        self.kernel_region_phys + (id as u64) * obj_size
    }

    fn is_kernel_region_slot(&self, id: u16) -> bool {
        (id as u32) < self.kernel_region_capacity
    }

    pub fn allocate(&mut self, data: ObjectData) -> Result<(ObjectId, Generation), KernelError> {
        debug_assert!(self.kernel_region_phys != 0, "ObjectPool not initialized");

        let id = self.table.peek_free_head().ok_or(KernelError::PoolExhausted)?;
        match self.is_kernel_region_slot(id) {
            true => {
                let phys = self.slot_phys(id);
                let (id, generation) = self.table.allocate(phys)?;
                let virt = addr::phys_to_virt(PhysAddr::new(phys));
                unsafe {
                    core::ptr::write(virt.as_mut_ptr::<ObjectData>(), data);
                }
                Ok((ObjectId::new(id), Generation::new(generation as u64)))
            }
            false => Err(KernelError::PoolExhausted),
        }
    }

    pub fn get(&self, id: ObjectId, generation: Generation) -> Result<&ObjectData, KernelError> {
        let phys = self.table.get(id.raw(), generation.raw() as u32)?;
        Ok(unsafe { &*(addr::phys_to_virt(PhysAddr::new(phys)).as_ptr::<ObjectData>()) })
    }

    pub fn get_mut(
        &mut self,
        id: ObjectId,
        generation: Generation,
    ) -> Result<&mut ObjectData, KernelError> {
        let phys = self.table.get(id.raw(), generation.raw() as u32)?;
        Ok(unsafe { &mut *(addr::phys_to_virt(PhysAddr::new(phys)).as_mut_ptr::<ObjectData>()) })
    }

    pub fn inc_ref(&mut self, id: ObjectId, generation: Generation) -> Result<(), KernelError> {
        self.table.inc_ref(id.raw(), generation.raw() as u32)
    }

    pub fn dec_ref(&mut self, id: ObjectId, generation: Generation) -> Option<ObjectData> {
        let outcome = self.table.dec_ref(id.raw(), generation.raw() as u32);
        match outcome {
            DecRefOutcome::Underflow => {
                crate::kprintln!(
                    "[cap] BUG: refcount underflow on object {}. indicates double-free",
                    id.raw()
                );
                None
            }
            DecRefOutcome::Freed(phys) => {
                let data = unsafe {
                    core::ptr::read(addr::phys_to_virt(PhysAddr::new(phys)).as_ptr::<ObjectData>())
                };
                Some(data)
            }
            DecRefOutcome::Alive | DecRefOutcome::Stale => None,
        }
    }

    pub fn revoke(
        &mut self,
        id: ObjectId,
        generation: Generation,
    ) -> Result<(Generation, Option<ObjectData>), KernelError> {
        let (new_gen, old_phys) = self.table.revoke(id.raw(), generation.raw() as u32)?;
        let data = unsafe {
            core::ptr::read(addr::phys_to_virt(PhysAddr::new(old_phys)).as_ptr::<ObjectData>())
        };
        Ok((Generation::new(new_gen as u64), Some(data)))
    }

    pub fn free(
        &mut self,
        id: ObjectId,
        generation: Generation,
    ) -> Result<Option<ObjectData>, KernelError> {
        let (old_phys, _new_gen) = self.table.free(id.raw(), generation.raw() as u32)?;
        match old_phys {
            0 => Ok(None),
            _ => {
                let data = unsafe {
                    core::ptr::read(
                        addr::phys_to_virt(PhysAddr::new(old_phys)).as_ptr::<ObjectData>(),
                    )
                };
                Ok(Some(data))
            }
        }
    }

    pub fn for_each_active_mut(&mut self, mut f: impl FnMut(&mut ObjectData)) {
        self.table.for_each_active(|_id, phys| {
            let data = unsafe {
                &mut *(addr::phys_to_virt(PhysAddr::new(phys)).as_mut_ptr::<ObjectData>())
            };
            f(data);
        });
    }

    pub fn invalidate_process_object(
        &mut self,
        target_pid: crate::types::Pid,
    ) -> crate::static_vec::StaticVec<(ObjectId, Generation), 4> {
        let mut freed = crate::static_vec::StaticVec::<(ObjectId, Generation), 4>::new();
        self.table.free_where(
            |_id, phys, _gen| {
                let data = unsafe {
                    &*(addr::phys_to_virt(PhysAddr::new(phys)).as_ptr::<ObjectData>())
                };
                matches!(data, ObjectData::Process(p) if p.pid == target_pid)
            },
            |id, stale_gen, phys| {
                let _ = freed.push((ObjectId::new(id), Generation::new(stale_gen as u64)));
                let data = unsafe {
                    core::ptr::read(
                        addr::phys_to_virt(PhysAddr::new(phys)).as_ptr::<ObjectData>(),
                    )
                };
                crate::cap::ops::cleanup_object_data(&data);
            },
        );
        freed
    }

    #[allow(dead_code)]
    pub fn table(&self) -> &ObjectTable {
        &self.table
    }

    #[allow(dead_code)]
    pub fn table_mut(&mut self) -> &mut ObjectTable {
        &mut self.table
    }
}

pub static POOL: IrqMutex<ObjectPool, 1> = IrqMutex::new(ObjectPool::new());