use core::{
    alloc::Layout,
    mem::size_of,
    ptr,
    sync::atomic::{AtomicU32, Ordering},
};

use libk_mm::{address::Virtualize, process::ProcessAddressSpace};
use yggdrasil_abi::error::Error;

// XXX
#[cfg(any(target_pointer_width = "64", rust_analyzer))]
const KERNEL_VIRT_OFFSET: usize = 0xFFFFFF8000000000;

#[cfg(any(target_pointer_width = "32", rust_analyzer))]
const KERNEL_VIRT_OFFSET: usize = 0xC0000000;

/// Helper trait to allow cross-address space access to pointers
pub trait ForeignPointer: Sized {
    /// Perform a volatile pointer write without dropping the old value.
    ///
    /// # Panics
    ///
    /// The function panics if any of the following conditions is met:
    ///
    /// * The address of the pointer is not mapped in the `space`.
    /// * The pointer is not writable.
    /// * The pointer is misaligned.
    ///
    /// # Safety
    ///
    /// As this function allows direct memory writes, it is inherently unsafe.
    unsafe fn write_foreign_volatile(self: *mut Self, space: &ProcessAddressSpace, value: Self);

    /// # Safety
    ///
    /// As this function allows direct memory writes, it is inherently unsafe.
    unsafe fn try_write_foreign_volatile(
        self: *mut Self,
        space: &ProcessAddressSpace,
        value: Self,
    ) -> Result<(), Error>;

    /// # Safety
    ///
    /// Even though this function does the necessary checks, it is still a raw pointer to reference
    /// conversion, and thus is unsafe.
    unsafe fn try_read_foreign_volatile(
        self: *const Self,
        space: &ProcessAddressSpace,
    ) -> Result<Self, Error>;

    /// Performs pointer validation for given address space:
    ///
    /// * Checks if the pointer has proper alignment for the type.
    /// * Checks if the pointer is mapped in the address space.
    /// * Checks if the pointer is above the userspace memory boundary.
    ///
    /// # Safety
    ///
    /// Even though this function does the necessary checks, it is still a raw pointer to reference
    /// conversion, and thus is unsafe.
    unsafe fn validate_user_ptr<'a>(
        self: *const Self,
        space: &ProcessAddressSpace,
    ) -> Result<&'a Self, Error>;

    /// [ForeignPointer::validate_user_ptr], with extra "writability" check.
    ///
    /// # Safety
    ///
    /// Even though this function does the necessary checks, it is still a raw pointer to reference
    /// conversion, and thus is unsafe.
    unsafe fn validate_user_mut<'a>(
        self: *mut Self,
        space: &ProcessAddressSpace,
    ) -> Result<&'a mut Self, Error>;

    /// [ForeignPointer::validate_user_ptr], but for slices
    ///
    /// # Safety
    ///
    /// Even though this function does the necessary checks, it is still a raw pointer to reference
    /// conversion, and thus is unsafe.
    unsafe fn validate_user_slice<'a>(
        self: *const Self,
        len: usize,
        space: &ProcessAddressSpace,
    ) -> Result<&'a [Self], Error>;

    /// [ForeignPointer::validate_user_slice], but for mutable slices
    ///
    /// # Safety
    ///
    /// Even though this function does the necessary checks, it is still a raw pointer to reference
    /// conversion, and thus is unsafe.
    unsafe fn validate_user_slice_mut<'a>(
        self: *mut Self,
        len: usize,
        space: &ProcessAddressSpace,
    ) -> Result<&'a mut [Self], Error>;
}

pub trait ForeignAtomic<T> {
    fn atomic_load_foreign(
        self: *const Self,
        space: &ProcessAddressSpace,
        ordering: Ordering,
    ) -> Result<T, Error>;
}

impl<T> ForeignPointer for T {
    unsafe fn write_foreign_volatile(self: *mut Self, space: &ProcessAddressSpace, value: T) {
        self.try_write_foreign_volatile(space, value)
            .expect("Invalid foreign pointer, could not write")
    }

    unsafe fn try_write_foreign_volatile(
        self: *mut Self,
        space: &ProcessAddressSpace,
        value: Self,
    ) -> Result<(), Error> {
        // TODO check align
        let addr = self as usize;
        let start_page = addr & !0xFFF;
        let end_page = (addr + size_of::<T>() - 1) & !0xFFF;
        let page_offset = addr & 0xFFF;

        if start_page != end_page {
            todo!("Foreign pointer write crossed a page boundary");
        }

        let phys_page = space.translate(start_page)?;

        let virt_ptr = phys_page.add(page_offset).virtualize() as *mut T;
        virt_ptr.write_unaligned(value);
        Ok(())
    }

    unsafe fn try_read_foreign_volatile(
        self: *const Self,
        space: &ProcessAddressSpace,
    ) -> Result<Self, Error> {
        let addr = self as usize;
        let start_page = addr & !0xFFF;
        let end_page = (addr + size_of::<T>() - 1) & !0xFFF;
        let page_offset = addr & 0xFFF;

        if start_page != end_page {
            todo!("Foreign pointer write crossed a page boundary");
        }

        let phys_page = space.translate(start_page)?;
        let virt_ptr = phys_page.add(page_offset).virtualize() as *const Self;

        Ok(virt_ptr.read_unaligned())
    }

    unsafe fn validate_user_slice_mut<'a>(
        self: *mut Self,
        len: usize,
        space: &ProcessAddressSpace,
    ) -> Result<&'a mut [Self], Error> {
        let base = self as usize;
        let layout = Layout::array::<T>(len).unwrap();

        validate_user_align_size(base, &layout)?;
        validate_user_region(space, base, layout.size(), true)?;

        Ok(core::slice::from_raw_parts_mut(self, len))
    }

    unsafe fn validate_user_slice<'a>(
        self: *const Self,
        len: usize,
        space: &ProcessAddressSpace,
    ) -> Result<&'a [Self], Error> {
        let base = self as usize;
        let layout = Layout::array::<T>(len).unwrap();

        validate_user_align_size(base, &layout)?;
        validate_user_region(space, base, layout.size(), false)?;

        Ok(core::slice::from_raw_parts(self, len))
    }

    unsafe fn validate_user_mut<'a>(
        self: *mut Self,
        space: &ProcessAddressSpace,
    ) -> Result<&'a mut Self, Error> {
        let addr = self as usize;
        let layout = Layout::new::<T>();

        // Common validation
        validate_user_align_size(addr, &layout)?;

        // Validate that the pages covered by this address are mapped as writable by the process
        // TODO for CoW this may differ
        validate_user_region(space, addr, layout.size(), true)?;

        Ok(&mut *self)
    }

    unsafe fn validate_user_ptr<'a>(
        self: *const Self,
        space: &ProcessAddressSpace,
    ) -> Result<&'a Self, Error> {
        let addr = self as usize;
        let layout = Layout::new::<T>();

        // Common validation
        validate_user_align_size(addr, &layout)?;
        validate_user_region(space, addr, layout.size(), false)?;

        Ok(&*self)
    }
}

impl ForeignAtomic<u32> for AtomicU32 {
    fn atomic_load_foreign(
        self: *const Self,
        space: &ProcessAddressSpace,
        ordering: Ordering,
    ) -> Result<u32, Error> {
        let virt = self.addr();
        if virt % size_of::<u32>() != 0 {
            // Misaligned atomic
            return Err(Error::InvalidMemoryOperation);
        }
        let phys = space.translate_lock(virt)?;
        let phys_virt: *const Self = ptr::with_exposed_provenance(phys.virtualize());
        Ok(unsafe { &*phys_virt }.load(ordering))
    }
}

fn validate_user_align_size(addr: usize, layout: &Layout) -> Result<(), Error> {
    // Explicitly disallow NULL
    if addr == 0 {
        return Err(Error::InvalidArgument);
    }
    // Validate alignment
    if addr % layout.align() != 0 {
        return Err(Error::InvalidArgument);
    }
    if addr + layout.size() > KERNEL_VIRT_OFFSET {
        todo!();
    }

    Ok(())
}

/// Validates access to given userspace memory region with given constraints
pub fn validate_user_region(
    space: &ProcessAddressSpace,
    base: usize,
    len: usize,
    _need_write: bool,
) -> Result<(), Error> {
    if base + len > crate::task::mem::KERNEL_VIRT_OFFSET {
        panic!("Invalid argument");
    }

    let aligned_start = base & !0xFFF;
    let aligned_end = (base + len + 0xFFF) & !0xFFF;

    for page in (aligned_start..aligned_end).step_by(0x1000) {
        // TODO check writability
        space.translate(page)?;
    }

    Ok(())
}
