use crate::error::KernelError;
use crate::mem::addr;
use crate::sync::IrqMutex;
use crate::types::{Generation, ObjPhys};
use lancer_core::header::KernelObjectHeader;
use lancer_core::object_layout::KernelObject;
use lancer_core::object_tag::ObjectTag;
use x86_64::PhysAddr;

pub struct ObjectPool {
    generation_counter: u32,
    #[cfg(lancer_test)]
    active_count: u32,
}

impl ObjectPool {
    const fn new() -> Self {
        Self {
            generation_counter: 0,
            #[cfg(lancer_test)]
            active_count: 0,
        }
    }

    pub fn init(&mut self) {
        self.generation_counter = 0;
        #[cfg(lancer_test)]
        {
            self.active_count = 0;
        }
    }

    #[cfg(lancer_test)]
    pub fn active_count(&self) -> u32 {
        self.active_count
    }

    pub fn register_object(
        &mut self,
        phys: u64,
        _tag: ObjectTag,
    ) -> Result<(ObjPhys, Generation), KernelError> {
        debug_assert!(
            phys.is_multiple_of(core::mem::align_of::<KernelObjectHeader>() as u64),
            "register_object: misaligned phys {phys:#x}"
        );
        self.generation_counter = self.generation_counter.wrapping_add(1);
        let generation = self.generation_counter;
        unsafe {
            let hdr_ptr =
                addr::phys_to_virt(PhysAddr::new(phys)).as_mut_ptr::<KernelObjectHeader>();
            (*hdr_ptr).generation = generation;
            (*hdr_ptr).ref_count = 1;
        }
        #[cfg(lancer_test)]
        {
            self.active_count += 1;
        }
        Ok((ObjPhys::new(phys), Generation::new(generation)))
    }

    fn read_header(&self, phys: ObjPhys) -> &'static KernelObjectHeader {
        unsafe {
            &*(addr::phys_to_virt(PhysAddr::new(phys.raw())).as_ptr::<KernelObjectHeader>())
        }
    }

    fn write_header(&mut self, phys: ObjPhys) -> &'static mut KernelObjectHeader {
        unsafe {
            &mut *(addr::phys_to_virt(PhysAddr::new(phys.raw()))
                .as_mut_ptr::<KernelObjectHeader>())
        }
    }

    fn validate_generation(
        &self,
        phys: ObjPhys,
        generation: Generation,
    ) -> Result<(), KernelError> {
        let hdr = self.read_header(phys);
        match hdr.generation != generation.raw() {
            true => Err(KernelError::StaleGeneration),
            false => match hdr.ref_count {
                0 => Err(KernelError::InvalidObject),
                _ => Ok(()),
            },
        }
    }

    pub fn inc_ref(&mut self, phys: ObjPhys, generation: Generation) -> Result<(), KernelError> {
        self.validate_generation(phys, generation)?;
        let hdr = self.write_header(phys);
        hdr.inc_ref().ok_or(KernelError::ResourceExhausted)?;
        Ok(())
    }

    pub fn get_tag(
        &self,
        phys: ObjPhys,
        generation: Generation,
    ) -> Result<ObjectTag, KernelError> {
        self.validate_generation(phys, generation)?;
        let hdr = self.read_header(phys);
        ObjectTag::try_from(hdr.tag).map_err(|_| KernelError::InvalidType)
    }

    pub fn read_as<T: KernelObject + 'static>(
        &self,
        phys: ObjPhys,
        generation: Generation,
    ) -> Result<&T, KernelError> {
        unsafe { super::accessor::read_object::<T>(PhysAddr::new(phys.raw()), generation.raw()) }
    }

    pub fn write_as<T: KernelObject>(
        &mut self,
        phys: ObjPhys,
        generation: Generation,
    ) -> Result<&mut T, KernelError> {
        let ptr = unsafe {
            super::accessor::write_object::<T>(PhysAddr::new(phys.raw()), generation.raw())?
        };
        Ok(unsafe { &mut *ptr })
    }

    pub fn dec_ref_phys(
        &mut self,
        phys: ObjPhys,
        generation: Generation,
    ) -> Option<(u64, ObjectTag)> {
        let hdr = self.read_header(phys);
        match hdr.generation != generation.raw() || hdr.ref_count == 0 {
            true => return None,
            false => {}
        }

        let tag_byte = hdr.tag;
        let tag = ObjectTag::try_from(tag_byte).ok()?;

        let hdr = self.write_header(phys);
        let new_rc = hdr.dec_ref();
        match new_rc {
            0 => {
                hdr.bump_generation();
                #[cfg(lancer_test)]
                {
                    self.active_count = self.active_count.saturating_sub(1);
                }
                Some((phys.raw(), tag))
            }
            _ => None,
        }
    }

    pub fn free_phys(
        &mut self,
        phys: ObjPhys,
        generation: Generation,
    ) -> Result<Option<(u64, ObjectTag)>, KernelError> {
        let hdr = self.read_header(phys);
        match hdr.generation != generation.raw() {
            true => return Err(KernelError::StaleGeneration),
            false => {}
        }
        match hdr.ref_count {
            0 => return Ok(None),
            _ => {}
        }

        let tag = ObjectTag::try_from(hdr.tag).ok().map(|t| (phys.raw(), t));
        let hdr = self.write_header(phys);
        hdr.invalidate();
        #[cfg(lancer_test)]
        {
            self.active_count = self.active_count.saturating_sub(1);
        }
        Ok(tag)
    }

    pub fn revoke_phys(
        &mut self,
        phys: ObjPhys,
        generation: Generation,
    ) -> Result<(Generation, Option<(u64, ObjectTag)>), KernelError> {
        let hdr = self.read_header(phys);
        match hdr.generation != generation.raw() || hdr.ref_count == 0 {
            true => return Err(KernelError::StaleGeneration),
            false => {}
        }

        let tag = ObjectTag::try_from(hdr.tag).ok().map(|t| (phys.raw(), t));
        let hdr = self.write_header(phys);
        hdr.invalidate();
        let new_gen = hdr.generation;
        #[cfg(lancer_test)]
        {
            self.active_count = self.active_count.saturating_sub(1);
        }
        Ok((Generation::new(new_gen), tag))
    }

    pub fn generation_of(&self, phys: ObjPhys) -> Option<u32> {
        let hdr = self.read_header(phys);
        match hdr.ref_count > 0 {
            true => Some(hdr.generation),
            false => None,
        }
    }
}

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