use alloc::sync::Arc;
use core::{intrinsics::unlikely, panic};
use log::error;
use system_error::SystemError;
use x86::{bits64::rflags::RFlags, controlregs::Cr4};

use crate::{
    arch::{
        interrupt::{trap::X86PfErrorCode, TrapFrame},
        ipc::signal::Signal,
        mm::{MemoryManagementArch, X86_64MMArch},
        CurrentIrqArch, MMArch,
    },
    exception::{extable::ExceptionTableManager, InterruptArch},
    ipc::{
        signal::force_sig_fault_to_current,
        signal_types::{SignalArch, BUS_ADRERR, SEGV_ACCERR, SEGV_MAPERR},
    },
    mm::{
        fault::{FaultFlags, PageFaultHandler, PageFaultMessage},
        oom::{OomContext, OomOutcome},
        ucontext::{AddressSpace, LockedVMA},
        VirtAddr, VirtRegion, VmFaultReason, VmFlags,
    },
    process::ProcessManager,
};

use super::LockedFrameAllocator;

pub type PageMapper =
    crate::mm::page::PageMapper<crate::arch::x86_64::mm::X86_64MMArch, LockedFrameAllocator>;

impl X86_64MMArch {
    /// A protection fault can block behind the address-space write lock while
    /// mprotect/mremap publishes a new leaf PTE. The hardware error code then
    /// describes the old PTE, so revalidate write/instruction permission from
    /// the stable page table before converting that stale fault into SIGSEGV.
    fn stale_protection_fault_resolved(
        mapper: &PageMapper,
        address: VirtAddr,
        error_code: X86PfErrorCode,
        fault_from_user: bool,
    ) -> bool {
        const EXPECTED_FLAGS: u32 = X86PfErrorCode::X86_PF_PROT.bits()
            | X86PfErrorCode::X86_PF_WRITE.bits()
            | X86PfErrorCode::X86_PF_USER.bits()
            | X86PfErrorCode::X86_PF_INSTR.bits();

        if !fault_from_user
            || !error_code.contains(X86PfErrorCode::X86_PF_PROT | X86PfErrorCode::X86_PF_USER)
            || error_code.bits() & !EXPECTED_FLAGS != 0
            || !error_code.intersects(X86PfErrorCode::X86_PF_WRITE | X86PfErrorCode::X86_PF_INSTR)
        {
            return false;
        }

        let Some(entry) = mapper.get_entry(address, 0) else {
            return false;
        };
        entry.present()
            && entry.flags().has_user()
            && (!error_code.contains(X86PfErrorCode::X86_PF_WRITE) || entry.write())
            && (!error_code.contains(X86PfErrorCode::X86_PF_INSTR) || entry.flags().has_execute())
    }

    pub fn vma_access_error(vma: Arc<LockedVMA>, error_code: X86PfErrorCode) -> bool {
        let vm_flags = *vma.lock().vm_flags();
        let foreign = false;
        if error_code.contains(X86PfErrorCode::X86_PF_PK) {
            return true;
        }

        if unlikely(error_code.contains(X86PfErrorCode::X86_PF_SGX)) {
            return true;
        }

        if !Self::vma_access_permitted(
            vma.clone(),
            error_code.contains(X86PfErrorCode::X86_PF_WRITE),
            error_code.contains(X86PfErrorCode::X86_PF_INSTR),
            foreign,
        ) {
            return true;
        }

        if error_code.contains(X86PfErrorCode::X86_PF_WRITE) {
            if unlikely(!vm_flags.contains(VmFlags::VM_WRITE)) {
                return true;
            }
            return false;
        }

        if unlikely(error_code.contains(X86PfErrorCode::X86_PF_PROT)) {
            return true;
        }

        if unlikely(!vma.is_accessible()) {
            return true;
        }
        false
    }

    pub fn show_fault_oops(regs: &TrapFrame, error_code: X86PfErrorCode, address: VirtAddr) {
        let mapper =
            unsafe { PageMapper::current(crate::mm::PageTableKind::User, LockedFrameAllocator) };
        if let Some(entry) = mapper.get_entry(address, 0) {
            if entry.present() {
                if !entry.flags().has_execute() {
                    error!("kernel tried to execute NX-protected page - exploit attempt?");
                } else if mapper.table().phys().data() & MMArch::ENTRY_FLAG_USER != 0
                    && unsafe { x86::controlregs::cr4().contains(Cr4::CR4_ENABLE_SMEP) }
                {
                    error!("unable to execute userspace code (SMEP?)");
                }
            }
        }
        if address.data() < X86_64MMArch::PAGE_SIZE && !regs.is_from_user() {
            error!(
                "BUG: kernel NULL pointer dereference, address: {:#x}",
                address.data()
            );
        } else {
            error!(
                "BUG: unable to handle page fault for address: {:#x}",
                address.data()
            );
        }

        error!(
            "#PF: {} {} in {} mode\n",
            if error_code.contains(X86PfErrorCode::X86_PF_USER) {
                "user"
            } else {
                "supervisor"
            },
            if error_code.contains(X86PfErrorCode::X86_PF_INSTR) {
                "instruction fetch"
            } else if error_code.contains(X86PfErrorCode::X86_PF_WRITE) {
                "write access"
            } else {
                "read access"
            },
            if regs.is_from_user() {
                "user"
            } else {
                "kernel"
            }
        );
        error!(
            "#PF: error_code({:#04x}) - {}\n",
            error_code,
            if !error_code.contains(X86PfErrorCode::X86_PF_PROT) {
                "not-present page"
            } else if error_code.contains(X86PfErrorCode::X86_PF_RSVD) {
                "reserved bit violation"
            } else if error_code.contains(X86PfErrorCode::X86_PF_PK) {
                "protection keys violation"
            } else {
                "permissions violation"
            }
        );
    }

    pub fn page_fault_oops(regs: &TrapFrame, error_code: X86PfErrorCode, address: VirtAddr) {
        if regs.is_from_user() {
            Self::show_fault_oops(regs, error_code, address);
        }
        panic!()
    }

    /// 内核态缺页异常处理
    /// ## 参数
    ///
    /// - `regs`: 中断栈帧
    /// - `error_code`: 错误标志
    /// - `address`: 发生缺页异常的虚拟地址
    pub fn do_kern_addr_fault(
        regs: &'static mut TrapFrame,
        error_code: X86PfErrorCode,
        address: VirtAddr,
    ) {
        // 尝试异常表修复
        if Self::try_fixup_exception(regs, error_code, address) {
            // 成功修复,直接返回
            return;
        }

        unsafe { crate::debug::traceback::lookup_kallsyms(regs.rip, 0xff) };
        let pcb = crate::process::ProcessManager::current_pcb();
        let kstack_guard_addr = pcb.kernel_stack().guard_page_address();
        if let Some(guard_page) = kstack_guard_addr {
            let guard_page_size = pcb.kernel_stack().guard_page_size().unwrap();
            if address.data() >= guard_page.data()
                && address.data() < guard_page.data() + guard_page_size
            {
                // 发生在内核栈保护页上
                error!(
                    "kernel stack guard page fault at {:#x}, guard page range: {:#x} - {:#x}",
                    address.data(),
                    guard_page.data(),
                    guard_page.data() + guard_page_size
                );
            }
        }
        panic!(
            "do_kern_addr_fault has not yet been implemented, 
        fault address: {:#x},
        rip: {:#x},
        error_code: {:#b}, 
        pid: {}\n",
            address.data(),
            regs.rip,
            error_code,
            pcb.raw_pid().data()
        );
        //TODO https://code.dragonos.org.cn/xref/linux-6.6.21/arch/x86/mm/fault.c#do_kern_addr_fault
    }

    /// 尝试使用异常表修复页错误
    ///
    /// ## 返回值
    /// - `true`: 成功修复,可以继续执行
    /// - `false`: 无法修复,是真正的内核错误
    #[inline(never)]
    fn try_fixup_exception(
        regs: &mut TrapFrame,
        _error_code: X86PfErrorCode,
        address: VirtAddr,
    ) -> bool {
        // 只处理用户空间地址的访问错误
        if !address.check_user() {
            return false;
        }

        // 在异常表中查找修复代码
        if let Some(fixup_addr) = ExceptionTableManager::search_exception_table(regs.rip as usize) {
            // log::debug!(
            //     "Page fault at {:#x} accessing user address {:#x}, fixed up to {:#x}",
            //     regs.rip,
            //     address.data(),
            //     fixup_addr
            // );

            // 修改trap frame的RIP到修复代码
            regs.rip = fixup_addr as u64;

            return true;
        }

        false
    }

    /// 用户态缺页异常处理
    /// ## 参数
    ///
    /// - `regs`: 中断栈帧
    /// - `error_code`: 错误标志
    /// - `address`: 发生缺页异常的虚拟地址
    pub unsafe fn do_user_addr_fault(
        regs: &'static mut TrapFrame,
        error_code: X86PfErrorCode,
        address: VirtAddr,
    ) {
        // log::debug!("fault at {:?}:{:?}",
        // address,
        // error_code,
        // );
        let rflags = RFlags::from_bits_truncate(regs.rflags);
        let mut flags: FaultFlags = FaultFlags::FAULT_FLAG_ALLOW_RETRY
            | FaultFlags::FAULT_FLAG_KILLABLE
            | FaultFlags::FAULT_FLAG_INTERRUPTIBLE;

        if error_code & (X86PfErrorCode::X86_PF_USER | X86PfErrorCode::X86_PF_INSTR)
            == X86PfErrorCode::X86_PF_INSTR
        {
            Self::page_fault_oops(regs, error_code, address);
        }

        if unlikely(
            !error_code.contains(X86PfErrorCode::X86_PF_USER)
                && unsafe { x86::controlregs::cr4().contains(Cr4::CR4_ENABLE_SMAP) }
                && !rflags.contains(RFlags::FLAGS_AC),
        ) {
            Self::page_fault_oops(regs, error_code, address);
        }

        if unlikely(error_code.contains(X86PfErrorCode::X86_PF_RSVD)) {
            // TODO https://code.dragonos.org.cn/xref/linux-6.6.21/arch/x86/mm/fault.c#pgtable_bad
            panic!(
                "Reserved bits are never expected to be set, error_code: {:#b}, address: {:#x}",
                error_code,
                address.data()
            );
        }

        if regs.is_from_user() {
            unsafe { CurrentIrqArch::interrupt_enable() };
            flags |= FaultFlags::FAULT_FLAG_USER;
        } else if rflags.contains(RFlags::FLAGS_IF) {
            unsafe { CurrentIrqArch::interrupt_enable() };
        }

        if error_code.contains(X86PfErrorCode::X86_PF_SHSTK) {
            flags |= FaultFlags::FAULT_FLAG_WRITE;
        }
        if error_code.contains(X86PfErrorCode::X86_PF_WRITE) {
            flags |= FaultFlags::FAULT_FLAG_WRITE;
        }
        if error_code.contains(X86PfErrorCode::X86_PF_INSTR) {
            flags |= FaultFlags::FAULT_FLAG_INSTRUCTION;
        }

        let fault_from_user = regs.is_from_user();
        let send_fault_signal = |sig: Signal, code: i32, addr: VirtAddr| {
            // Only mark Trapped when reaching the final fault->signal branch;
            // recoverable page faults, retries, and exception-table fixups
            // stay Running to continue completing the same XOL.
            if fault_from_user {
                crate::exception::uprobe::mark_current_xol_trapped();
            }
            if let Err(e) = force_sig_fault_to_current(sig, code, addr) {
                error!(
                    "failed to force {:?} fault to current process: pid={:?}, code={}, addr={:#x}, err={:?}",
                    sig,
                    ProcessManager::current_pid(),
                    code,
                    addr.data(),
                    e
                );
            }
        };
        let send_segv_maperr = || send_fault_signal(Signal::SIGSEGV, SEGV_MAPERR, address);
        let send_segv_accerr = || send_fault_signal(Signal::SIGSEGV, SEGV_ACCERR, address);
        let send_bus_adrerr = || send_fault_signal(Signal::SIGBUS, BUS_ADRERR, address);

        // 辅助函数：处理内核访问用户地址失败的情况
        let handle_kernel_access_failed = |r: &mut TrapFrame| {
            // 如果是内核代码访问用户地址，尝试异常表修复
            if !r.is_from_user() {
                if Self::try_fixup_exception(r, error_code, address) {
                    return true; // 成功修复
                }
                // 如果异常表中没有，说明是bug
                error!(
                    "Kernel code at {:#x} illegally accessed user address {:#x} \
                     without exception table entry",
                    r.rip,
                    address.data()
                );
                panic!("Illegal user space access from kernel");
            }
            false // 不是内核访问，继续正常流程
        };

        if !regs.is_from_user()
            && ProcessManager::current_pcb().pagefault_disabled() > 0
            && handle_kernel_access_failed(regs)
        {
            return;
        }

        let fault_region = VirtRegion::new(
            VirtAddr::new(address.data() & !MMArch::PAGE_OFFSET_MASK),
            MMArch::PAGE_SIZE,
        );
        let current_address_space: Arc<AddressSpace> = AddressSpace::current().unwrap();

        let mut tried_direct_reclaim = false;

        'fault_retry: loop {
            let mut space_guard = current_address_space.write();
            let mut fault;
            loop {
                let vma = space_guard.mappings.find_nearest(address);
                let vma = match vma {
                    Some(vma) => vma,
                    None => {
                        if space_guard
                            .mappings
                            .first_reservation_conflict(fault_region)
                            .is_some()
                        {
                            drop(space_guard);
                            current_address_space.wait_for_no_reservation_conflict(fault_region);
                            space_guard = current_address_space.write();
                            continue;
                        }

                        // No VMA exists, so try an exception-table fixup first.
                        if handle_kernel_access_failed(regs) {
                            return;
                        }

                        log::error!(
                            "pid:{}, can not find nearest vma, \n\terror_code: {:?}, address: {:#x}, rip: {:#x}",
                            ProcessManager::current_pid().data(),
                            error_code,
                            address.data(),
                            regs.rip,
                        );

                        send_segv_maperr();
                        return;
                    }
                };
                let guard = vma.lock();
                let region = *guard.region();
                let vm_flags = *guard.vm_flags();
                drop(guard);

                if !region.contains(address) {
                    if space_guard
                        .mappings
                        .first_reservation_conflict(fault_region)
                        .is_some()
                    {
                        drop(space_guard);
                        current_address_space.wait_for_no_reservation_conflict(fault_region);
                        space_guard = current_address_space.write();
                        continue;
                    }

                    if vm_flags.contains(VmFlags::VM_GROWSDOWN) {
                        let extension_size = region.start() - address;
                        let extension_region = VirtRegion::new(address, extension_size);
                        if space_guard
                            .mappings
                            .first_reservation_conflict(extension_region)
                            .is_some()
                        {
                            drop(space_guard);
                            current_address_space
                                .wait_for_no_reservation_conflict(extension_region);
                            space_guard = current_address_space.write();
                            continue;
                        }

                        // 首先检查地址是否在栈的合理扩展范围内
                        // 如果地址距离栈底太远（超过最大栈限制），则这不是一个栈扩展请求，
                        // 而是一个无关的无效内存访问（例如空指针解引用）
                        let max_stack_limit = space_guard
                            .user_stack
                            .as_ref()
                            .map(|s| s.max_limit())
                            .unwrap_or(0);

                        if extension_size > max_stack_limit {
                            // 地址距离栈太远，这不是栈扩展请求，而是普通的无效内存访问
                            // 检查是否需要异常表修复
                            if handle_kernel_access_failed(regs) {
                                return; // 已通过异常表修复
                            }

                            send_segv_maperr();
                            return;
                        }

                        if !space_guard.can_extend_stack(extension_size) {
                            // 栈扩展超过限制
                            log::warn!(
                            "pid {} user stack limit exceeded, error_code: {:?}, address: {:#x}",
                            ProcessManager::current_pid().data(),
                            error_code,
                            address.data(),
                        );

                            // 栈溢出，检查是否需要异常表修复
                            if handle_kernel_access_failed(regs) {
                                return; // 已通过异常表修复
                            }

                            send_segv_maperr();
                            return;
                        }
                        let post_commit_population = match space_guard.extend_stack(extension_size)
                        {
                            Ok(request) => request,
                            Err(err) => {
                                log::warn!(
                                    "pid {} user stack extend failed: {:?}, error_code: {:?}, address: {:#x}",
                                    ProcessManager::current_pid().data(),
                                    err,
                                    error_code,
                                    address.data(),
                                );
                                if handle_kernel_access_failed(regs) {
                                    return;
                                }
                                send_segv_maperr();
                                return;
                            }
                        };
                        if let Some(request) = post_commit_population {
                            // MCL_FUTURE can make this new anonymous stack
                            // VMA locked.  Population owns a retry boundary,
                            // so it must happen only after the MM write guard
                            // is released; then restart the original fault
                            // against the revalidated VMA set.
                            drop(space_guard);
                            let _ = current_address_space.populate_locked_vma_post_commit(request);
                            continue 'fault_retry;
                        }
                        // extend_stack() inserts a distinct VMA. The nearest
                        // VMA sampled above does not contain `address`, so
                        // never feed it to the fault handler after a successful
                        // extension. Re-enter through the normal lookup path
                        // after releasing the write guard.
                        drop(space_guard);
                        continue 'fault_retry;
                    } else {
                        // The address is outside every VMA; try an exception-table fixup first.
                        if handle_kernel_access_failed(regs) {
                            return;
                        }

                        log::error!(
                            "pid: {} No mapped vma, error_code: {:?},rip:{:#x}, address: {:#x}, flags: {:?}",
                            ProcessManager::current_pid().data(),
                            error_code,
                            regs.rip,
                            address.data(),
                            flags
                        );

                        send_segv_maperr();
                        return;
                    }
                }

                if Self::stale_protection_fault_resolved(
                    &space_guard.user_mapper.utable,
                    address,
                    error_code,
                    fault_from_user,
                ) {
                    return;
                }

                if unlikely(Self::vma_access_error(vma.clone(), error_code)) {
                    // VMA权限错误，检查是否需要异常表修复
                    if handle_kernel_access_failed(regs) {
                        return; // 已通过异常表修复
                    }

                    // log::error!(
                    //     "vma access error, error_code: {:?}, address: {:#x}",
                    //     error_code,
                    //     address.data(),
                    // );

                    send_segv_accerr();
                    return;
                }
                let mapper = &mut space_guard.user_mapper.utable;
                let message = PageFaultMessage::new(
                    vma.clone(),
                    address,
                    flags,
                    mapper,
                    current_address_space.clone(),
                );

                let outcome = PageFaultHandler::handle_mm_fault(message);
                fault = outcome.reason;

                if fault.contains(VmFaultReason::VM_FAULT_COMPLETED) {
                    return;
                }

                if unlikely(fault.contains(VmFaultReason::VM_FAULT_RETRY)) {
                    flags |= FaultFlags::FAULT_FLAG_TRIED;
                    let wait = outcome.retry_wait;
                    drop(space_guard);
                    if let Some(wait) = wait {
                        match wait.wait() {
                            Ok(()) => {}
                            Err(SystemError::EINTR | SystemError::ERESTARTSYS) => {
                                if regs.is_from_user() {
                                    <crate::arch::ipc::signal::X86_64SignalArch as SignalArch>::do_signal_or_restart(regs);
                                } else {
                                    handle_kernel_access_failed(regs);
                                }
                                return;
                            }
                            Err(SystemError::ENOMEM) => {
                                fault = VmFaultReason::VM_FAULT_OOM;
                                crate::mm::fault::account_fault_result(flags, fault);
                                space_guard = current_address_space.write();
                                break;
                            }
                            Err(_) => {
                                fault = VmFaultReason::VM_FAULT_SIGBUS;
                                crate::mm::fault::account_fault_result(flags, fault);
                                space_guard = current_address_space.write();
                                break;
                            }
                        }
                    }
                    continue 'fault_retry;
                } else {
                    break;
                }
            }

            let vm_fault_error = VmFaultReason::VM_FAULT_OOM
                | VmFaultReason::VM_FAULT_SIGBUS
                | VmFaultReason::VM_FAULT_SIGSEGV
                | VmFaultReason::VM_FAULT_HWPOISON
                | VmFaultReason::VM_FAULT_HWPOISON_LARGE
                | VmFaultReason::VM_FAULT_FALLBACK;

            if fault.intersects(vm_fault_error) {
                // 内核态访问用户地址（如 copy_from_user）应走异常表修复，返回 -EFAULT，而不是发送信号/崩溃
                if !regs.is_from_user() {
                    if Self::try_fixup_exception(regs, error_code, address) {
                        return;
                    }
                    panic!(
                        "kernel access to user addr failed without fixup, fault: {:?}, addr: {:#x}, rip: {:#x}",
                        fault,
                        address.data(),
                        regs.rip
                    );
                }

                if fault.contains(VmFaultReason::VM_FAULT_OOM) {
                    // Skip direct reclaim for fault-injection OOM: the injection
                    // consumes a fail count on the first fault, so reclaiming
                    // then retrying would bypass the killer entirely.
                    if !tried_direct_reclaim
                        && !fault.contains(VmFaultReason::VM_FAULT_OOM_INJECTED)
                    {
                        tried_direct_reclaim = true;
                        drop(space_guard);
                        crate::mm::page::PageReclaimer::shrink_list(
                            <crate::mm::allocator::page_frame::PageFrameCount>::new(64),
                        );
                        continue;
                    }

                    error!(
                        "page fault OOM: pid={:?}, addr={:#x}, rip={:#x}, fault={:?}",
                        ProcessManager::current_pid(),
                        address.data(),
                        regs.rip,
                        fault
                    );
                    drop(space_guard);
                    let ctx = OomContext {
                        trigger_pid: ProcessManager::current_pid(),
                        trigger_tgid: ProcessManager::current_pcb().raw_tgid(),
                        fault_address: address,
                        fault_ip: regs.rip as usize,
                        order: 0,
                    };
                    match crate::mm::oom::pagefault_out_of_memory(ctx) {
                        OomOutcome::Retry => {
                            continue;
                        }
                        OomOutcome::CurrentTaskKilled => {
                            <crate::arch::ipc::signal::X86_64SignalArch as SignalArch>::do_signal_or_restart(regs);
                            return;
                        }
                        OomOutcome::NoVictim => {
                            panic!(
                                "fatal user page-fault OOM: pid={:?}, tgid={:?}, addr={:#x}, rip={:#x}",
                                ProcessManager::current_pid(),
                                ProcessManager::current_pcb().raw_tgid(),
                                address.data(),
                                regs.rip
                            );
                        }
                    }
                } else if fault.contains(VmFaultReason::VM_FAULT_SIGBUS)
                    || fault.contains(VmFaultReason::VM_FAULT_HWPOISON)
                    || fault.contains(VmFaultReason::VM_FAULT_HWPOISON_LARGE)
                {
                    // Linux x86 maps these fault handler errors to SIGBUS/BUS_ADRERR
                    // except for hwpoison's MCE-specific si_code, which DragonOS does not model yet.
                    send_bus_adrerr();
                } else if fault.contains(VmFaultReason::VM_FAULT_SIGSEGV) {
                    send_segv_maperr();
                } else {
                    panic!("unexpected fault error: {:?}", fault);
                }
                return;
            }

            panic!("fault error: {:?}", fault)
        }
    }
}
