use lancer_core::object_tag::ObjectTag;

use super::object::{
    CNodeData, EndpointData, FrameObjectData, IrqHandlerData, IrqSource, NotificationData,
    ObjectData, PortRange, SchedContextData,
};
use super::pool::ObjectPool;
use super::table::{CapRef, Rights};
use crate::error::KernelError;
use crate::types::{Generation, ObjectId, Priority};
use x86_64::PhysAddr;

const MAX_RETYPE_COUNT: u32 = 64;

fn make_object_data(
    tag: ObjectTag,
    frame_phys: Option<PhysAddr>,
    size_bits: u8,
) -> Result<ObjectData, KernelError> {
    match tag {
        ObjectTag::Endpoint => Ok(ObjectData::Endpoint(EndpointData::new())),
        ObjectTag::Notification => Ok(ObjectData::Notification(NotificationData::new())),
        ObjectTag::Frame => {
            let phys = frame_phys.ok_or(KernelError::InvalidParameter)?;
            Ok(ObjectData::Frame(FrameObjectData::new_from_untyped(phys)))
        }
        ObjectTag::CNode => {
            let phys = frame_phys.ok_or(KernelError::InvalidParameter)?;
            let frame_count = lancer_core::cnode::frames_for_cnode(size_bits);
            Ok(ObjectData::CNode(CNodeData {
                slots_phys: phys,
                size_bits,
                frame_count,
            }))
        }
        ObjectTag::SchedContext => Ok(ObjectData::SchedContext(SchedContextData::new(
            0,
            0,
            Priority::IDLE,
        ))),
        ObjectTag::IrqHandler => Ok(ObjectData::IrqHandler(IrqHandlerData {
            vector: crate::arch::x86_64::idt::IrqVector::new(32),
            source: IrqSource::Ioapic {
                gsi: crate::arch::ioapic::Gsi::new(0),
            },
            port_range: match PortRange::new(0, 1) {
                Some(pr) => pr,
                None => return Err(KernelError::InvalidParameter),
            },
        })),
        ObjectTag::Untyped
        | ObjectTag::Process
        | ObjectTag::Framebuffer
        | ObjectTag::PciDevice
        | ObjectTag::MemoryRegion => Err(KernelError::InvalidType),

    }
}

#[allow(clippy::too_many_arguments)]
pub fn kernel_retype(
    pool: &mut ObjectPool,
    untyped_id: ObjectId,
    untyped_gen: Generation,
    obj_tag: ObjectTag,
    size_bits: u8,
    dest_cnode_id: ObjectId,
    dest_cnode_gen: Generation,
    dest_addr: u64,
    dest_depth: u8,
    count: u32,
) -> Result<(), KernelError> {
    match count > MAX_RETYPE_COUNT {
        true => return Err(KernelError::InvalidParameter),
        false => {}
    }

    let untyped_data = pool.get(untyped_id, untyped_gen)?;
    let untyped = untyped_data.as_untyped()?;
    let phys_base = untyped.phys_base;

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

    let mut allocated: crate::static_vec::StaticVec<(ObjectId, 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;

        let phys_for_data = match obj_tag {
            ObjectTag::Frame | ObjectTag::CNode => Some(PhysAddr::new(obj_phys)),
            _ => None,
        };

        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 data = make_object_data(obj_tag, phys_for_data, size_bits)?;

        let kern_result = pool
            .allocate(data)
            .map_err(|_| KernelError::PoolExhausted)?;
        let (obj_id, obj_gen) = kern_result;

        let cap = CapRef::new(obj_tag, obj_id, Rights::ALL, obj_gen);
        let slot_addr = dest_addr + i as u64;
        match super::cnode::resolve_and_insert(
            pool,
            dest_cnode_id,
            dest_cnode_gen,
            slot_addr,
            dest_depth,
            cap,
        ) {
            Ok(()) => {
                let _ = allocated.push((obj_id, obj_gen));
                Ok(())
            }
            Err(e) => {
                let _ = pool.free(obj_id, obj_gen);
                Err(e)
            }
        }
    });

    match rollback_result {
        Ok(()) => {
            let untyped_mut = pool.get_mut(untyped_id, untyped_gen)?;
            let ut = untyped_mut.as_untyped_mut()?;
            ut.state.commit_retype(&result);

            match obj_tag {
                ObjectTag::Frame => {
                    (0..count).for_each(|i| {
                        let frame_phys = phys_base.as_u64()
                            + result.start_offset() as u64
                            + (i as u64) * result.stride() as u64;
                        let frame_idx = (frame_phys / 4096) as usize;
                        crate::mem::phys::BitmapFrameAllocator::mark_used(frame_idx);
                    });
                }
                ObjectTag::CNode => {
                    let frame_count = lancer_core::cnode::frames_for_cnode(size_bits) as u32;
                    (0..count).for_each(|i| {
                        let cnode_phys = phys_base.as_u64()
                            + result.start_offset() as u64
                            + (i as u64) * result.stride() as u64;
                        (0..frame_count).for_each(|f| {
                            let frame_idx = ((cnode_phys + 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, &(id, generation))| {
                let slot_addr = dest_addr + i as u64;
                let _ = super::cnode::resolve_and_clear(
                    pool, dest_cnode_id, dest_cnode_gen, slot_addr, dest_depth,
                );
                let _ = pool.free(id, generation);
            });
            Err(e)
        }
    }
}

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,
    }
}