use lancer_core::object_layout::CNodeObject;

use super::object::{CNodeData, ObjectTag};
use super::pool::ObjectPool;
use super::table::{CapRef, CapSlot, Rights};
use crate::error::KernelError;
use crate::mem::addr;
use crate::mem::phys::BitmapFrameAllocator;
use crate::proc::ProcessManager;
use crate::types::{Generation, ObjPhys, Pid};
use x86_64::PhysAddr;

pub use lancer_core::cnode::{MAX_CNODE_BITS, MIN_CNODE_BITS};
use lancer_core::cnode::{
    MAX_RESOLVE_DEPTH, PAGE_SIZE, extract_guard, extract_index, frames_for_cnode,
};

pub fn create_cnode(
    size_bits: u8,
    allocator: &BitmapFrameAllocator,
) -> Result<CNodeData, KernelError> {
    if !(MIN_CNODE_BITS..=MAX_CNODE_BITS).contains(&size_bits) {
        return Err(KernelError::InvalidParameter);
    }
    let frame_count = frames_for_cnode(size_bits);
    let slots_phys = allocator
        .allocate_contiguous(frame_count as usize)
        .ok_or(KernelError::ResourceExhausted)?;

    let slot_count = 1usize << size_bits;
    let base_ptr = addr::phys_to_virt(slots_phys).as_mut_ptr::<CapSlot>();
    (0..slot_count).for_each(|i| unsafe {
        base_ptr.add(i).write(CapSlot::Empty);
    });

    Ok(CNodeData {
        slots_phys,
        size_bits,
        frame_count,
    })
}

pub fn destroy_cnode(cnode: &CNodeData, allocator: &BitmapFrameAllocator) {
    let base_frame_idx = (cnode.slots_phys.as_u64() / PAGE_SIZE as u64) as usize;
    (0..cnode.frame_count as usize).for_each(|i| {
        let frame_phys = PhysAddr::new((base_frame_idx + i) as u64 * PAGE_SIZE as u64);
        crate::mem::addr::zero_frame(frame_phys);
        let frame = unsafe {
            x86_64::structures::paging::PhysFrame::from_start_address_unchecked(frame_phys)
        };
        allocator.deallocate_frame(frame);
    });
}

unsafe fn read_slot(slots_phys: PhysAddr, index: usize) -> CapSlot {
    let base = addr::phys_to_virt(slots_phys).as_ptr::<CapSlot>();
    unsafe { base.add(index).read() }
}

unsafe fn slot_ptr(slots_phys: PhysAddr, index: usize) -> *mut CapSlot {
    let base = addr::phys_to_virt(slots_phys).as_mut_ptr::<CapSlot>();
    unsafe { base.add(index) }
}

pub fn resolve(
    pool: &ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
) -> Result<CapSlot, KernelError> {
    resolve_inner(
        pool,
        cnode_phys,
        cnode_gen,
        address,
        depth,
        guard_value,
        guard_bits,
        0,
    )
}

#[allow(clippy::too_many_arguments)]
fn resolve_inner(
    pool: &ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
    recursion: u8,
) -> Result<CapSlot, KernelError> {
    if recursion >= MAX_RESOLVE_DEPTH {
        return Err(KernelError::InvalidSlot);
    }

    if depth < guard_bits {
        return Err(KernelError::InvalidSlot);
    }

    let extracted = extract_guard(address, depth, guard_bits);
    if extracted != guard_value {
        return Err(KernelError::GuardMismatch);
    }

    let post_guard_depth = depth - guard_bits;

    let (slots_phys, size_bits) = {
        let cnode = pool.read_as::<CNodeObject>(cnode_phys, cnode_gen)?;
        (PhysAddr::new(cnode.slots_phys), cnode.size_bits)
    };

    if post_guard_depth < size_bits {
        return Err(KernelError::InvalidSlot);
    }

    let index = extract_index(address, post_guard_depth, size_bits) as usize;
    let remaining = post_guard_depth - size_bits;
    let slot = unsafe { read_slot(slots_phys, index) };

    match remaining {
        0 => Ok(slot),
        _ => match slot {
            CapSlot::Active(cap) if cap.tag() == ObjectTag::CNode => resolve_inner(
                pool,
                cap.phys(),
                cap.generation(),
                address,
                remaining,
                cap.guard_value(),
                cap.guard_bits(),
                recursion + 1,
            ),
            _ => Err(KernelError::InvalidSlot),
        },
    }
}

fn resolve_slot_ptr(
    pool: &ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
) -> Result<*mut CapSlot, KernelError> {
    resolve_slot_ptr_inner(
        pool,
        cnode_phys,
        cnode_gen,
        address,
        depth,
        guard_value,
        guard_bits,
        0,
    )
}

#[allow(clippy::too_many_arguments)]
fn resolve_slot_ptr_inner(
    pool: &ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
    recursion: u8,
) -> Result<*mut CapSlot, KernelError> {
    if recursion >= MAX_RESOLVE_DEPTH {
        return Err(KernelError::InvalidSlot);
    }

    if depth < guard_bits {
        return Err(KernelError::InvalidSlot);
    }

    let extracted = extract_guard(address, depth, guard_bits);
    if extracted != guard_value {
        return Err(KernelError::GuardMismatch);
    }

    let post_guard_depth = depth - guard_bits;

    let (slots_phys, size_bits) = {
        let cnode = pool.read_as::<CNodeObject>(cnode_phys, cnode_gen)?;
        (PhysAddr::new(cnode.slots_phys), cnode.size_bits)
    };

    if post_guard_depth < size_bits {
        return Err(KernelError::InvalidSlot);
    }

    let index = extract_index(address, post_guard_depth, size_bits) as usize;
    let remaining = post_guard_depth - size_bits;

    match remaining {
        0 => Ok(unsafe { slot_ptr(slots_phys, index) }),
        _ => {
            let slot = unsafe { read_slot(slots_phys, index) };
            match slot {
                CapSlot::Active(cap) if cap.tag() == ObjectTag::CNode => resolve_slot_ptr_inner(
                    pool,
                    cap.phys(),
                    cap.generation(),
                    address,
                    remaining,
                    cap.guard_value(),
                    cap.guard_bits(),
                    recursion + 1,
                ),
                _ => Err(KernelError::InvalidSlot),
            }
        }
    }
}

#[allow(clippy::too_many_arguments)]
pub fn resolve_and_insert(
    pool: &ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
    cap: CapRef,
) -> Result<(), KernelError> {
    let ptr = resolve_slot_ptr(
        pool,
        cnode_phys,
        cnode_gen,
        address,
        depth,
        guard_value,
        guard_bits,
    )?;
    let slot = unsafe { &mut *ptr };
    match slot {
        CapSlot::Active(_) => Err(KernelError::SlotOccupied),
        CapSlot::Empty => {
            *slot = CapSlot::Active(cap);
            Ok(())
        }
    }
}

#[allow(clippy::too_many_arguments)]
pub fn resolve_and_validate(
    pool: &ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
    expected_tag: ObjectTag,
    required_rights: Rights,
) -> Result<CapRef, KernelError> {
    match resolve(
        pool,
        cnode_phys,
        cnode_gen,
        address,
        depth,
        guard_value,
        guard_bits,
    )? {
        CapSlot::Empty => Err(KernelError::SlotEmpty),
        CapSlot::Active(cap) => {
            if cap.tag() != expected_tag {
                return Err(KernelError::InvalidType);
            }
            if !cap.rights().contains(required_rights) {
                return Err(KernelError::InsufficientRights);
            }
            Ok(cap)
        }
    }
}

#[allow(clippy::too_many_arguments)]
pub fn resolve_and_read(
    pool: &ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
) -> Result<CapRef, KernelError> {
    match resolve(
        pool,
        cnode_phys,
        cnode_gen,
        address,
        depth,
        guard_value,
        guard_bits,
    )? {
        CapSlot::Empty => Err(KernelError::SlotEmpty),
        CapSlot::Active(cap) => Ok(cap),
    }
}

pub fn resolve_and_clear(
    pool: &ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
) -> Result<CapRef, KernelError> {
    let ptr = resolve_slot_ptr(
        pool,
        cnode_phys,
        cnode_gen,
        address,
        depth,
        guard_value,
        guard_bits,
    )?;
    let slot = unsafe { &mut *ptr };
    match slot {
        CapSlot::Empty => Err(KernelError::SlotEmpty),
        CapSlot::Active(cap) => {
            let cap = *cap;
            *slot = CapSlot::Empty;
            Ok(cap)
        }
    }
}

pub fn walk_cnode_slots(
    pool: &ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    mut f: impl FnMut(&mut CapSlot),
) -> Result<(), KernelError> {
    walk_cnode_slots_recursive(pool, cnode_phys, cnode_gen, &mut f, 0)
}

fn walk_cnode_slots_recursive(
    pool: &ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    f: &mut impl FnMut(&mut CapSlot),
    depth: u8,
) -> Result<(), KernelError> {
    if depth >= MAX_RESOLVE_DEPTH {
        return Ok(());
    }

    let (slots_phys, size_bits) = {
        let cnode = pool.read_as::<CNodeObject>(cnode_phys, cnode_gen)?;
        (PhysAddr::new(cnode.slots_phys), cnode.size_bits)
    };

    let slot_count = 1usize << size_bits;
    (0..slot_count).for_each(|i| {
        let slot = unsafe { &mut *slot_ptr(slots_phys, i) };
        match slot {
            CapSlot::Active(cap) if cap.tag() == ObjectTag::CNode => {
                let child_phys = cap.phys();
                let child_gen = cap.generation();
                f(slot);
                let _ = walk_cnode_slots_recursive(pool, child_phys, child_gen, f, depth + 1);
            }
            _ => f(slot),
        }
    });
    Ok(())
}

pub fn cnode_coords(
    pid: Pid,
    ptable: &ProcessManager,
) -> Result<(ObjPhys, Generation, u8, u64, u8), KernelError> {
    let _sched = ptable.get(pid).ok_or(KernelError::InvalidObject)?;
    let exec = ptable.exec(pid).ok_or(KernelError::InvalidObject)?;
    let (cnode_phys, cnode_gen) = exec.root_cnode().ok_or(KernelError::InvalidObject)?;
    Ok((
        cnode_phys,
        cnode_gen,
        exec.cnode_depth(),
        exec.root_guard_value(),
        exec.root_guard_bits(),
    ))
}

pub fn resolve_caller_validate(
    pid: Pid,
    address: u64,
    expected_tag: ObjectTag,
    required_rights: Rights,
    ptable: &ProcessManager,
    pool: &ObjectPool,
) -> Result<CapRef, KernelError> {
    let (cnode_phys, cnode_gen, depth, gv, gb) = cnode_coords(pid, ptable)?;
    resolve_and_validate(
        pool,
        cnode_phys,
        cnode_gen,
        address,
        depth,
        gv,
        gb,
        expected_tag,
        required_rights,
    )
}

pub fn resolve_caller_read(
    pid: Pid,
    address: u64,
    ptable: &ProcessManager,
    pool: &ObjectPool,
) -> Result<CapRef, KernelError> {
    let (cnode_phys, cnode_gen, depth, gv, gb) = cnode_coords(pid, ptable)?;
    resolve_and_read(pool, cnode_phys, cnode_gen, address, depth, gv, gb)
}

pub fn resolve_caller_insert(
    pid: Pid,
    address: u64,
    cap: CapRef,
    ptable: &ProcessManager,
    pool: &ObjectPool,
) -> Result<(), KernelError> {
    let (cnode_phys, cnode_gen, depth, gv, gb) = cnode_coords(pid, ptable)?;
    resolve_and_insert(pool, cnode_phys, cnode_gen, address, depth, gv, gb, cap)
}

pub fn resolve_caller_clear(
    pid: Pid,
    address: u64,
    ptable: &ProcessManager,
    pool: &ObjectPool,
) -> Result<CapRef, KernelError> {
    let (cnode_phys, cnode_gen, depth, gv, gb) = cnode_coords(pid, ptable)?;
    resolve_and_clear(pool, cnode_phys, cnode_gen, address, depth, gv, gb)
}

#[cfg(lancer_test)]
pub fn drain_cnode_tree(
    pool: &mut ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    callback: &mut impl FnMut(CapRef, &mut ObjectPool),
) -> Result<(), KernelError> {
    drain_cnode_tree_inner(pool, cnode_phys, cnode_gen, callback, 0)
}

#[cfg(lancer_test)]
fn drain_cnode_tree_inner(
    pool: &mut ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    callback: &mut impl FnMut(CapRef, &mut ObjectPool),
    recursion: u8,
) -> Result<(), KernelError> {
    if recursion >= MAX_RESOLVE_DEPTH {
        return Err(KernelError::InvalidSlot);
    }

    let (slots_phys, size_bits) = {
        let cnode = pool.read_as::<CNodeObject>(cnode_phys, cnode_gen)?;
        (PhysAddr::new(cnode.slots_phys), cnode.size_bits)
    };

    let slot_count = 1usize << size_bits;
    (0..slot_count).for_each(|i| {
        let slot = unsafe { &mut *slot_ptr(slots_phys, i) };
        match slot {
            CapSlot::Active(cap) => {
                let cap = *cap;
                *slot = CapSlot::Empty;
                if cap.tag() == ObjectTag::CNode {
                    let _ = drain_cnode_tree_inner(
                        pool,
                        cap.phys(),
                        cap.generation(),
                        callback,
                        recursion + 1,
                    );
                }
                callback(cap, pool);
            }
            CapSlot::Empty => {}
        }
    });
    Ok(())
}

pub fn drain_cnode_phys(
    slots_phys: PhysAddr,
    size_bits: u8,
    frame_count: u8,
    pool: &mut ObjectPool,
    ptable: &mut crate::proc::ProcessManager,
) {
    drain_cnode_phys_inner(slots_phys, size_bits, pool, ptable, 0);
    let cnode_data = super::object::CNodeData {
        slots_phys,
        size_bits,
        frame_count,
    };
    destroy_cnode(&cnode_data, &BitmapFrameAllocator);
}

fn drain_cnode_phys_inner(
    slots_phys: PhysAddr,
    size_bits: u8,
    pool: &mut ObjectPool,
    ptable: &mut crate::proc::ProcessManager,
    recursion: u8,
) {
    if recursion >= MAX_RESOLVE_DEPTH {
        return;
    }

    let slot_count = 1usize << size_bits;
    (0..slot_count).for_each(|i| {
        let slot = unsafe { &mut *slot_ptr(slots_phys, i) };
        match slot {
            CapSlot::Active(cap) => {
                let cap = *cap;
                *slot = CapSlot::Empty;
                let freed_phys = cap.phys();
                match cap.tag() == ObjectTag::CNode {
                    true => match pool.dec_ref_phys(cap.phys(), cap.generation()) {
                        Some((nested_phys, ObjectTag::CNode)) => {
                            super::derivation::unlink_child(pool, freed_phys);
                            let nested = unsafe {
                                &*(crate::mem::addr::phys_to_virt(PhysAddr::new(nested_phys))
                                    .as_ptr::<CNodeObject>())
                            };
                            drain_cnode_phys_inner(
                                PhysAddr::new(nested.slots_phys),
                                nested.size_bits,
                                pool,
                                ptable,
                                recursion + 1,
                            );
                            let nested_data = super::object::CNodeData {
                                slots_phys: PhysAddr::new(nested.slots_phys),
                                size_bits: nested.size_bits,
                                frame_count: nested.frame_count,
                            };
                            destroy_cnode(&nested_data, &BitmapFrameAllocator);
                            super::kernel_objects::free_slot(nested_phys);
                        }
                        Some((other_phys, other_tag)) => {
                            super::derivation::unlink_child(pool, freed_phys);
                            super::ops::cleanup_by_tag_with_ptable(
                                other_tag,
                                other_phys,
                                ptable,
                            );
                        }
                        None => {}
                    },
                    false => match pool.dec_ref_phys(cap.phys(), cap.generation()) {
                        Some((obj_phys, tag)) => {
                            super::derivation::unlink_child(pool, freed_phys);
                            super::ops::cleanup_by_tag_with_ptable(
                                tag,
                                obj_phys,
                                ptable,
                            );
                        }
                        None => {}
                    },
                }
            }
            CapSlot::Empty => {}
        }
    });
}

pub fn invalidate_stale_in_cnode(
    pool: &ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    target_phys: ObjPhys,
    target_gen: Generation,
) -> Result<(), KernelError> {
    invalidate_stale_inner(pool, cnode_phys, cnode_gen, target_phys, target_gen, 0)
}

fn invalidate_stale_inner(
    pool: &ObjectPool,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    target_phys: ObjPhys,
    target_gen: Generation,
    recursion: u8,
) -> Result<(), KernelError> {
    if recursion >= MAX_RESOLVE_DEPTH {
        return Err(KernelError::InvalidSlot);
    }

    let (slots_phys, size_bits) = {
        let cnode = pool.read_as::<CNodeObject>(cnode_phys, cnode_gen)?;
        (PhysAddr::new(cnode.slots_phys), cnode.size_bits)
    };

    let slot_count = 1usize << size_bits;
    (0..slot_count).for_each(|i| {
        let slot = unsafe { &mut *slot_ptr(slots_phys, i) };
        match slot {
            CapSlot::Active(cap) => {
                if cap.tag() == ObjectTag::CNode {
                    let _ = invalidate_stale_inner(
                        pool,
                        cap.phys(),
                        cap.generation(),
                        target_phys,
                        target_gen,
                        recursion + 1,
                    );
                }
                if cap.phys() == target_phys && cap.generation() == target_gen {
                    *slot = CapSlot::Empty;
                }
            }
            CapSlot::Empty => {}
        }
    });
    Ok(())
}