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

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

const MAX_RETYPE_COUNT: u32 = 64;

fn write_inline<T: KernelObject>(virt: x86_64::VirtAddr, header: KernelObjectHeader) {
    unsafe { core::ptr::write(virt.as_mut_ptr::<T>(), T::init_default(header)) };
}

fn write_repr_c_object(tag: ObjectTag, obj_phys: u64, size_bits: u8) -> Result<u64, KernelError> {
    let header = KernelObjectHeader::new(tag, 0, 64);
    let virt = addr::phys_to_virt(PhysAddr::new(obj_phys));

    match tag {
        ObjectTag::Endpoint => {
            write_inline::<EndpointObject>(virt, header);
            Ok(obj_phys)
        }
        ObjectTag::Notification => {
            write_inline::<NotificationObject>(virt, header);
            Ok(obj_phys)
        }
        ObjectTag::SchedContext => {
            write_inline::<SchedContextObject>(virt, header);
            Ok(obj_phys)
        }
        ObjectTag::IrqHandler => {
            write_inline::<IrqHandlerObject>(virt, header);
            Ok(obj_phys)
        }
        ObjectTag::Frame => {
            let kobj_phys =
                super::kernel_objects::alloc_slot().ok_or(KernelError::PoolExhausted)?;
            let mut obj = FrameObject::init_default(header);
            obj.phys_addr = obj_phys;
            obj.size_bits = match size_bits {
                0 => 12,
                b => b,
            };
            super::kernel_objects::write_at(kobj_phys, obj);
            Ok(kobj_phys)
        }
        ObjectTag::CNode => {
            let frame_count = lancer_core::cnode::frames_for_cnode(size_bits);
            let mut obj = CNodeObject::init_default(header);
            obj.slots_phys = obj_phys;
            obj.size_bits = size_bits;
            obj.frame_count = frame_count;
            unsafe { core::ptr::write(virt.as_mut_ptr::<CNodeObject>(), obj) };
            Ok(obj_phys)
        }
        ObjectTag::Untyped
        | ObjectTag::Process
        | ObjectTag::Framebuffer
        | ObjectTag::PciDevice
        | ObjectTag::VRegion => Ok(obj_phys),
    }
}

fn validate_retype_tag(tag: ObjectTag) -> Result<(), KernelError> {
    match tag {
        ObjectTag::Endpoint
        | ObjectTag::Notification
        | ObjectTag::Frame
        | ObjectTag::CNode
        | ObjectTag::SchedContext
        | ObjectTag::IrqHandler => Ok(()),
        ObjectTag::Untyped
        | ObjectTag::Process
        | ObjectTag::Framebuffer
        | ObjectTag::PciDevice
        | ObjectTag::VRegion => Err(KernelError::InvalidType),
    }
}

struct TcbInitResult {
    pid: crate::types::Pid,
    pml4_phys: crate::mem::typed_addr::Pml4Phys,
    root_cnode: Option<(ObjPhys, Generation)>,
}

enum AddressSpaceMode {
    New,
    Shared { parent_pid: crate::types::Pid },
}

fn init_tcb_memory(
    ptable: &mut crate::proc::ProcessManager,
    phys_base: PhysAddr,
    offset: u64,
    mode: AddressSpaceMode,
) -> Result<TcbInitResult, KernelError> {
    let tcb_phys = PhysAddr::new(phys_base.as_u64() + offset);
    let header_size = lancer_core::untyped::HEADER_SIZE as u64;
    let tcb_offset = (header_size + 63) & !63;

    let exec_phys = PhysAddr::new(tcb_phys.as_u64() + tcb_offset);

    let hhdm = crate::mem::addr::hhdm_offset();
    let exec_virt = x86_64::VirtAddr::new(exec_phys.as_u64() + hhdm);

    match mode {
        AddressSpaceMode::New => {
            let mut allocator = crate::mem::phys::BitmapFrameAllocator;
            let pml4_raw = crate::proc::address_space::create_user_pml4(&mut allocator)
                .ok_or(KernelError::ResourceExhausted)?;
            let pml4_phys = crate::mem::typed_addr::Pml4Phys::from_create(pml4_raw);

            let sched = crate::proc::manager::make_default_sched(
                crate::types::Pid::new(0),
                Generation::new(0),
                exec_phys.as_u64(),
            );
            let (pid, _) =
                ptable.pid_table_mut().allocate(sched, &allocator).ok_or_else(|| {
                    crate::proc::address_space::teardown_user_space(
                        pml4_phys.raw(),
                        &mut allocator,
                    );
                    KernelError::PoolExhausted
                })?;

            if crate::proc::address_space::pml4_ref_create(pml4_phys.raw(), pid).is_err() {
                let _ = ptable.pid_table_mut().free(pid);
                crate::proc::address_space::teardown_user_space(pml4_phys.raw(), &mut allocator);
                return Err(KernelError::ResourceExhausted);
            }

            let exec = crate::proc::manager::make_default_exec(pml4_phys, true);
            finish_tcb_init(ptable, pid, exec_virt, exec, pml4_phys, None)
        }
        AddressSpaceMode::Shared { parent_pid } => {
            let parent_exec = ptable.exec(parent_pid).ok_or(KernelError::InvalidObject)?;
            let pml4_phys = parent_exec.pml4_phys;
            let root_cnode = parent_exec.root_cnode;
            let cnode_depth = parent_exec.cnode_depth;
            let guard_value = parent_exec.root_guard_value;
            let guard_bits = parent_exec.root_guard_bits;

            crate::proc::address_space::pml4_ref_share(pml4_phys.raw())?;

            let sched = crate::proc::manager::make_default_sched(
                crate::types::Pid::new(0),
                Generation::new(0),
                exec_phys.as_u64(),
            );
            let allocator = crate::mem::phys::BitmapFrameAllocator;
            let (pid, _) =
                ptable.pid_table_mut().allocate(sched, &allocator).ok_or_else(|| {
                    let _ = crate::proc::address_space::pml4_ref_release(pml4_phys.raw());
                    KernelError::PoolExhausted
                })?;

            let mut exec = crate::proc::manager::make_default_exec(pml4_phys, true);
            exec.root_cnode = root_cnode;
            exec.cnode_depth = cnode_depth;
            exec.root_guard_value = guard_value;
            exec.root_guard_bits = guard_bits;

            finish_tcb_init(ptable, pid, exec_virt, exec, pml4_phys, root_cnode)
        }
    }
}

fn finish_tcb_init(
    ptable: &mut crate::proc::ProcessManager,
    pid: crate::types::Pid,
    exec_virt: x86_64::VirtAddr,
    exec: crate::proc::ExecContext,
    pml4_phys: crate::mem::typed_addr::Pml4Phys,
    root_cnode: Option<(ObjPhys, Generation)>,
) -> Result<TcbInitResult, KernelError> {
    unsafe {
        core::ptr::write(exec_virt.as_mut_ptr::<crate::proc::ExecContext>(), exec);
    }

    let exec_ref = ptable.exec_mut(pid).ok_or(KernelError::InvalidObject)?;
    exec_ref.seal_context();

    Ok(TcbInitResult {
        pid,
        pml4_phys,
        root_cnode,
    })
}

#[allow(clippy::too_many_arguments)]
pub fn kernel_retype(
    pool: &mut ObjectPool,
    ptable: Option<&mut crate::proc::ProcessManager>,
    untyped_phys: ObjPhys,
    untyped_gen: Generation,
    obj_tag: ObjectTag,
    size_bits: u8,
    dest_cnode_phys: ObjPhys,
    dest_cnode_gen: Generation,
    dest_addr: u64,
    dest_depth: u8,
    dest_guard_value: u64,
    dest_guard_bits: u8,
    count: u32,
) -> Result<(), KernelError> {
    if count > MAX_RETYPE_COUNT {
        return Err(KernelError::InvalidParameter);
    }

    if obj_tag == ObjectTag::Process {
        match count {
            0 => return Ok(()),
            1 => {}
            _ => return Err(KernelError::InvalidParameter),
        }
        let pt = ptable.ok_or(KernelError::InvalidParameter)?;
        retype_tcb_from_untyped(
            pool,
            pt,
            untyped_phys,
            untyped_gen,
            AddressSpaceMode::New,
            size_bits,
            dest_cnode_phys,
            dest_cnode_gen,
            dest_addr,
            dest_depth,
            dest_guard_value,
            dest_guard_bits,
        )?;
        return Ok(());
    }

    if obj_tag == ObjectTag::VRegion {
        const MAX_PAGES: u32 = lancer_core::types::MAX_VREGION_PAGES as u32;
        let page_count = match count {
            0 => return Ok(()),
            1..=MAX_PAGES => count as u16,
            _ => return Err(KernelError::InvalidParameter),
        };
        retype_vregion_from_untyped(
            pool,
            untyped_phys,
            untyped_gen,
            page_count,
            dest_cnode_phys,
            dest_cnode_gen,
            dest_addr,
            dest_depth,
            dest_guard_value,
            dest_guard_bits,
        )?;
        return Ok(());
    }

    validate_retype_tag(obj_tag)?;

    let ut = pool.read_as::<UntypedObject>(untyped_phys, untyped_gen)?;
    let phys_base = PhysAddr::new(ut.phys_base);
    let state = ut.to_state();

    let result = match obj_tag {
        ObjectTag::Frame => state
            .try_retype(ObjectTag::Frame, size_bits, count)
            .map_err(retype_to_kernel_error)?,
        _ => {
            let (obj_size, obj_align) = lancer_core::untyped::object_layout(obj_tag, size_bits)
                .map_err(retype_to_kernel_error)?;
            state
                .try_allocate_raw(obj_size, obj_align, count)
                .map_err(retype_to_kernel_error)?
        }
    };

    let mut allocated: crate::static_vec::StaticVec<
        (ObjPhys, Generation),
        { MAX_RETYPE_COUNT as usize },
    > = crate::static_vec::StaticVec::new();

    let rollback_result = (0..count).try_for_each(|i| {
        let obj_phys =
            phys_base.as_u64() + result.start_offset() as u64 + (i as u64) * result.stride() as u64;

        if obj_tag == ObjectTag::CNode {
            let slot_count = 1usize << size_bits;
            let hhdm = crate::mem::addr::hhdm_offset();
            let base_ptr = (obj_phys + hhdm) as *mut u8;
            let byte_count = slot_count * core::mem::size_of::<super::table::CapSlot>();
            unsafe { core::ptr::write_bytes(base_ptr, 0, byte_count) };
        }

        let register_phys = write_repr_c_object(obj_tag, obj_phys, size_bits)?;

        let (obj_registered_phys, obj_gen) = match pool.register_object(register_phys, obj_tag) {
            Ok(pair) => pair,
            Err(_) => {
                if obj_tag == ObjectTag::Frame {
                    super::kernel_objects::free_slot(register_phys);
                }
                return Err(KernelError::PoolExhausted);
            }
        };

        if obj_tag == ObjectTag::Frame {
            let mut ft = super::frame_table::FRAME_TABLE.lock();
            if let Some(idx) = ft.alloc_idx() {
                let frame_obj = unsafe {
                    &mut *(addr::phys_to_virt(PhysAddr::new(obj_registered_phys.raw()))
                        .as_mut_ptr::<FrameObject>())
                };
                frame_obj.frame_table_idx = idx;
            }
        }

        let cap = CapRef::new(obj_tag, obj_registered_phys, Rights::ALL, obj_gen);
        let slot_addr = dest_addr + i as u64;
        match super::cnode::resolve_and_insert(
            pool,
            dest_cnode_phys,
            dest_cnode_gen,
            slot_addr,
            dest_depth,
            dest_guard_value,
            dest_guard_bits,
            cap,
        ) {
            Ok(()) => {
                let _ = allocated.push((obj_registered_phys, obj_gen));
                Ok(())
            }
            Err(e) => {
                let _ = pool.free_phys(obj_registered_phys, obj_gen);
                if obj_tag == ObjectTag::Frame {
                    super::kernel_objects::free_slot(register_phys);
                }
                Err(e)
            }
        }
    });

    match rollback_result {
        Ok(()) => {
            let ut_mut = pool.write_as::<UntypedObject>(untyped_phys, untyped_gen)?;
            let mut state = ut_mut.to_state();
            state.commit_retype(&result);
            ut_mut.apply_state(&state);

            allocated
                .as_slice()
                .iter()
                .for_each(|&(child_phys, _child_gen)| {
                    let link_result =
                        super::derivation::link_child(pool, untyped_phys, untyped_gen, child_phys);
                    debug_assert!(
                        link_result.is_ok(),
                        "link_child failed after commit: {:?}",
                        link_result.err()
                    );
                });

            let frames_per_obj: u32 = match obj_tag {
                ObjectTag::Frame => 1,
                ObjectTag::CNode => lancer_core::cnode::frames_for_cnode(size_bits) as u32,
                _ => 0,
            };
            if frames_per_obj > 0 {
                (0..count).for_each(|i| {
                    let base = phys_base.as_u64()
                        + result.start_offset() as u64
                        + (i as u64) * result.stride() as u64;
                    (0..frames_per_obj).for_each(|f| {
                        let frame_idx = ((base + f as u64 * 4096) / 4096) as usize;
                        crate::mem::phys::BitmapFrameAllocator::mark_used(frame_idx);
                    });
                });
            }
            Ok(())
        }
        Err(e) => {
            allocated
                .as_slice()
                .iter()
                .enumerate()
                .for_each(|(i, &(phys, generation))| {
                    let slot_addr = dest_addr + i as u64;
                    let _ = super::cnode::resolve_and_clear(
                        pool,
                        dest_cnode_phys,
                        dest_cnode_gen,
                        slot_addr,
                        dest_depth,
                        dest_guard_value,
                        dest_guard_bits,
                    );
                    if let Ok(Some((old_phys, _))) = pool.free_phys(phys, generation)
                        && obj_tag == ObjectTag::Frame
                    {
                        super::kernel_objects::free_slot(old_phys);
                    }
                });
            Err(e)
        }
    }
}

#[allow(clippy::too_many_arguments)]
fn retype_tcb_from_untyped(
    pool: &mut ObjectPool,
    ptable: &mut crate::proc::ProcessManager,
    untyped_phys: ObjPhys,
    untyped_gen: Generation,
    mode: AddressSpaceMode,
    size_bits: u8,
    dest_cnode_phys: ObjPhys,
    dest_cnode_gen: Generation,
    dest_addr: u64,
    dest_depth: u8,
    dest_guard_value: u64,
    dest_guard_bits: u8,
) -> Result<crate::types::Pid, KernelError> {
    let ut = pool.read_as::<UntypedObject>(untyped_phys, untyped_gen)?;
    if matches!(mode, AddressSpaceMode::Shared { .. }) && ut.is_device != 0 {
        return Err(KernelError::InvalidType);
    }

    let phys_base = PhysAddr::new(ut.phys_base);
    let (obj_size, obj_align) = lancer_core::untyped::object_layout(ObjectTag::Process, size_bits)
        .map_err(retype_to_kernel_error)?;
    let result = ut
        .to_state()
        .try_allocate_raw(obj_size, obj_align, 1)
        .map_err(retype_to_kernel_error)?;

    let tcb_init = init_tcb_memory(ptable, phys_base, result.start_offset() as u64, mode)?;
    let pid = tcb_init.pid;
    let pml4_phys = tcb_init.pml4_phys;
    let shared_cnode = tcb_init.root_cnode;

    if let Some((cnode_phys, cnode_gen)) = shared_cnode
        && let Err(e) = pool.inc_ref(cnode_phys, cnode_gen)
    {
        let _ = ptable.pid_table_mut().free(pid);
        let _ = crate::proc::address_space::pml4_ref_release(pml4_phys.raw());
        return Err(e);
    }

    let teardown_tcb = |pt: &mut crate::proc::ProcessManager, pool: &mut ObjectPool| {
        if let Some((cnode_phys, cnode_gen)) = shared_cnode {
            pool.dec_ref_phys(cnode_phys, cnode_gen);
        }
        let _ = pt.pid_table_mut().free(pid);
        match shared_cnode {
            Some(_) => {
                let _ = crate::proc::address_space::pml4_ref_release(pml4_phys.raw());
            }
            None => {
                let mut allocator = crate::mem::phys::BitmapFrameAllocator;
                crate::proc::address_space::teardown_user_space(pml4_phys.raw(), &mut allocator);
            }
        }
    };

    let obj_phys = phys_base.as_u64() + result.start_offset() as u64;
    let header = KernelObjectHeader::new(ObjectTag::Process, 0, 64);
    let mut proc_obj = lancer_core::object_layout::ProcessObject::init_default(header);
    proc_obj.pid = pid.raw();
    unsafe {
        core::ptr::write(
            addr::phys_to_virt(PhysAddr::new(obj_phys))
                .as_mut_ptr::<lancer_core::object_layout::ProcessObject>(),
            proc_obj,
        );
    }

    let (obj_registered_phys, obj_gen) = match pool.register_object(obj_phys, ObjectTag::Process) {
        Ok(pair) => pair,
        Err(_) => {
            teardown_tcb(ptable, pool);
            return Err(KernelError::PoolExhausted);
        }
    };

    let cap = CapRef::new(ObjectTag::Process, obj_registered_phys, Rights::ALL, obj_gen);
    if let Err(e) = super::cnode::resolve_and_insert(
        pool,
        dest_cnode_phys,
        dest_cnode_gen,
        dest_addr,
        dest_depth,
        dest_guard_value,
        dest_guard_bits,
        cap,
    ) {
        let _ = pool.free_phys(obj_registered_phys, obj_gen);
        teardown_tcb(ptable, pool);
        return Err(e);
    }

    let ut_mut = pool.write_as::<UntypedObject>(untyped_phys, untyped_gen)?;
    let mut state = ut_mut.to_state();
    state.commit_retype(&result);
    ut_mut.apply_state(&state);

    let link_result =
        super::derivation::link_child(pool, untyped_phys, untyped_gen, obj_registered_phys);
    debug_assert!(
        link_result.is_ok(),
        "link_child failed after tcb commit: {:?}",
        link_result.err()
    );

    let total_frames = result.stride().div_ceil(4096);
    (0..total_frames).for_each(|f| {
        let frame_idx = ((obj_phys + f as u64 * 4096) / 4096) as usize;
        crate::mem::phys::BitmapFrameAllocator::mark_used(frame_idx);
    });

    Ok(pid)
}

#[allow(clippy::too_many_arguments)]
pub fn retype_thread_from_untyped(
    pool: &mut ObjectPool,
    ptable: &mut crate::proc::ProcessManager,
    untyped_phys: ObjPhys,
    untyped_gen: Generation,
    parent_pid: crate::types::Pid,
    dest_cnode_phys: ObjPhys,
    dest_cnode_gen: Generation,
    dest_addr: u64,
    dest_depth: u8,
    dest_guard_value: u64,
    dest_guard_bits: u8,
) -> Result<crate::types::Pid, KernelError> {
    retype_tcb_from_untyped(
        pool,
        ptable,
        untyped_phys,
        untyped_gen,
        AddressSpaceMode::Shared { parent_pid },
        0,
        dest_cnode_phys,
        dest_cnode_gen,
        dest_addr,
        dest_depth,
        dest_guard_value,
        dest_guard_bits,
    )
}

#[allow(clippy::too_many_arguments)]
pub(crate) fn retype_vregion_from_untyped(
    pool: &mut ObjectPool,
    untyped_phys: ObjPhys,
    untyped_gen: Generation,
    page_count: u16,
    dest_cnode_phys: ObjPhys,
    dest_cnode_gen: Generation,
    dest_addr: u64,
    dest_depth: u8,
    dest_guard_value: u64,
    dest_guard_bits: u8,
) -> Result<(), KernelError> {
    let ut = pool.read_as::<UntypedObject>(untyped_phys, untyped_gen)?;
    if ut.is_device != 0 {
        return Err(KernelError::InvalidType);
    }

    let phys_base = PhysAddr::new(ut.phys_base);
    let backing_size = page_count as u32 * 4096;
    let result = ut
        .to_state()
        .try_allocate_raw(backing_size, 4096, 1)
        .map_err(retype_to_kernel_error)?;

    let backing_phys = phys_base.as_u64() + result.start_offset() as u64;

    let ut_mut = pool.write_as::<UntypedObject>(untyped_phys, untyped_gen)?;
    let mut state = ut_mut.to_state();
    state.commit_retype(&result);
    ut_mut.apply_state(&state);

    let kobj_phys =
        super::kernel_objects::alloc_slot().ok_or(KernelError::PoolExhausted)?;

    let header = KernelObjectHeader::new(ObjectTag::VRegion, 0, 64);
    let mut obj = VRegionObject::init_default(header);
    obj.phys_base = backing_phys;
    obj.page_count = page_count;
    super::kernel_objects::write_at(kobj_phys, obj);

    let (obj_registered_phys, obj_gen) = match pool.register_object(kobj_phys, ObjectTag::VRegion) {
        Ok(pair) => pair,
        Err(_) => {
            super::kernel_objects::free_slot(kobj_phys);
            return Err(KernelError::PoolExhausted);
        }
    };

    let cap = CapRef::new(ObjectTag::VRegion, obj_registered_phys, Rights::ALL, obj_gen);
    if let Err(e) = super::cnode::resolve_and_insert(
        pool,
        dest_cnode_phys,
        dest_cnode_gen,
        dest_addr,
        dest_depth,
        dest_guard_value,
        dest_guard_bits,
        cap,
    ) {
        let _ = pool.free_phys(obj_registered_phys, obj_gen);
        super::kernel_objects::free_slot(kobj_phys);
        return Err(e);
    }

    super::derivation::link_child(pool, untyped_phys, untyped_gen, obj_registered_phys)
        .expect("link_child failed after vregion commit");

    (0..page_count as u32).for_each(|f| {
        let frame_idx = ((backing_phys + f as u64 * 4096) / 4096) as usize;
        crate::mem::phys::BitmapFrameAllocator::mark_used(frame_idx);
    });

    Ok(())
}

fn retype_to_kernel_error(e: lancer_core::untyped::RetypeError) -> KernelError {
    match e {
        lancer_core::untyped::RetypeError::InsufficientSpace => KernelError::ResourceExhausted,
        lancer_core::untyped::RetypeError::DeviceRejectsNonFrame => KernelError::InvalidType,
        lancer_core::untyped::RetypeError::ZeroCount => KernelError::InvalidParameter,
        lancer_core::untyped::RetypeError::Overflow => KernelError::InvalidParameter,
        lancer_core::untyped::RetypeError::InvalidSizeBits => KernelError::InvalidParameter,
        lancer_core::untyped::RetypeError::InvalidType => KernelError::InvalidType,
    }
}