use lancer_core::intrusive_list::ListLinks;
use lancer_core::run_queue::RunQueueTag;
use lancer_core::slot_map::{SlotEntry, SlotMap};
use lancer_core::timer_wheel::{TimerEntryStorage, TimerTag};

use crate::mem::phys::BitmapFrameAllocator;
use crate::types::{Generation, MAX_PIDS, Pid};

use super::{NONE_SENTINEL, SchedEntity};

const PID_SLOT_SIZE: usize = core::mem::size_of::<SlotEntry<SchedEntity>>();
const PID_LEAF_ENTRIES: usize = 4096 / PID_SLOT_SIZE;
const PID_MAX_LEAVES: usize = MAX_PIDS.div_ceil(PID_LEAF_ENTRIES);

const _: () = assert!(PID_LEAF_ENTRIES >= 1);
const _: () = assert!(PID_LEAF_ENTRIES * PID_SLOT_SIZE <= 4096);
const _: () = assert!(PID_MAX_LEAVES * PID_LEAF_ENTRIES >= MAX_PIDS);

pub struct PidTable {
    slots: SlotMap<SchedEntity, PID_LEAF_ENTRIES, PID_MAX_LEAVES>,
}

const _: () = assert!(MAX_PIDS <= NONE_SENTINEL as usize);

impl PidTable {
    pub const fn empty() -> Self {
        Self {
            slots: SlotMap::new(),
        }
    }

    pub fn init(&mut self) {
        crate::kprintln!(
            "  PidTable: max={} pids, {}/leaf, {}/page waste, slot_entry={}B, sched_entity={}B",
            MAX_PIDS,
            PID_LEAF_ENTRIES,
            4096 - PID_LEAF_ENTRIES * PID_SLOT_SIZE,
            PID_SLOT_SIZE,
            core::mem::size_of::<SchedEntity>(),
        );
    }

    fn try_expand(&mut self, allocator: &BitmapFrameAllocator) -> bool {
        match self.slots.leaf_count() as usize >= PID_MAX_LEAVES {
            true => false,
            false => {
                let frame = match allocator.allocate_contiguous(1) {
                    Some(f) => f,
                    None => return false,
                };
                let virt = crate::mem::addr::phys_to_virt(frame);
                unsafe {
                    core::ptr::write_bytes(virt.as_mut_ptr::<u8>(), 0, 4096);
                    let leaf = virt.as_mut_ptr::<[SlotEntry<SchedEntity>; PID_LEAF_ENTRIES]>();
                    self.slots.expand(leaf);
                }
                true
            }
        }
    }

    pub fn try_reclaim(&mut self) {
        core::iter::from_fn(|| self.slots.try_reclaim_trailing())
            .take(PID_MAX_LEAVES)
            .for_each(|leaf_ptr| {
                let virt = x86_64::VirtAddr::from_ptr(leaf_ptr);
                let phys = crate::mem::addr::virt_to_phys(virt)
                    .expect("PidTable leaf not in HHDM");
                BitmapFrameAllocator::free_frame_by_addr(phys);
            });
    }

    pub fn allocate(
        &mut self,
        sched: SchedEntity,
        allocator: &BitmapFrameAllocator,
    ) -> Option<(Pid, Generation)> {
        let state = sched.state;
        let exec_phys = sched.exec_phys;
        match self.slots.allocate(sched) {
            Some(result) => self.finish_allocate(result),
            None => {
                self.try_expand(allocator);
                let retry = SchedEntity {
                    state,
                    exec_phys,
                    ..SchedEntity::EMPTY
                };
                self.slots.allocate(retry).and_then(|r| self.finish_allocate(r))
            }
        }
    }

    fn finish_allocate(&mut self, (idx, generation): (u32, u32)) -> Option<(Pid, Generation)> {
        let pid = Pid::try_new(idx)?;
        let entry = self.slots.get_mut(idx)?;
        entry.pid = pid;
        entry.generation = Generation::new(generation);
        Some((pid, Generation::new(generation)))
    }

    pub fn free(&mut self, pid: Pid) -> Result<Generation, crate::error::KernelError> {
        let idx = pid.raw();
        let generation = self
            .slots
            .generation_of(idx)
            .ok_or(crate::error::KernelError::InvalidObject)?;
        let freed = self.slots.free(idx, generation)?;
        Ok(freed.generation)
    }

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

    pub fn get(&self, pid: Pid) -> Option<&SchedEntity> {
        self.slots.get(pid.raw())
    }

    pub fn get_mut(&mut self, pid: Pid) -> Option<&mut SchedEntity> {
        self.slots.get_mut(pid.raw())
    }

    fn entry(&self, id: u32) -> &SchedEntity {
        self.slots
            .get(id)
            .expect("PidTable::entry: id not active")
    }

    fn entry_mut(&mut self, id: u32) -> &mut SchedEntity {
        self.slots
            .get_mut(id)
            .expect("PidTable::entry_mut: id not active")
    }

    pub fn capacity(&self) -> u32 {
        self.slots.capacity()
    }

    #[cfg(lancer_test)]
    pub fn leaf_count(&self) -> u16 {
        self.slots.leaf_count()
    }
}

impl ListLinks<RunQueueTag> for PidTable {
    fn link_next(&self, id: u32) -> u32 {
        self.entry(id).run_next
    }
    fn link_prev(&self, id: u32) -> u32 {
        self.entry(id).run_prev
    }
    fn set_link_next(&mut self, id: u32, next: u32) {
        self.entry_mut(id).run_next = next;
    }
    fn set_link_prev(&mut self, id: u32, prev: u32) {
        self.entry_mut(id).run_prev = prev;
    }
}

impl ListLinks<TimerTag> for PidTable {
    fn link_next(&self, id: u32) -> u32 {
        self.entry(id).timer_next
    }
    fn link_prev(&self, id: u32) -> u32 {
        self.entry(id).timer_prev
    }
    fn set_link_next(&mut self, id: u32, next: u32) {
        self.entry_mut(id).timer_next = next;
    }
    fn set_link_prev(&mut self, id: u32, prev: u32) {
        self.entry_mut(id).timer_prev = prev;
    }
}

impl TimerEntryStorage for PidTable {
    fn tw_slot(&self, id: u32) -> u16 {
        self.entry(id).timer_slot
    }
    fn set_tw_slot(&mut self, id: u32, slot: u16) {
        self.entry_mut(id).timer_slot = slot;
    }
    fn tw_deadline(&self, id: u32) -> u64 {
        self.entry(id).timer_deadline
    }
    fn set_tw_deadline(&mut self, id: u32, deadline: u64) {
        self.entry_mut(id).timer_deadline = deadline;
    }
}