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, ObjectId, 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_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_id: ObjectId,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
) -> Result<CapSlot, KernelError> {
    resolve_inner(pool, cnode_id, cnode_gen, address, depth, 0)
}

fn resolve_inner(
    pool: &ObjectPool,
    cnode_id: ObjectId,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    recursion: u8,
) -> Result<CapSlot, KernelError> {
    if recursion >= MAX_RESOLVE_DEPTH {
        return Err(KernelError::InvalidSlot);
    }

    let (slots_phys, size_bits) = {
        let data = pool.get(cnode_id, cnode_gen)?;
        let cnode = data.as_cnode()?;
        (cnode.slots_phys, cnode.size_bits)
    };

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

    let index = extract_index(address, depth, size_bits) as usize;
    let remaining = 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.object_id(),
                cap.generation(),
                address,
                remaining,
                recursion + 1,
            ),
            _ => Err(KernelError::InvalidSlot),
        },
    }
}

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

fn resolve_slot_ptr_inner(
    pool: &ObjectPool,
    cnode_id: ObjectId,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    recursion: u8,
) -> Result<*mut CapSlot, KernelError> {
    if recursion >= MAX_RESOLVE_DEPTH {
        return Err(KernelError::InvalidSlot);
    }

    let (slots_phys, size_bits) = {
        let data = pool.get(cnode_id, cnode_gen)?;
        let cnode = data.as_cnode()?;
        (cnode.slots_phys, cnode.size_bits)
    };

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

    let index = extract_index(address, depth, size_bits) as usize;
    let remaining = 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.object_id(),
                        cap.generation(),
                        address,
                        remaining,
                        recursion + 1,
                    )
                }
                _ => Err(KernelError::InvalidSlot),
            }
        }
    }
}

pub fn resolve_and_insert(
    pool: &ObjectPool,
    cnode_id: ObjectId,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    cap: CapRef,
) -> Result<(), KernelError> {
    let ptr = resolve_slot_ptr(pool, cnode_id, cnode_gen, address, depth)?;
    let slot = unsafe { &mut *ptr };
    match slot {
        CapSlot::Active(_) => Err(KernelError::SlotOccupied),
        CapSlot::Empty => {
            *slot = CapSlot::Active(cap);
            Ok(())
        }
    }
}

pub fn resolve_and_validate(
    pool: &ObjectPool,
    cnode_id: ObjectId,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    expected_tag: ObjectTag,
    required_rights: Rights,
) -> Result<CapRef, KernelError> {
    match resolve(pool, cnode_id, cnode_gen, address, depth)? {
        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)
        }
    }
}

pub fn resolve_and_read(
    pool: &ObjectPool,
    cnode_id: ObjectId,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
) -> Result<CapRef, KernelError> {
    match resolve(pool, cnode_id, cnode_gen, address, depth)? {
        CapSlot::Empty => Err(KernelError::SlotEmpty),
        CapSlot::Active(cap) => Ok(cap),
    }
}

pub fn resolve_and_clear(
    pool: &ObjectPool,
    cnode_id: ObjectId,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
) -> Result<CapRef, KernelError> {
    let ptr = resolve_slot_ptr(pool, cnode_id, cnode_gen, address, depth)?;
    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_id: ObjectId,
    cnode_gen: Generation,
    mut f: impl FnMut(&mut CapSlot),
) -> Result<(), KernelError> {
    let (slots_phys, size_bits) = {
        let data = pool.get(cnode_id, cnode_gen)?;
        let cnode = data.as_cnode()?;
        (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) };
        f(slot);
    });
    Ok(())
}

pub fn cnode_coords(
    pid: Pid,
    ptable: &ProcessManager,
) -> Result<(ObjectId, Generation, u8), KernelError> {
    let proc = ptable.get(pid).ok_or(KernelError::InvalidObject)?;
    let (cnode_id, cnode_gen) = proc.root_cnode().ok_or(KernelError::InvalidObject)?;
    Ok((cnode_id, cnode_gen, proc.cnode_depth()))
}

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

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

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

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

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

#[cfg(lancer_test)]
fn drain_cnode_tree_inner(
    pool: &mut ObjectPool,
    cnode_id: ObjectId,
    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 data = pool.get(cnode_id, cnode_gen)?;
        let cnode = data.as_cnode()?;
        (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.object_id(),
                        cap.generation(),
                        callback,
                        recursion + 1,
                    );
                }
                callback(cap, pool);
            }
            CapSlot::Empty => {}
        }
    });
    Ok(())
}

pub fn drain_cnode_data(
    cnode: &super::object::CNodeData,
    pool: &mut ObjectPool,
    callback: &mut impl FnMut(CapRef, &mut ObjectPool),
) {
    drain_cnode_data_inner(cnode, pool, callback, 0);
}

fn drain_cnode_data_inner(
    cnode: &super::object::CNodeData,
    pool: &mut ObjectPool,
    callback: &mut impl FnMut(CapRef, &mut ObjectPool),
    recursion: u8,
) {
    if recursion >= MAX_RESOLVE_DEPTH {
        return;
    }

    let slot_count = 1usize << cnode.size_bits;
    (0..slot_count).for_each(|i| {
        let slot = unsafe { &mut *slot_ptr(cnode.slots_phys, i) };
        match slot {
            CapSlot::Active(cap) => {
                let cap = *cap;
                *slot = CapSlot::Empty;
                match cap.tag() == ObjectTag::CNode {
                    true => match pool.dec_ref(cap.object_id(), cap.generation()) {
                        Some(super::object::ObjectData::CNode(ref nested)) => {
                            drain_cnode_data_inner(nested, pool, callback, recursion + 1);
                            destroy_cnode(nested, &BitmapFrameAllocator);
                        }
                        Some(ref data) => {
                            super::ops::cleanup_object_data(data);
                        }
                        None => {}
                    },
                    false => callback(cap, pool),
                }
            }
            CapSlot::Empty => {}
        }
    });
}

pub fn invalidate_stale_in_cnode(
    pool: &ObjectPool,
    cnode_id: ObjectId,
    cnode_gen: Generation,
    target_oid: ObjectId,
    target_gen: Generation,
) -> Result<(), KernelError> {
    invalidate_stale_inner(pool, cnode_id, cnode_gen, target_oid, target_gen, 0)
}

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

    let (slots_phys, size_bits) = {
        let data = pool.get(cnode_id, cnode_gen)?;
        let cnode = data.as_cnode()?;
        (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.object_id(),
                        cap.generation(),
                        target_oid,
                        target_gen,
                        recursion + 1,
                    );
                }
                if cap.object_id() == target_oid && cap.generation() == target_gen {
                    *slot = CapSlot::Empty;
                }
            }
            CapSlot::Empty => {}
        }
    });
    Ok(())
}