use lancer_core::arena::{ArenaEntry, SlabArena};
use x86_64::structures::paging::page_table::PageTableFlags;
use x86_64::structures::paging::{Mapper, OffsetPageTable, Page, Size4KiB};
use x86_64::{PhysAddr, VirtAddr};

use crate::mem::addr;
use crate::mem::phys::BitmapFrameAllocator;
use crate::sync::IrqMutex;
use crate::types::{BlockedPid, CreatedPid, Generation, MAX_PIDS, Pid, Priority};

use super::address_space::{self, Pml4ReleaseResult};
use super::context::{CpuContext, FpuState, IpcMessage};
use super::{
    Process, ProcessState, GUARD_PAGE_SIZE, KERNEL_STACK_PAGES, KERNEL_STACK_SIZE, PROC_NAME_LEN,
    STACK_CANARY, STACK_PAINT_BYTE,
};
use crate::mem::typed_addr::Pml4Phys;
use crate::ring::RingIndex;

const SLAB_FRAMES: usize = 128;
const STACK_TOTAL_FRAMES: usize = KERNEL_STACK_PAGES + 1;

fn try_unmap_guard_page(mapper: &mut OffsetPageTable, target_virt: VirtAddr) -> bool {
    let page_4k: Page<Size4KiB> = Page::containing_address(target_virt);
    match mapper.unmap(page_4k) {
        Ok((_, flush)) => {
            flush.flush();
            true
        }
        Err(x86_64::structures::paging::mapper::UnmapError::ParentEntryHugePage) => false,
        Err(_) => false,
    }
}

fn remap_guard_and_free_stack(stack_base_phys: PhysAddr, allocator: &mut BitmapFrameAllocator) {
    let guard_virt = addr::phys_to_virt(stack_base_phys);
    let hhdm_offset = addr::hhdm_offset();
    let mut mapper = unsafe { crate::arch::paging::init(hhdm_offset) };
    let guard_page: Page<Size4KiB> = Page::containing_address(guard_virt);
    let frame = x86_64::structures::paging::PhysFrame::containing_address(stack_base_phys);
    let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE | PageTableFlags::NO_EXECUTE;
    match unsafe { mapper.map_to(guard_page, frame, flags, &mut *allocator) } {
        Ok(flush) => flush.flush(),
        Err(_) => {}
    }
    (0..STACK_TOTAL_FRAMES).for_each(|i| {
        BitmapFrameAllocator::free_frame_by_addr(PhysAddr::new(
            stack_base_phys.as_u64() + (i as u64) * 4096,
        ));
    });
}


pub struct ProcessManager {
    arena: SlabArena<Process>,
    slab_base_phys: u64,
}

impl ProcessManager {
    pub const fn empty() -> Self {
        Self {
            arena: SlabArena::new_null(),
            slab_base_phys: 0,
        }
    }

    pub fn init(&mut self, allocator: &mut BitmapFrameAllocator) {
        let entry_size = core::mem::size_of::<ArenaEntry<Process>>();
        let total_bytes = SLAB_FRAMES * 4096;
        let cap = total_bytes / entry_size;
        let cap = cap.min(MAX_PIDS) as u16;

        let base_frame = allocator
            .allocate_contiguous(SLAB_FRAMES)
            .expect("ProcessManager slab allocation failed");

        self.slab_base_phys = base_frame.as_u64();
        let base_virt = addr::phys_to_virt(base_frame);
        let ptr = base_virt.as_mut_ptr::<ArenaEntry<Process>>();

        unsafe {
            core::ptr::write_bytes(ptr as *mut u8, 0, total_bytes);
            self.arena.init_from_raw_parts(ptr, cap);
        }

        crate::kprintln!(
            "  ProcessManager: {} slots ({} frames, entry={}B)",
            cap,
            SLAB_FRAMES,
            entry_size,
        );
    }

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

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

    pub fn allocate(&mut self, allocator: &mut BitmapFrameAllocator) -> Option<CreatedPid> {
        let pml4_phys = Pml4Phys::from_create(address_space::create_user_pml4(allocator)?);

        let stack_base_phys = match allocator.allocate_contiguous(STACK_TOTAL_FRAMES) {
            Some(phys) => phys,
            None => {
                address_space::teardown_user_space(pml4_phys.raw(), allocator);
                return None;
            }
        };

        let stack_base_virt = addr::phys_to_virt(stack_base_phys);
        let hhdm_offset = addr::hhdm_offset();
        let mut mapper = unsafe { crate::arch::paging::init(hhdm_offset) };
        let _guard_unmapped = try_unmap_guard_page(&mut mapper, stack_base_virt);

        let data_base_virt = VirtAddr::new(stack_base_virt.as_u64() + GUARD_PAGE_SIZE as u64);
        unsafe {
            let data_ptr = data_base_virt.as_mut_ptr::<u8>();
            core::ptr::write_bytes(data_ptr, STACK_PAINT_BYTE, KERNEL_STACK_SIZE);
            core::ptr::copy_nonoverlapping(
                STACK_CANARY.to_le_bytes().as_ptr(),
                data_ptr,
                8,
            );
        }

        let kernel_stack_top =
            VirtAddr::new(stack_base_virt.as_u64() + (STACK_TOTAL_FRAMES as u64) * 4096);

        let process = Process {
            pid: Pid::new(0),
            generation: Generation::new(0),
            state: ProcessState::Created,
            pml4_phys,
            kernel_stack_top,
            saved_context: CpuContext::zero(),
            fpu_state: FpuState::default_init(),
            priority: Priority::IDLE,
            next_ipc: None,
            blocked_reason: None,
            ipc_message: IpcMessage::zero(),
            ipc_badge: 0,
            reply_target: None,
            sched_context: None,
            death_notification: None,
            bound_notification: None,
            effective_priority: Priority::IDLE,
            ring_region_id: None,
            ring_sq_head: RingIndex::new(0),
            ring_cq_tail: RingIndex::new(0),
            allocated_frames: 0,
            run_cpu: None,
            context_checksum: 0,
            name: [0u8; PROC_NAME_LEN],
            name_len: 0,
            stack_phys_base: stack_base_phys,
            fs_base: 0,
            root_cnode: None,
            cnode_depth: 0,
        };

        let (idx, generation) = match self.arena.allocate(process) {
            Ok(pair) => pair,
            Err(_) => {
                remap_guard_and_free_stack(stack_base_phys, allocator);
                address_space::teardown_user_space(pml4_phys.raw(), allocator);
                return None;
            }
        };

        let pid = Pid::new(idx);

        if address_space::pml4_ref_create(pml4_phys.raw(), pid).is_err() {
            self.arena
                .free(idx, generation)
                .expect("pml4_ref_create rollback: arena free failed");
            remap_guard_and_free_stack(stack_base_phys, allocator);
            address_space::teardown_user_space(pml4_phys.raw(), allocator);
            return None;
        }

        let proc = self.arena.get_by_index_mut(idx).unwrap();
        proc.pid = pid;
        proc.generation = generation;
        proc.seal_context();

        Some(CreatedPid::trust(pid))
    }

    pub fn allocate_thread(
        &mut self,
        parent_pid: Pid,
        allocator: &mut BitmapFrameAllocator,
    ) -> Option<CreatedPid> {
        let parent = self.get(parent_pid)?;
        let parent_pml4 = parent.pml4_phys;
        let parent_root_cnode = parent.root_cnode;
        let parent_cnode_depth = parent.cnode_depth;

        let stack_base_phys = allocator.allocate_contiguous(STACK_TOTAL_FRAMES)?;

        let stack_base_virt = addr::phys_to_virt(stack_base_phys);
        let hhdm_offset = addr::hhdm_offset();
        let mut mapper = unsafe { crate::arch::paging::init(hhdm_offset) };
        let _guard_unmapped = try_unmap_guard_page(&mut mapper, stack_base_virt);

        let data_base_virt = VirtAddr::new(stack_base_virt.as_u64() + GUARD_PAGE_SIZE as u64);
        unsafe {
            let data_ptr = data_base_virt.as_mut_ptr::<u8>();
            core::ptr::write_bytes(data_ptr, STACK_PAINT_BYTE, KERNEL_STACK_SIZE);
            core::ptr::copy_nonoverlapping(
                STACK_CANARY.to_le_bytes().as_ptr(),
                data_ptr,
                8,
            );
        }

        let kernel_stack_top =
            VirtAddr::new(stack_base_virt.as_u64() + (STACK_TOTAL_FRAMES as u64) * 4096);

        let process = Process {
            pid: Pid::new(0),
            generation: Generation::new(0),
            state: ProcessState::Created,
            pml4_phys: parent_pml4,
            kernel_stack_top,
            saved_context: CpuContext::zero(),
            fpu_state: FpuState::default_init(),
            priority: Priority::IDLE,
            next_ipc: None,
            blocked_reason: None,
            ipc_message: IpcMessage::zero(),
            ipc_badge: 0,
            reply_target: None,
            sched_context: None,
            death_notification: None,
            bound_notification: None,
            effective_priority: Priority::IDLE,
            ring_region_id: None,
            ring_sq_head: RingIndex::new(0),
            ring_cq_tail: RingIndex::new(0),
            allocated_frames: 0,
            run_cpu: None,
            context_checksum: 0,
            name: [0u8; PROC_NAME_LEN],
            name_len: 0,
            stack_phys_base: stack_base_phys,
            fs_base: 0,
            root_cnode: parent_root_cnode,
            cnode_depth: parent_cnode_depth,
        };

        let (idx, generation) = match self.arena.allocate(process) {
            Ok(pair) => pair,
            Err(_) => {
                remap_guard_and_free_stack(stack_base_phys, allocator);
                return None;
            }
        };

        let pid = Pid::new(idx);

        if address_space::pml4_ref_share(parent_pml4.raw()).is_err() {
            self.arena
                .free(idx, generation)
                .expect("allocate_thread rollback: arena free failed");
            remap_guard_and_free_stack(stack_base_phys, allocator);
            return None;
        }

        let proc = self.arena.get_by_index_mut(idx).unwrap();
        proc.pid = pid;
        proc.generation = generation;
        proc.seal_context();

        Some(CreatedPid::trust(pid))
    }

    pub fn get(&self, pid: Pid) -> Option<&Process> {
        self.arena.get_by_index(pid.raw()).inspect(|p| {
            debug_assert!(
                p.state() != ProcessState::Free,
                "occupied arena slot has Free state for pid {}",
                pid.raw()
            );
        })
    }

    pub fn get_mut(&mut self, pid: Pid) -> Option<&mut Process> {
        self.arena.get_by_index_mut(pid.raw()).inspect(|p| {
            debug_assert!(
                p.state() != ProcessState::Free,
                "occupied arena slot has Free state for pid {}",
                pid.raw()
            );
        })
    }

    pub fn as_created(&self, pid: Pid) -> Option<CreatedPid> {
        self.get(pid)
            .filter(|p| p.state() == ProcessState::Created)
            .map(|_| CreatedPid::trust(pid))
    }

    pub fn start(&mut self, created: CreatedPid) -> Result<(), crate::error::KernelError> {
        self[created.pid()].transition_to(ProcessState::Ready)
    }

    pub fn clear_reply_targets_for(&mut self, target: Pid) {
        self.arena.for_each_active_mut(|proc| {
            if proc.reply_target == Some(target) {
                proc.reply_target = None;
            }
        });
    }

    pub fn zombify(&mut self, pid: Pid) -> bool {
        let proc = match self.arena.get_by_index_mut(pid.raw()) {
            Some(p) => p,
            None => return false,
        };

        if matches!(proc.state(), ProcessState::Free | ProcessState::Zombie) {
            return false;
        }

        if proc.zombify_state().is_err() {
            crate::kprintln!(
                "[proc] BUG: pid {} failed -> Zombie (state={:?})",
                pid.raw(),
                proc.state()
            );
            return false;
        }

        let mut pool = crate::cap::pool::POOL.lock();
        crate::ipc::endpoint::remove_from_queues(pid, &mut pool, self);

        match self[pid].death_notification() {
            Some((notif_id, notif_gen, bits)) => {
                match pool
                    .get_mut(notif_id, notif_gen)
                    .and_then(|d| d.as_notification_mut())
                {
                    Ok(notif) => {
                        let has_waiters = crate::ipc::notification::signal_inner(notif, bits);
                        if has_waiters {
                            crate::ipc::notification::drain_and_wake(notif, self);
                        }
                    }
                    Err(_) => {}
                }
                let _ = crate::ipc::notification::wake_bound_receivers_with_pool(
                    notif_id, notif_gen, bits, self, &mut pool,
                );
            }
            None => {}
        }

        drop(pool);
        self.clear_reply_targets_for(pid);

        true
    }

    pub fn reap(&mut self, pid: Pid, allocator: &mut BitmapFrameAllocator) {
        let proc = match self.arena.get_by_index(pid.raw()) {
            Some(p) => p,
            None => return,
        };

        if proc.state() != ProcessState::Zombie {
            return;
        }

        let pml4 = proc.pml4_phys.raw();
        let stack_base = proc.stack_phys_base;
        let root_cnode = proc.root_cnode;

        self.arena.get_by_index_mut(pid.raw()).unwrap().root_cnode = None;

        {
            let mut pool = crate::cap::pool::POOL.lock();

            let freed_proc_objs = pool.invalidate_process_object(pid);
            freed_proc_objs.iter().for_each(|&(oid, stale_gen)| {
                crate::cap::ops::invalidate_stale_caps_via_cnode(self, &pool, oid, stale_gen);
            });

            if let Some((sc_id, sc_gen)) = self[pid].sched_context()
                && let Ok(sc) = pool
                    .get_mut(sc_id, sc_gen)
                    .and_then(|d| d.as_sched_context_mut())
            {
                sc.attached_pid = None;
            }

            if let Some((cnode_id, cnode_gen)) = root_cnode {
                match pool.dec_ref(cnode_id, cnode_gen) {
                    Some(crate::cap::object::ObjectData::CNode(ref cnode_data)) => {
                        crate::cap::cnode::drain_cnode_data(
                            cnode_data,
                            &mut pool,
                            &mut |cap, pool| {
                                if let Some(ref data) =
                                    pool.dec_ref(cap.object_id(), cap.generation())
                                {
                                    crate::cap::ops::cleanup_object_data_with_ptable(data, self);
                                }
                            },
                        );
                        crate::cap::cnode::destroy_cnode(
                            cnode_data,
                            &crate::mem::phys::BitmapFrameAllocator,
                        );
                    }
                    Some(ref data) => {
                        crate::cap::ops::cleanup_object_data(data);
                    }
                    None => {}
                }
            }
        }

        let entry_gen = self
            .arena
            .generation_of(pid.raw())
            .expect("reap: arena slot empty before free");
        match address_space::pml4_ref_release(pml4) {
            Pml4ReleaseResult::LastRef => address_space::teardown_user_space(pml4, allocator),
            Pml4ReleaseResult::StillShared => {}
        }
        remap_guard_and_free_stack(stack_base, allocator);

        self.arena
            .free(pid.raw(), entry_gen)
            .expect("reap: arena free failed (generation mismatch)");
    }

    pub fn destroy(&mut self, pid: Pid, allocator: &mut BitmapFrameAllocator) -> bool {
        let proc = match self.arena.get_by_index(pid.raw()) {
            Some(p) => p,
            None => return true,
        };

        if proc.state() == ProcessState::Free {
            return true;
        }

        if proc.state() != ProcessState::Zombie && !self.zombify(pid) {
            return false;
        }

        self.reap(pid, allocator);
        true
    }

    pub fn identify_guard_stack(&self, rsp: u64) -> Option<usize> {
        let cap = self.arena.capacity();
        (0..cap as u16).find(|&idx| {
            self.arena
                .get_by_index(idx)
                .filter(|p| p.state() != ProcessState::Free)
                .is_some_and(|p| {
                    let guard_virt = addr::phys_to_virt(p.stack_phys_base).as_u64();
                    rsp >= guard_virt && rsp < guard_virt + GUARD_PAGE_SIZE as u64
                })
        })
        .map(|idx| idx as usize)
    }

    pub fn verify_stack_canary(&self, pid: Pid) {
        let proc = self
            .arena
            .get_by_index(pid.raw())
            .unwrap_or_else(|| panic!("verify_stack_canary: pid {} not in arena", pid.raw()));
        let data_virt =
            addr::phys_to_virt(proc.stack_phys_base).as_u64() + GUARD_PAGE_SIZE as u64;
        let actual = unsafe {
            let ptr = data_virt as *const u8;
            u64::from_le_bytes(core::slice::from_raw_parts(ptr, 8).try_into().unwrap())
        };
        assert!(
            actual == STACK_CANARY,
            "stack canary corrupted for pid {} (expected={:#x}, found={:#x})",
            pid.raw(),
            STACK_CANARY,
            actual
        );
    }

    #[cfg(lancer_test)]
    pub fn stack_high_water_mark(&self, pid: Pid) -> usize {
        let proc = match self.arena.get_by_index(pid.raw()) {
            Some(p) => p,
            None => return 0,
        };
        let data_virt =
            addr::phys_to_virt(proc.stack_phys_base).as_u64() + GUARD_PAGE_SIZE as u64;
        let data = unsafe { core::slice::from_raw_parts(data_virt as *const u8, KERNEL_STACK_SIZE) };
        let unpainted = data.iter().skip(8).take_while(|&&b| b == STACK_PAINT_BYTE).count();
        KERNEL_STACK_SIZE - 8 - unpainted
    }

    #[cfg(lancer_test)]
    pub fn stack_data_mut(&mut self, pid: Pid) -> Option<&'static mut [u8; KERNEL_STACK_SIZE]> {
        let proc = self.arena.get_by_index(pid.raw())?;
        let data_virt =
            addr::phys_to_virt(proc.stack_phys_base).as_u64() + GUARD_PAGE_SIZE as u64;
        Some(unsafe { &mut *(data_virt as *mut [u8; KERNEL_STACK_SIZE]) })
    }
}

impl core::ops::Index<Pid> for ProcessManager {
    type Output = Process;
    fn index(&self, pid: Pid) -> &Process {
        self.arena
            .get_by_index(pid.raw())
            .unwrap_or_else(|| panic!("ProcessManager: invalid pid {}", pid.raw()))
    }
}

impl core::ops::IndexMut<Pid> for ProcessManager {
    fn index_mut(&mut self, pid: Pid) -> &mut Process {
        self.arena
            .get_by_index_mut(pid.raw())
            .unwrap_or_else(|| panic!("ProcessManager: invalid pid {}", pid.raw()))
    }
}

impl core::ops::Index<&BlockedPid> for ProcessManager {
    type Output = Process;
    fn index(&self, bp: &BlockedPid) -> &Process {
        &self[bp.pid()]
    }
}

impl core::ops::IndexMut<&BlockedPid> for ProcessManager {
    fn index_mut(&mut self, bp: &BlockedPid) -> &mut Process {
        &mut self[bp.pid()]
    }
}

impl core::ops::Index<CreatedPid> for ProcessManager {
    type Output = Process;
    fn index(&self, cp: CreatedPid) -> &Process {
        &self[cp.pid()]
    }
}

impl core::ops::IndexMut<CreatedPid> for ProcessManager {
    fn index_mut(&mut self, cp: CreatedPid) -> &mut Process {
        &mut self[cp.pid()]
    }
}

pub static PROCESSES: IrqMutex<ProcessManager, 0> = IrqMutex::new(ProcessManager::empty());