use lancer_core::header::{KernelObjectHeader, NONE_SENTINEL};
use lancer_core::object_layout::{
    CNodeObject, EndpointObject, FrameObject, IrqHandlerObject, KernelObject, NotificationObject,
    SchedContextObject,
};
use lancer_core::object_tag::ObjectTag;

use super::object::PidQueue;
use super::pool::POOL;
use super::table::{CapRef, Rights};
use crate::error::KernelError;
use crate::mem::addr;
use crate::types::{Generation, ObjPhys, Pid};
use x86_64::PhysAddr;

pub fn cleanup_by_tag(tag: ObjectTag, phys: u64) {
    match tag {
        ObjectTag::IrqHandler => {
            let obj =
                unsafe { &*(addr::phys_to_virt(PhysAddr::new(phys)).as_ptr::<IrqHandlerObject>()) };
            match crate::arch::idt::IrqVector::try_new(obj.vector) {
                Some(vec) => crate::irq::unbind_by_vector(vec),
                None => {}
            }
        }
        ObjectTag::CNode => {
            let obj =
                unsafe { &*(addr::phys_to_virt(PhysAddr::new(phys)).as_ptr::<CNodeObject>()) };
            let cnode_data = super::object::CNodeData {
                slots_phys: PhysAddr::new(obj.slots_phys),
                size_bits: obj.size_bits,
                frame_count: obj.frame_count,
            };
            super::cnode::destroy_cnode(&cnode_data, &crate::mem::phys::BitmapFrameAllocator);
            super::kernel_objects::free_slot(phys);
        }
        ObjectTag::Frame => {
            let obj =
                unsafe { &*(addr::phys_to_virt(PhysAddr::new(phys)).as_ptr::<FrameObject>()) };
            let idx = obj.frame_table_idx;
            match idx != NONE_SENTINEL as u16 {
                true => {
                    let mut ft = super::frame_table::FRAME_TABLE.lock();
                    ft.get(idx).iommu_mapping().inspect(|mapping| {
                        crate::iommu::unmap_frame_iommu(mapping.bus, mapping.devfn, mapping.iova);
                    });
                    ft.clear(idx);
                }
                false => {}
            }
            super::kernel_objects::free_slot(phys);
        }
        ObjectTag::Process => {
            super::kernel_objects::free_slot(phys);
        }
        ObjectTag::VRegion => {
            super::kernel_objects::free_slot(phys);
        }
        ObjectTag::Endpoint
        | ObjectTag::Notification
        | ObjectTag::SchedContext
        | ObjectTag::Framebuffer
        | ObjectTag::PciDevice
        | ObjectTag::Untyped => {}
    }
}

pub fn cleanup_by_tag_with_ptable(
    tag: ObjectTag,
    phys: u64,
    ptable: &mut crate::proc::ProcessManager,
) {
    match tag {
        ObjectTag::Frame | ObjectTag::VRegion => {}
        _ => cleanup_by_tag(tag, phys),
    }
    match tag {
        ObjectTag::Endpoint => {
            let obj =
                unsafe { &*(addr::phys_to_virt(PhysAddr::new(phys)).as_ptr::<EndpointObject>()) };
            let senders = PidQueue::from_repr_c(obj.sender_head, obj.sender_tail, obj.sender_len);
            let receivers =
                PidQueue::from_repr_c(obj.receiver_head, obj.receiver_tail, obj.receiver_len);
            unblock_queue(&senders, ptable);
            unblock_queue(&receivers, ptable);
        }
        ObjectTag::Notification => {
            let obj = unsafe {
                &*(addr::phys_to_virt(PhysAddr::new(phys)).as_ptr::<NotificationObject>())
            };
            (0..obj.waiter_count as usize)
                .filter_map(|i| Pid::try_new(obj.waiters[i]))
                .for_each(|pid| {
                    let proof = ptable[pid].blocked_proof();
                    match ptable.unblock_and_enqueue(pid, proof) {
                        Ok(()) => {
                            let exec = ptable.exec_mut(pid).unwrap();
                            exec.saved_context.rax =
                                crate::error::KernelError::InvalidObject.to_errno() as u64;
                            exec.seal_context();
                        }
                        Err(e) => {
                            crate::kprintln!(
                                "[cap] BUG: notification cleanup failed to unblock pid {}: {:?}",
                                pid.raw(),
                                e
                            );
                        }
                    }
                });
        }
        ObjectTag::SchedContext => {
            let obj = unsafe {
                &*(addr::phys_to_virt(PhysAddr::new(phys)).as_ptr::<SchedContextObject>())
            };
            Pid::try_new(obj.attached_pid).inspect(|&pid| {
                if let Some(sched) = ptable.get_mut(pid) {
                    sched.detach_sched_context();
                }
            });
        }
        ObjectTag::Frame => {
            let obj =
                unsafe { &*(addr::phys_to_virt(PhysAddr::new(phys)).as_ptr::<FrameObject>()) };
            let frame_phys = PhysAddr::new(obj.phys_addr);
            let owned_by_untyped = obj.header.parent_untyped != NONE_SENTINEL;
            let idx = obj.frame_table_idx;
            match idx != NONE_SENTINEL as u16 {
                true => {
                    let mut ft = super::frame_table::FRAME_TABLE.lock();
                    let entry = ft.get_mut(idx);
                    entry.for_each_mapping(|pid, vaddr| {
                        if let Some(exec) = ptable.exec(pid) {
                            let _ =
                                crate::proc::address_space::unmap_user_page(exec.pml4_phys, vaddr);
                            match crate::mem::refcount::decrement(frame_phys) {
                                Ok(0) if !owned_by_untyped => {
                                    crate::mem::phys::BitmapFrameAllocator::free_frame_by_addr(
                                        frame_phys,
                                    )
                                }
                                Ok(_) => {}
                                Err(e) => crate::kprintln!(
                                    "[cap] frame refcount decrement failed: {:#x} {:?}",
                                    frame_phys.as_u64(),
                                    e
                                ),
                            }
                        }
                    });
                    entry.iommu_mapping().inspect(|mapping| {
                        crate::iommu::unmap_frame_iommu(mapping.bus, mapping.devfn, mapping.iova);
                    });
                    ft.clear(idx);
                }
                false => {}
            }
            super::kernel_objects::free_slot(phys);
        }
        ObjectTag::VRegion => {
            let obj = unsafe {
                &*(addr::phys_to_virt(PhysAddr::new(phys))
                    .as_ptr::<lancer_core::object_layout::VRegionObject>())
            };
            let page_count = obj.page_count;
            unmap_vregion_mapping(obj.owner_pid, obj.owner_vaddr, page_count, ptable);
            unmap_vregion_mapping(obj.child_pid, obj.child_vaddr, page_count, ptable);
            super::kernel_objects::free_slot(phys);
        }
        ObjectTag::Process
        | ObjectTag::Framebuffer
        | ObjectTag::PciDevice
        | ObjectTag::Untyped
        | ObjectTag::CNode
        | ObjectTag::IrqHandler => {}
    }
}

pub(crate) fn unmap_region(
    pml4_phys: crate::mem::typed_addr::Pml4Phys,
    vaddr_base: u64,
    page_count: u16,
) {
    (0..page_count as u64).for_each(|i| {
        let vaddr = x86_64::VirtAddr::new(vaddr_base + i * 4096);
        let _ = crate::proc::address_space::unmap_user_page(pml4_phys, vaddr);
    });
}

fn unmap_vregion_mapping(
    pid_raw: u32,
    vaddr_base: u64,
    page_count: u16,
    ptable: &mut crate::proc::ProcessManager,
) {
    Pid::try_new(pid_raw)
        .and_then(|pid| ptable.exec(pid).map(|e| e.pml4_phys))
        .inspect(|&pml4_phys| {
            unmap_region(pml4_phys, vaddr_base, page_count);
        });
}

fn unblock_queue(queue: &super::object::PidQueue, ptable: &mut crate::proc::ProcessManager) {
    let mut cursor = queue.head;
    let mut steps = 0u32;
    let max = crate::types::MAX_PIDS as u32;
    core::iter::from_fn(|| {
        cursor.filter(|_| steps < max).inspect(|&pid| {
            steps += 1;
            cursor = ptable[pid].next_ipc;
            ptable[pid].next_ipc = None;
            let proof = ptable[pid].blocked_proof();
            match ptable.unblock_and_enqueue(pid, proof) {
                Ok(()) => {
                    let exec = ptable.exec_mut(pid).unwrap();
                    exec.saved_context.rax =
                        crate::error::KernelError::InvalidObject.to_errno() as u64;
                    exec.seal_context();
                }
                Err(e) => {
                    crate::kprintln!(
                        "[cap] BUG: unblock_queue failed to unblock pid {}: {:?}",
                        pid.raw(),
                        e
                    );
                }
            }
            ptable.clear_reply_targets_for(pid);
        })
    })
    .count();
}

pub fn resolve_process_cap(
    cap: &CapRef,
    pool: &super::pool::ObjectPool,
) -> Result<crate::types::Pid, KernelError> {
    let obj = pool
        .read_as::<lancer_core::object_layout::ProcessObject>(cap.phys(), cap.generation())?;
    crate::types::Pid::try_new(obj.pid).ok_or(KernelError::InvalidObject)
}

#[allow(clippy::too_many_arguments)]
pub fn create_via_cnode(
    pool: &mut super::pool::ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
    tag: ObjectTag,
) -> Result<ObjPhys, KernelError> {
    let obj_phys = match tag {
        ObjectTag::Endpoint => {
            let phys = super::kernel_objects::alloc_slot().ok_or(KernelError::PoolExhausted)?;
            let header = KernelObjectHeader::new(ObjectTag::Endpoint, 0, 64);
            super::kernel_objects::write_at(phys, EndpointObject::init_default(header));
            phys
        }
        ObjectTag::Notification => {
            let phys = super::kernel_objects::alloc_slot().ok_or(KernelError::PoolExhausted)?;
            let header = KernelObjectHeader::new(ObjectTag::Notification, 0, 64);
            super::kernel_objects::write_at(phys, NotificationObject::init_default(header));
            phys
        }
        ObjectTag::Process
        | ObjectTag::SchedContext
        | ObjectTag::IrqHandler
        | ObjectTag::Framebuffer
        | ObjectTag::PciDevice
        | ObjectTag::CNode
        | ObjectTag::Untyped
        | ObjectTag::Frame
        | ObjectTag::VRegion => {
            return Err(KernelError::InvalidType);
        }
    };

    let (registered_phys, obj_gen) = match pool.register_object(obj_phys, tag) {
        Ok(pair) => pair,
        Err(e) => {
            super::kernel_objects::free_slot(obj_phys);
            return Err(e);
        }
    };
    let cap = CapRef::new(tag, registered_phys, Rights::ALL, obj_gen);
    match super::cnode::resolve_and_insert(
        pool,
        cnode_phys,
        cnode_gen,
        address,
        depth,
        guard_value,
        guard_bits,
        cap,
    ) {
        Ok(()) => Ok(registered_phys),
        Err(e) => {
            let _ = pool.free_phys(registered_phys, obj_gen);
            Err(e)
        }
    }
}

#[allow(clippy::too_many_arguments)]
pub fn derive_via_cnode(
    pool: &mut super::pool::ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    src_addr: u64,
    dest_addr: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
    rights_mask: Rights,
) -> Result<(), KernelError> {
    let src = super::cnode::resolve_and_read(
        pool,
        cnode_phys,
        cnode_gen,
        src_addr,
        depth,
        guard_value,
        guard_bits,
    )?;
    if !src.rights().contains(Rights::GRANT) {
        return Err(KernelError::InsufficientRights);
    }
    pool.inc_ref(src.phys(), src.generation())?;
    let derived = src.with_rights(src.rights() & rights_mask);
    match super::cnode::resolve_and_insert(
        pool,
        cnode_phys,
        cnode_gen,
        dest_addr,
        depth,
        guard_value,
        guard_bits,
        derived,
    ) {
        Ok(()) => Ok(()),
        Err(e) => {
            pool.dec_ref_phys(src.phys(), src.generation());
            Err(e)
        }
    }
}

pub fn identify_via_cnode(
    pool: &super::pool::ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
) -> Result<(ObjectTag, Rights), KernelError> {
    let cap = super::cnode::resolve_and_read(
        pool,
        cnode_phys,
        cnode_gen,
        address,
        depth,
        guard_value,
        guard_bits,
    )?;
    match pool.get_tag(cap.phys(), cap.generation()) {
        Ok(_) => Ok((cap.tag(), cap.rights())),
        Err(KernelError::StaleGeneration) => {
            let _ = super::cnode::resolve_and_clear(
                pool,
                cnode_phys,
                cnode_gen,
                address,
                depth,
                guard_value,
                guard_bits,
            );
            Err(KernelError::StaleGeneration)
        }
        Err(e) => Err(e),
    }
}

#[allow(clippy::too_many_arguments)]
pub fn insert_phys_cap_via_cnode(
    pool: &mut super::pool::ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
    tag: ObjectTag,
    obj_phys: u64,
    rights: Rights,
) -> Result<ObjPhys, KernelError> {
    let (registered_phys, obj_gen) = match pool.register_object(obj_phys, tag) {
        Ok(pair) => pair,
        Err(e) => {
            super::kernel_objects::free_slot(obj_phys);
            return Err(e);
        }
    };
    if tag == ObjectTag::Frame {
        let mut ft = super::frame_table::FRAME_TABLE.lock();
        if let Some(idx) = ft.alloc_idx() {
            let frame_obj = unsafe {
                &mut *(crate::mem::addr::phys_to_virt(x86_64::PhysAddr::new(
                    registered_phys.raw(),
                ))
                .as_mut_ptr::<lancer_core::object_layout::FrameObject>())
            };
            frame_obj.frame_table_idx = idx;
        }
    }
    let cap = CapRef::new(tag, registered_phys, rights, obj_gen);
    match super::cnode::resolve_and_insert(
        pool,
        cnode_phys,
        cnode_gen,
        address,
        depth,
        guard_value,
        guard_bits,
        cap,
    ) {
        Ok(()) => Ok(registered_phys),
        Err(e) => {
            let _ = pool.free_phys(registered_phys, obj_gen);
            Err(e)
        }
    }
}

pub fn revoke_via_cnode(
    pid: crate::types::Pid,
    address: u64,
    ptable: &mut crate::proc::ProcessManager,
) -> Result<(), KernelError> {
    let (cnode_phys, cnode_gen, depth, guard_value, guard_bits) =
        super::cnode::cnode_coords(pid, ptable)?;
    let cap_snapshot = {
        let pool = POOL.lock();
        super::cnode::resolve_and_read(
            &pool,
            cnode_phys,
            cnode_gen,
            address,
            depth,
            guard_value,
            guard_bits,
        )?
    };
    if !cap_snapshot.rights().contains(Rights::REVOKE) {
        return Err(KernelError::InsufficientRights);
    }
    let stale_phys = cap_snapshot.phys();
    let stale_gen = cap_snapshot.generation();

    let is_untyped = cap_snapshot.tag() == ObjectTag::Untyped;

    match is_untyped {
        true => {
            let mut pool = POOL.lock();
            super::derivation::destroy_children(&mut pool, ptable, stale_phys, stale_gen)?;
            Ok(())
        }
        false => {
            {
                let mut pool = POOL.lock();
                super::derivation::unlink_child(&mut pool, stale_phys);
            }
            let (_new_gen, old_phys_tag) = POOL.lock().revoke_phys(stale_phys, stale_gen)?;
            {
                let pool = POOL.lock();
                let _ = super::cnode::resolve_and_clear(
                    &pool,
                    cnode_phys,
                    cnode_gen,
                    address,
                    depth,
                    guard_value,
                    guard_bits,
                );
                invalidate_stale_caps_via_cnode(ptable, &pool, stale_phys, stale_gen);
            }
            old_phys_tag.inspect(|&(phys, tag)| {
                cleanup_by_tag_with_ptable(tag, phys, ptable);
            });
            Ok(())
        }
    }
}

pub fn invalidate_stale_caps_via_cnode(
    ptable: &crate::proc::ProcessManager,
    pool: &super::pool::ObjectPool,
    phys: ObjPhys,
    stale_gen: Generation,
) {
    let cap = ptable.capacity();
    (0..cap as u32)
        .filter_map(crate::types::Pid::try_new)
        .filter_map(|pid| ptable.exec(pid).and_then(|e| e.root_cnode()))
        .for_each(|(cphys, cgen)| {
            let _ = super::cnode::invalidate_stale_in_cnode(pool, cphys, cgen, phys, stale_gen);
        });
}