pub mod address_space;
#[cfg(not(lancer_test))]
pub mod bootstrap;
pub mod context;
pub mod elf;
pub mod loader;
pub mod manager;
pub mod pid_table;

use crate::mem::typed_addr::Pml4Phys;
use crate::ring::RingIndex;
use crate::types::{BlockedPid, Generation, ObjPhys, Pid, Priority};

pub(crate) const ROOT_CNODE_SIZE_BITS: u8 = 10;
use context::{CpuContext, FpuState, IpcMessage};

pub use lancer_core::process_state::ProcessState;
pub use manager::{PROCESSES, ProcessManager};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BlockedReason {
    Sending(ObjPhys, Generation),
    Receiving(ObjPhys, Generation),
    Calling(ObjPhys, Generation),
    WaitingNotification(ObjPhys, Generation),
}

pub(crate) const PROC_NAME_LEN: usize = 32;

#[repr(C, align(64))]
pub struct SchedEntity {
    pub(crate) pid: Pid,
    pub(crate) generation: Generation,
    pub(crate) state: ProcessState,
    pub(crate) priority: Priority,
    pub(crate) effective_priority: Priority,
    pub(crate) sched_context: Option<(ObjPhys, Generation)>,
    pub(crate) next_ipc: Option<Pid>,
    pub(crate) blocked_reason: Option<BlockedReason>,
    pub(crate) run_cpu: Option<u16>,
    pub(crate) death_notification: Option<(ObjPhys, Generation, u64)>,
    pub(crate) bound_notification: Option<(ObjPhys, Generation)>,
    pub(crate) exec_phys: u64,
    pub(crate) next_free: u32,
    pub(crate) run_prev: u32,
    pub(crate) run_next: u32,
    pub(crate) run_priority: u8,
    pub(crate) timer_prev: u32,
    pub(crate) timer_next: u32,
    pub(crate) timer_slot: u16,
    pub(crate) timer_deadline: u64,
}

#[repr(C, align(64))]
pub struct ExecContext {
    pub(crate) saved_context: CpuContext,
    pub(crate) fpu_state: FpuState,
    pub(crate) pml4_phys: Pml4Phys,
    pub(crate) fs_base: u64,
    pub(crate) context_checksum: u64,
    pub(crate) ipc_message: IpcMessage,
    pub(crate) ipc_badge: u64,
    pub(crate) reply_target: Option<Pid>,
    pub(crate) ring_region_id: Option<(ObjPhys, Generation)>,
    pub(crate) ring_sq_head: RingIndex,
    pub(crate) ring_cq_tail: RingIndex,
    pub(crate) allocated_frames: u16,
    pub(crate) name: [u8; PROC_NAME_LEN],
    pub(crate) name_len: u8,
    pub(crate) root_cnode: Option<(ObjPhys, Generation)>,
    pub(crate) cnode_depth: u8,
    pub(crate) root_guard_value: u64,
    pub(crate) root_guard_bits: u8,
    pub(crate) from_untyped: bool,
}

pub const MAX_FRAMES_PER_PROCESS: u16 = 16384;

const _: () = assert!(core::mem::size_of::<SchedEntity>() <= 256);
const _: () = assert!(core::mem::size_of::<ExecContext>() <= 1536);

use lancer_core::header::NONE_SENTINEL;

impl SchedEntity {
    pub const EMPTY: Self = Self {
        pid: Pid::new(0),
        generation: Generation::new(0),
        state: ProcessState::Free,
        priority: Priority::IDLE,
        effective_priority: Priority::IDLE,
        sched_context: None,
        next_ipc: None,
        blocked_reason: None,
        run_cpu: None,
        death_notification: None,
        bound_notification: None,
        exec_phys: 0,
        next_free: NONE_SENTINEL,
        run_prev: NONE_SENTINEL,
        run_next: NONE_SENTINEL,
        run_priority: 0,
        timer_prev: NONE_SENTINEL,
        timer_next: NONE_SENTINEL,
        timer_slot: 0,
        timer_deadline: 0,
    };

    pub fn state(&self) -> ProcessState {
        self.state
    }

    pub fn transition_to(&mut self, to: ProcessState) -> Result<(), crate::error::KernelError> {
        match to {
            ProcessState::Dead => Err(crate::error::KernelError::BadState),
            _ => self.state.transition(to),
        }
    }

    pub fn zombify_state(&mut self) -> Result<(), crate::error::KernelError> {
        self.state.transition(ProcessState::Zombie)
    }

    pub fn is_runnable(&self) -> bool {
        matches!(self.state, ProcessState::Ready | ProcessState::Running)
    }

    pub fn blocked_reason(&self) -> Option<BlockedReason> {
        self.blocked_reason
    }

    pub fn blocked_proof(&self) -> BlockedPid {
        assert!(
            self.state == ProcessState::Blocked,
            "blocked_proof called on non-blocked process (pid={}, state={:?})",
            self.pid.raw(),
            self.state
        );
        BlockedPid::from_blocked(self.pid, self.generation)
    }

    pub fn block_on(
        &mut self,
        reason: BlockedReason,
    ) -> Result<BlockedPid, crate::error::KernelError> {
        self.state.transition(ProcessState::Blocked)?;
        self.blocked_reason = Some(reason);
        Ok(BlockedPid::from_blocked(self.pid, self.generation))
    }

    pub fn unblock(&mut self, _proof: BlockedPid) -> Result<(), crate::error::KernelError> {
        self.state.transition(ProcessState::Ready)?;
        self.blocked_reason = None;
        Ok(())
    }

    pub fn effective_priority(&self) -> Priority {
        self.effective_priority
    }

    pub fn boost_effective_priority(&mut self, donor: Priority) {
        if donor > self.effective_priority {
            self.effective_priority = donor;
        }
    }

    pub fn reset_effective_priority(&mut self) {
        self.effective_priority = self.priority;
    }

    pub fn sched_context(&self) -> Option<(ObjPhys, Generation)> {
        self.sched_context
    }

    pub fn attach_sched_context(
        &mut self,
        id: ObjPhys,
        generation: Generation,
        priority: Priority,
    ) {
        self.sched_context = Some((id, generation));
        self.priority = priority;
        self.effective_priority = priority;
    }

    pub fn detach_sched_context(&mut self) {
        self.sched_context = None;
    }

    pub fn death_notification(&self) -> Option<(ObjPhys, Generation, u64)> {
        self.death_notification
    }

    pub fn set_death_notification(
        &mut self,
        notif_id: ObjPhys,
        generation: Generation,
        bits: u64,
    ) {
        self.death_notification = Some((notif_id, generation, bits));
    }

    pub fn bound_notification(&self) -> Option<(ObjPhys, Generation)> {
        self.bound_notification
    }

    pub fn bind_notification(&mut self, notif_id: ObjPhys, generation: Generation) {
        self.bound_notification = Some((notif_id, generation));
    }

    pub fn unbind_notification(&mut self) {
        self.bound_notification = None;
    }
}

impl ExecContext {
    pub fn seal_context(&mut self) {
        self.context_checksum = self.saved_context.checksum();
    }

    pub fn verify_context(&self, pid: Pid) {
        let computed = self.saved_context.checksum();
        assert!(
            self.context_checksum == computed,
            "context checksum mismatch for pid {} (stored={:#x}, computed={:#x})",
            pid.raw(),
            self.context_checksum,
            computed
        );
    }

    pub fn set_name(&mut self, src: &[u8]) {
        let copy_len = src.len().min(PROC_NAME_LEN);
        self.name[..copy_len].copy_from_slice(&src[..copy_len]);
        (copy_len..PROC_NAME_LEN).for_each(|i| self.name[i] = 0);
        self.name_len = copy_len as u8;
    }

    pub fn name_str(&self) -> &str {
        let len = self.name_len as usize;
        match len {
            0 => "?",
            _ => core::str::from_utf8(&self.name[..len]).unwrap_or("?"),
        }
    }

    pub fn charge_frames(&mut self, count: u16) -> Result<(), crate::error::KernelError> {
        let new_total = self
            .allocated_frames
            .checked_add(count)
            .filter(|&t| t <= MAX_FRAMES_PER_PROCESS)
            .ok_or(crate::error::KernelError::ResourceExhausted)?;
        self.allocated_frames = new_total;
        Ok(())
    }

    pub fn root_cnode(&self) -> Option<(ObjPhys, Generation)> {
        self.root_cnode
    }

    pub fn cnode_depth(&self) -> u8 {
        self.cnode_depth
    }

    pub fn root_guard_value(&self) -> u64 {
        self.root_guard_value
    }

    pub fn root_guard_bits(&self) -> u8 {
        self.root_guard_bits
    }
}

#[cfg(lancer_test)]
pub fn report_refcount_audit() {
    let ptable = PROCESSES.lock();
    let cap = ptable.capacity();
    let (total, violations) = (0..cap as u32)
        .filter_map(crate::types::Pid::try_new)
        .filter_map(|pid| ptable.exec(pid).map(|e| e.pml4_phys.raw()))
        .fold((0usize, 0usize), |(t, v), pml4| {
            let (pt, pv) = address_space::audit_user_pages(pml4);
            (t + pt, v + pv)
        });
    crate::klog!(
        "test",
        "refcount audit {} PTEs checked, {} violations",
        total,
        violations
    );
}