// SPDX-License-Identifier: MPL-2.0

use core::{
    cmp::{max, min},
    num::NonZeroUsize,
    ops::Range,
};

use align_ext::AlignExt;
use ostd::{
    mm::{
        tlb::TlbFlushOp, vm_space::VmItem, CachePolicy, FrameAllocOptions, PageFlags, PageProperty,
        UFrame, VmSpace,
    },
    task::disable_preempt,
};

use super::interval_set::Interval;
use crate::{
    prelude::*,
    thread::exception::PageFaultInfo,
    vm::{
        perms::VmPerms,
        util::duplicate_frame,
        vmo::{CommitFlags, Vmo, VmoCommitError},
    },
};

/// Mapping a range of physical pages into a `Vmar`.
///
/// A `VmMapping` can bind with a `Vmo` which can provide physical pages for
/// mapping. Otherwise, it must be an anonymous mapping and will map any empty
/// physical page. A `VmMapping` binding with a `Vmo` is called VMO-backed
/// mapping. Generally, a VMO-backed mapping is a file-backed mapping. Yet
/// there are also some situations where specific pages that are not in a file
/// need to be mapped. e.g:
///  - Mappings to the VDSO data.
///  - Shared anonymous mappings. because the mapped pages need to be retained
///    and shared with other processes.
///
/// Such mappings will also be VMO-backed mappings.
///
/// This type controls the actual mapping in the [`VmSpace`]. It is a linear
/// type and cannot be [`Drop`]. To remove a mapping, use [`Self::unmap`].
#[derive(Debug)]
pub(super) struct VmMapping {
    /// The size of mapping, in bytes. The map size can even be larger than the
    /// size of VMO. Those pages outside VMO range cannot be read or write.
    ///
    /// Zero sized mapping is not allowed. So this field is always non-zero.
    map_size: NonZeroUsize,
    /// The base address relative to the root VMAR where the VMO is mapped.
    map_to_addr: Vaddr,
    /// Specific physical pages that need to be mapped. If this field is
    /// `None`, it means that the mapping is an independent anonymous mapping.
    ///
    /// The start of the virtual address maps to the start of the range
    /// specified in [`MappedVmo`].
    vmo: Option<MappedVmo>,
    /// Whether the mapping is shared.
    ///
    /// The updates to a shared mapping are visible among processes, or carried
    /// through to the underlying file for file-backed shared mappings.
    is_shared: bool,
    /// Whether the mapping needs to handle surrounding pages when handling
    /// page fault.
    handle_page_faults_around: bool,
    /// The permissions of pages in the mapping.
    ///
    /// All pages within the same `VmMapping` have the same permissions.
    perms: VmPerms,
}

impl Interval<Vaddr> for VmMapping {
    fn range(&self) -> Range<Vaddr> {
        self.map_to_addr..self.map_to_addr + self.map_size.get()
    }
}

/***************************** Basic methods *********************************/

impl VmMapping {
    pub(super) fn new(
        map_size: NonZeroUsize,
        map_to_addr: Vaddr,
        vmo: Option<MappedVmo>,
        is_shared: bool,
        handle_page_faults_around: bool,
        perms: VmPerms,
    ) -> Self {
        Self {
            map_size,
            map_to_addr,
            vmo,
            is_shared,
            handle_page_faults_around,
            perms,
        }
    }

    pub(super) fn new_fork(&self) -> Result<VmMapping> {
        Ok(VmMapping {
            vmo: self.vmo.as_ref().map(|vmo| vmo.dup()).transpose()?,
            ..*self
        })
    }

    /// Returns the mapping's start address.
    pub fn map_to_addr(&self) -> Vaddr {
        self.map_to_addr
    }

    /// Returns the mapping's end address.
    pub fn map_end(&self) -> Vaddr {
        self.map_to_addr + self.map_size.get()
    }

    /// Returns the mapping's size.
    pub fn map_size(&self) -> usize {
        self.map_size.get()
    }

    // Returns the permissions of pages in the mapping.
    pub fn perms(&self) -> VmPerms {
        self.perms
    }
}

/****************************** Page faults **********************************/

impl VmMapping {
    pub fn handle_page_fault(
        &self,
        vm_space: &VmSpace,
        page_fault_info: &PageFaultInfo,
    ) -> Result<()> {
        if !self.perms.contains(page_fault_info.required_perms) {
            trace!(
                "self.perms {:?}, page_fault_info.required_perms {:?}, self.range {:?}",
                self.perms,
                page_fault_info.required_perms,
                self.range()
            );
            return_errno_with_message!(Errno::EACCES, "perm check fails");
        }

        let address = page_fault_info.address;

        let page_aligned_addr = address.align_down(PAGE_SIZE);
        let is_write = page_fault_info.required_perms.contains(VmPerms::WRITE);

        if !is_write && self.vmo.is_some() && self.handle_page_faults_around {
            let res = self.handle_page_faults_around(vm_space, address);

            // Errors caused by the "around" pages should be ignored, so here we
            // only return the error if the faulting page is still not mapped.
            if res.is_err() {
                let preempt_guard = disable_preempt();
                let mut cursor = vm_space.cursor(
                    &preempt_guard,
                    &(page_aligned_addr..page_aligned_addr + PAGE_SIZE),
                )?;
                if let VmItem::Mapped { .. } = cursor.query().unwrap() {
                    return Ok(());
                }
            }

            return res;
        }

        'retry: loop {
            let preempt_guard = disable_preempt();
            let mut cursor = vm_space.cursor_mut(
                &preempt_guard,
                &(page_aligned_addr..page_aligned_addr + PAGE_SIZE),
            )?;

            match cursor.query().unwrap() {
                VmItem::Mapped {
                    va,
                    frame,
                    mut prop,
                } => {
                    if VmPerms::from(prop.flags).contains(page_fault_info.required_perms) {
                        // The page fault is already handled maybe by other threads.
                        // Just flush the TLB and return.
                        TlbFlushOp::Address(va).perform_on_current();
                        return Ok(());
                    }
                    assert!(is_write);
                    // Perform COW if it is a write access to a shared mapping.

                    // Skip if the page fault is already handled.
                    if prop.flags.contains(PageFlags::W) {
                        return Ok(());
                    }

                    // If the forked child or parent immediately unmaps the page after
                    // the fork without accessing it, we are the only reference to the
                    // frame. We can directly map the frame as writable without
                    // copying. In this case, the reference count of the frame is 2 (
                    // one for the mapping and one for the frame handle itself).
                    let only_reference = frame.reference_count() == 2;

                    let new_flags = PageFlags::W | PageFlags::ACCESSED | PageFlags::DIRTY;

                    if self.is_shared || only_reference {
                        cursor.protect_next(PAGE_SIZE, |p| p.flags |= new_flags);
                        cursor.flusher().issue_tlb_flush(TlbFlushOp::Address(va));
                        cursor.flusher().dispatch_tlb_flush();
                    } else {
                        let new_frame = duplicate_frame(&frame)?;
                        prop.flags |= new_flags;
                        cursor.map(new_frame.into(), prop);
                    }
                    cursor.flusher().sync_tlb_flush();
                }
                VmItem::NotMapped { .. } => {
                    // Map a new frame to the page fault address.
                    let (frame, is_readonly) = match self.prepare_page(address, is_write) {
                        Ok((frame, is_readonly)) => (frame, is_readonly),
                        Err(VmoCommitError::Err(e)) => return Err(e),
                        Err(VmoCommitError::NeedIo(index)) => {
                            drop(cursor);
                            drop(preempt_guard);
                            self.vmo
                                .as_ref()
                                .unwrap()
                                .commit_on(index, CommitFlags::empty())?;
                            continue 'retry;
                        }
                    };

                    let vm_perms = {
                        let mut perms = self.perms;
                        if is_readonly {
                            // COW pages are forced to be read-only.
                            perms -= VmPerms::WRITE;
                        }
                        perms
                    };

                    let mut page_flags = vm_perms.into();
                    page_flags |= PageFlags::ACCESSED;
                    if is_write {
                        page_flags |= PageFlags::DIRTY;
                    }
                    let map_prop = PageProperty::new(page_flags, CachePolicy::Writeback);

                    cursor.map(frame, map_prop);
                }
            }
            break 'retry;
        }
        Ok(())
    }

    fn prepare_page(
        &self,
        page_fault_addr: Vaddr,
        write: bool,
    ) -> core::result::Result<(UFrame, bool), VmoCommitError> {
        let mut is_readonly = false;
        let Some(vmo) = &self.vmo else {
            return Ok((FrameAllocOptions::new().alloc_frame()?.into(), is_readonly));
        };

        let page_offset = page_fault_addr.align_down(PAGE_SIZE) - self.map_to_addr;
        if !self.is_shared && page_offset >= vmo.size() {
            // The page index is outside the VMO. This is only allowed in private mapping.
            return Ok((FrameAllocOptions::new().alloc_frame()?.into(), is_readonly));
        }

        let page = vmo.get_committed_frame(page_offset)?;
        if !self.is_shared && write {
            // Write access to private VMO-backed mapping. Performs COW directly.
            Ok((duplicate_frame(&page)?.into(), is_readonly))
        } else {
            // Operations to shared mapping or read access to private VMO-backed mapping.
            // If read access to private VMO-backed mapping triggers a page fault,
            // the map should be readonly. If user next tries to write to the frame,
            // another page fault will be triggered which will performs a COW (Copy-On-Write).
            is_readonly = !self.is_shared;
            Ok((page, is_readonly))
        }
    }

    fn handle_page_faults_around(&self, vm_space: &VmSpace, page_fault_addr: Vaddr) -> Result<()> {
        const SURROUNDING_PAGE_NUM: usize = 16;
        const SURROUNDING_PAGE_ADDR_MASK: usize = !(SURROUNDING_PAGE_NUM * PAGE_SIZE - 1);

        let vmo = self.vmo.as_ref().unwrap();
        let around_page_addr = page_fault_addr & SURROUNDING_PAGE_ADDR_MASK;
        let size = min(vmo.size(), self.map_size.get());

        let mut start_addr = max(around_page_addr, self.map_to_addr);
        let end_addr = min(
            start_addr + SURROUNDING_PAGE_NUM * PAGE_SIZE,
            self.map_to_addr + size,
        );

        let vm_perms = self.perms - VmPerms::WRITE;
        'retry: loop {
            let preempt_guard = disable_preempt();
            let mut cursor = vm_space.cursor_mut(&preempt_guard, &(start_addr..end_addr))?;
            let operate =
                move |commit_fn: &mut dyn FnMut()
                    -> core::result::Result<UFrame, VmoCommitError>| {
                    if let VmItem::NotMapped { .. } = cursor.query().unwrap() {
                        // We regard all the surrounding pages as accessed, no matter
                        // if it is really so. Then the hardware won't bother to update
                        // the accessed bit of the page table on following accesses.
                        let page_flags = PageFlags::from(vm_perms) | PageFlags::ACCESSED;
                        let page_prop = PageProperty::new(page_flags, CachePolicy::Writeback);
                        let frame = commit_fn()?;
                        cursor.map(frame, page_prop);
                    } else {
                        let next_addr = cursor.virt_addr() + PAGE_SIZE;
                        if next_addr < end_addr {
                            let _ = cursor.jump(next_addr);
                        }
                    }
                    Ok(())
                };

            let start_offset = start_addr - self.map_to_addr;
            let end_offset = end_addr - self.map_to_addr;
            match vmo.try_operate_on_range(&(start_offset..end_offset), operate) {
                Ok(_) => return Ok(()),
                Err(VmoCommitError::NeedIo(index)) => {
                    drop(preempt_guard);
                    vmo.commit_on(index, CommitFlags::empty())?;
                    start_addr = index * PAGE_SIZE + self.map_to_addr;
                    continue 'retry;
                }
                Err(VmoCommitError::Err(e)) => return Err(e),
            }
        }
    }
}

/**************************** Transformations ********************************/

impl VmMapping {
    /// Enlarges the mapping by `extra_size` bytes to the high end.
    pub fn enlarge(self, extra_size: usize) -> Self {
        Self {
            map_size: NonZeroUsize::new(self.map_size.get() + extra_size).unwrap(),
            ..self
        }
    }

    /// Splits the mapping at the specified address.
    ///
    /// The address must be within the mapping and page-aligned. The address
    /// must not be either the start or the end of the mapping.
    fn split(self, at: Vaddr) -> Result<(Self, Self)> {
        debug_assert!(self.map_to_addr < at && at < self.map_end());
        debug_assert!(at % PAGE_SIZE == 0);

        let (mut l_vmo, mut r_vmo) = (None, None);

        if let Some(vmo) = self.vmo {
            let at_offset = vmo.range.start + at - self.map_to_addr;

            let l_range = vmo.range.start..at_offset;
            let r_range = at_offset..vmo.range.end;

            l_vmo = Some(MappedVmo::new(vmo.vmo.dup()?, l_range));
            r_vmo = Some(MappedVmo::new(vmo.vmo.dup()?, r_range));
        }

        let left_size = at - self.map_to_addr;
        let right_size = self.map_size.get() - left_size;
        let left = Self {
            map_to_addr: self.map_to_addr,
            map_size: NonZeroUsize::new(left_size).unwrap(),
            vmo: l_vmo,
            ..self
        };
        let right = Self {
            map_to_addr: at,
            map_size: NonZeroUsize::new(right_size).unwrap(),
            vmo: r_vmo,
            ..self
        };

        Ok((left, right))
    }

    /// Splits the mapping at the specified address.
    ///
    /// There are four conditions:
    /// 1. |-outside `range`-| + |------------within `range`------------|
    /// 2. |------------within `range`------------| + |-outside `range`-|
    /// 3. |-outside `range`-| + |-within `range`-| + |-outside `range`-|
    /// 4. |----------------------within `range` -----------------------|
    ///
    /// Returns (left outside, within, right outside) if successful.
    ///
    /// # Panics
    ///
    /// Panics if the mapping does not contain the range, or if the start or
    /// end of the range is not page-aligned.
    pub fn split_range(self, range: &Range<Vaddr>) -> Result<(Option<Self>, Self, Option<Self>)> {
        let mapping_range = self.range();
        if range.start <= mapping_range.start && mapping_range.end <= range.end {
            // Condition 4.
            return Ok((None, self, None));
        } else if mapping_range.start < range.start {
            let (left, within) = self.split(range.start).unwrap();
            if range.end < mapping_range.end {
                // Condition 3.
                let (within, right) = within.split(range.end).unwrap();
                return Ok((Some(left), within, Some(right)));
            } else {
                // Condition 1.
                return Ok((Some(left), within, None));
            }
        } else if mapping_range.contains(&range.end) {
            // Condition 2.
            let (within, right) = self.split(range.end).unwrap();
            return Ok((None, within, Some(right)));
        }
        panic!("The mapping does not contain the splitting range.");
    }
}

/************************** VM Space operations ******************************/

impl VmMapping {
    /// Unmaps the mapping from the VM space.
    pub(super) fn unmap(self, vm_space: &VmSpace) -> Result<()> {
        let preempt_guard = disable_preempt();
        let range = self.range();
        let mut cursor = vm_space.cursor_mut(&preempt_guard, &range)?;

        cursor.unmap(range.len());
        cursor.flusher().dispatch_tlb_flush();
        cursor.flusher().sync_tlb_flush();

        Ok(())
    }

    /// Change the perms of the mapping.
    pub(super) fn protect(self, vm_space: &VmSpace, perms: VmPerms) -> Self {
        let preempt_guard = disable_preempt();
        let range = self.range();
        let mut cursor = vm_space.cursor_mut(&preempt_guard, &range).unwrap();

        let op = |p: &mut PageProperty| p.flags = perms.into();
        while cursor.virt_addr() < range.end {
            if let Some(va) = cursor.protect_next(range.end - cursor.virt_addr(), op) {
                cursor.flusher().issue_tlb_flush(TlbFlushOp::Range(va));
            } else {
                break;
            }
        }
        cursor.flusher().dispatch_tlb_flush();
        cursor.flusher().sync_tlb_flush();

        Self { perms, ..self }
    }
}

/// A wrapper that represents a mapped [`Vmo`] and provide required functionalities
/// that need to be provided to mappings from the VMO.
#[derive(Debug)]
pub(super) struct MappedVmo {
    vmo: Vmo,
    /// Represents the accessible range in the VMO for mappings.
    range: Range<usize>,
}

impl MappedVmo {
    /// Creates a `MappedVmo` used for mapping.
    pub(super) fn new(vmo: Vmo, range: Range<usize>) -> Self {
        Self { vmo, range }
    }

    fn size(&self) -> usize {
        self.range.len()
    }

    /// Gets the committed frame at the input offset in the mapped VMO.
    ///
    /// If the VMO has not committed a frame at this index, it will commit
    /// one first and return it. If the commit operation needs to perform I/O,
    /// it will return a [`VmoCommitError::NeedIo`].
    fn get_committed_frame(
        &self,
        page_offset: usize,
    ) -> core::result::Result<UFrame, VmoCommitError> {
        debug_assert!(page_offset < self.range.len());
        debug_assert!(page_offset % PAGE_SIZE == 0);
        self.vmo.try_commit_page(self.range.start + page_offset)
    }

    /// Commits a page at a specific page index.
    ///
    /// This method may involve I/O operations if the VMO needs to fecth
    /// a page from the underlying page cache.
    pub fn commit_on(&self, page_idx: usize, commit_flags: CommitFlags) -> Result<UFrame> {
        debug_assert!(page_idx * PAGE_SIZE < self.range.len());
        self.vmo.commit_on(page_idx, commit_flags)
    }

    /// Traverses the indices within a specified range of a VMO sequentially.
    ///
    /// For each index position, you have the option to commit the page as well as
    /// perform other operations.
    ///
    /// Once a commit operation needs to perform I/O, it will return a [`VmoCommitError::NeedIo`].
    fn try_operate_on_range<F>(
        &self,
        range: &Range<usize>,
        operate: F,
    ) -> core::result::Result<(), VmoCommitError>
    where
        F: FnMut(
            &mut dyn FnMut() -> core::result::Result<UFrame, VmoCommitError>,
        ) -> core::result::Result<(), VmoCommitError>,
    {
        debug_assert!(range.start < self.range.len());
        debug_assert!(range.end <= self.range.len());

        let range = self.range.start + range.start..self.range.start + range.end;

        self.vmo.try_operate_on_range(&range, operate)
    }

    /// Duplicates the capability.
    pub fn dup(&self) -> Result<Self> {
        Ok(Self {
            vmo: self.vmo.dup()?,
            range: self.range.clone(),
        })
    }
}
