use crate::cap::cnode;
use crate::cap::frame_table::{self, FRAME_TABLE};
use crate::cap::object::ObjectTag;
use crate::cap::pool::POOL;
use crate::cap::table::Rights;
use crate::error::KernelError;
use crate::pci::DEVICE_TABLE;
use crate::proc::PROCESSES;
use crate::proc::context::CpuContext;
use lancer_core::object_layout::{FrameObject, PciDeviceObject};

use super::{SyscallResult, try_syscall};

pub fn sys_iommu_map(ctx: &mut CpuContext) {
    let pci_cap_addr = ctx.rdi;
    let frame_cap_addr = ctx.rsi;
    let pid = crate::arch::syscall::current_pid();

    let (pci_cap, frame_cap) = {
        let ptable = PROCESSES.lock();
        let pool = POOL.lock_after(&ptable);
        let pci = try_syscall!(
            ctx,
            cnode::resolve_caller_validate(
                pid,
                pci_cap_addr,
                ObjectTag::PciDevice,
                Rights::WRITE,
                &ptable,
                &pool,
            )
        );
        let frame = try_syscall!(
            ctx,
            cnode::resolve_caller_validate(
                pid,
                frame_cap_addr,
                ObjectTag::Frame,
                Rights::WRITE,
                &ptable,
                &pool,
            )
        );
        (pci, frame)
    };

    let device_table_idx = {
        let pool = POOL.lock();
        match pool.read_as::<PciDeviceObject>(pci_cap.phys(), pci_cap.generation()) {
            Ok(d) => d.device_table_idx,
            Err(e) => {
                ctx.rax = SyscallResult::error(e).raw();
                return;
            }
        }
    };

    let (bus, devfn) = {
        let dev_table = DEVICE_TABLE.lock();
        match dev_table.get(device_table_idx as usize) {
            Some(d) => (d.bus, d.devfn()),
            None => {
                ctx.rax = SyscallResult::error(KernelError::InvalidObject).raw();
                return;
            }
        }
    };

    let (frame_phys, ft_idx) = {
        let pool = POOL.lock();
        match pool.read_as::<FrameObject>(frame_cap.phys(), frame_cap.generation()) {
            Ok(f) => (f.phys_addr, f.frame_table_idx),
            Err(e) => {
                ctx.rax = SyscallResult::error(e).raw();
                return;
            }
        }
    };

    if ft_idx == lancer_core::header::NONE_SENTINEL as u16 {
        ctx.rax = SyscallResult::error(KernelError::ResourceExhausted).raw();
        return;
    }

    if !crate::iommu::is_available() {
        ctx.rax = SyscallResult::error(KernelError::NotInitialized).raw();
        return;
    }

    let hhdm_offset = crate::mem::addr::hhdm_offset();
    let mut allocator = crate::mem::phys::BitmapFrameAllocator;

    let (iova, domain_id) = {
        let mut devices = crate::iommu::DMA_DEVICES.lock();
        let dev = match devices
            .iter_mut()
            .find(|d| d.bus == bus && d.devfn == devfn)
        {
            Some(d) => d,
            None => {
                ctx.rax = SyscallResult::error(KernelError::NotFound).raw();
                return;
            }
        };

        let iova = match dev.iova_alloc.allocate(1) {
            Some(base) => base,
            None => {
                ctx.rax = SyscallResult::error(KernelError::ResourceExhausted).raw();
                return;
            }
        };

        match crate::iommu::tables::map_iova_to_phys(
            &dev.dma_pt,
            iova,
            x86_64::PhysAddr::new(frame_phys),
            &mut allocator,
            hhdm_offset,
        ) {
            Ok(()) => {}
            Err(e) => {
                let _ = dev.iova_alloc.free(iova, 1);
                ctx.rax = SyscallResult::error(e).raw();
                return;
            }
        }

        (iova, dev.domain_id)
    };

    if let Err(e) = crate::iommu::invalidate_iotlb_domain(domain_id) {
        crate::kprintln!(
            "[IOMMU] IOTLB invalidation warning during iommu_map: {:?}",
            e
        );
    }

    let mut ft = FRAME_TABLE.lock();
    let entry = ft.get_mut(ft_idx);
    match entry.iommu_mapping() {
        Some(_) => {
            drop(ft);
            crate::iommu::unmap_frame_iommu(bus, devfn, iova);
            ctx.rax = SyscallResult::error(KernelError::AlreadyExists).raw();
        }
        None => {
            entry.set_iommu(frame_table::IommuMapping { bus, devfn, iova });
            ctx.rax = SyscallResult::success(iova).raw();
        }
    }
}

pub fn sys_iommu_unmap(ctx: &mut CpuContext) {
    let pci_cap_addr = ctx.rdi;
    let frame_cap_addr = ctx.rsi;
    let pid = crate::arch::syscall::current_pid();

    let (pci_cap, frame_cap) = {
        let ptable = PROCESSES.lock();
        let pool = POOL.lock_after(&ptable);
        let pci = try_syscall!(
            ctx,
            cnode::resolve_caller_validate(
                pid,
                pci_cap_addr,
                ObjectTag::PciDevice,
                Rights::WRITE,
                &ptable,
                &pool,
            )
        );
        let frame = try_syscall!(
            ctx,
            cnode::resolve_caller_validate(
                pid,
                frame_cap_addr,
                ObjectTag::Frame,
                Rights::WRITE,
                &ptable,
                &pool,
            )
        );
        (pci, frame)
    };

    let device_table_idx = {
        let pool = POOL.lock();
        match pool.read_as::<PciDeviceObject>(pci_cap.phys(), pci_cap.generation()) {
            Ok(d) => d.device_table_idx,
            Err(e) => {
                ctx.rax = SyscallResult::error(e).raw();
                return;
            }
        }
    };

    let (bus, devfn) = {
        let dev_table = DEVICE_TABLE.lock();
        match dev_table.get(device_table_idx as usize) {
            Some(d) => (d.bus, d.devfn()),
            None => {
                ctx.rax = SyscallResult::error(KernelError::InvalidObject).raw();
                return;
            }
        }
    };

    let ft_idx = {
        let pool = POOL.lock();
        match pool.read_as::<FrameObject>(frame_cap.phys(), frame_cap.generation()) {
            Ok(f) => f.frame_table_idx,
            Err(e) => {
                ctx.rax = SyscallResult::error(e).raw();
                return;
            }
        }
    };
    if ft_idx == lancer_core::header::NONE_SENTINEL as u16 {
        ctx.rax = SyscallResult::error(KernelError::InvalidParameter).raw();
        return;
    }
    let mapping = {
        let mut ft = FRAME_TABLE.lock();
        let entry = ft.get_mut(ft_idx);
        match entry.iommu_mapping() {
            Some(m) if m.bus == bus && m.devfn == devfn => {
                let m_copy = frame_table::IommuMapping {
                    bus: m.bus,
                    devfn: m.devfn,
                    iova: m.iova,
                };
                entry.clear_iommu();
                m_copy
            }
            Some(_) => {
                ctx.rax = SyscallResult::error(KernelError::PermissionDenied).raw();
                return;
            }
            None => {
                ctx.rax = SyscallResult::error(KernelError::InvalidParameter).raw();
                return;
            }
        }
    };

    crate::iommu::unmap_frame_iommu(bus, devfn, mapping.iova);
    ctx.rax = SyscallResult::ok().raw();
}

pub fn sys_iommu_setup_device(ctx: &mut CpuContext) {
    let pid = crate::arch::syscall::current_pid();
    if pid != crate::types::Pid::new(0) {
        ctx.rax = SyscallResult::error(KernelError::PermissionDenied).raw();
        return;
    }

    let pci_cap_addr = ctx.rdi;
    let target_proc_addr = ctx.rsi;

    let (device_table_idx, target_pid) = {
        let ptable = PROCESSES.lock();
        let pool = POOL.lock_after(&ptable);
        let pci_cap = try_syscall!(
            ctx,
            cnode::resolve_caller_validate(
                pid,
                pci_cap_addr,
                ObjectTag::PciDevice,
                Rights::WRITE,
                &ptable,
                &pool,
            )
        );
        let proc_cap = try_syscall!(
            ctx,
            cnode::resolve_caller_validate(
                pid,
                target_proc_addr,
                ObjectTag::Process,
                Rights::READ,
                &ptable,
                &pool,
            )
        );
        let idx = match pool.read_as::<PciDeviceObject>(pci_cap.phys(), pci_cap.generation()) {
            Ok(d) => d.device_table_idx,
            Err(e) => {
                ctx.rax = SyscallResult::error(e).raw();
                return;
            }
        };
        let tpid = match pool.read_as::<lancer_core::object_layout::ProcessObject>(
            proc_cap.phys(),
            proc_cap.generation(),
        ) {
            Ok(p) => match crate::types::Pid::try_new(p.pid) {
                Some(pid) => pid,
                None => {
                    ctx.rax = SyscallResult::error(KernelError::InvalidObject).raw();
                    return;
                }
            },
            Err(e) => {
                ctx.rax = SyscallResult::error(e).raw();
                return;
            }
        };
        (idx, tpid)
    };

    if !crate::iommu::is_available() {
        ctx.rax = SyscallResult::ok().raw();
        return;
    }

    let hhdm_offset = crate::mem::addr::hhdm_offset();
    let mut allocator = crate::mem::phys::BitmapFrameAllocator;

    ctx.rax = match crate::iommu::setup_device_context(
        device_table_idx,
        target_pid,
        &mut allocator,
        hhdm_offset,
    ) {
        Ok(()) => SyscallResult::ok().raw(),
        Err(e) => SyscallResult::error(e).raw(),
    };
}