// SPDX-License-Identifier: GPL-2.0

// Copyright (C) 2025 Google LLC.

use kernel::{page::PAGE_SIZE, prelude::*, seq_file::SeqFile, task::Pid};

mod tree;
use self::tree::{FromArrayAllocs, ReserveNewTreeAlloc, TreeRangeAllocator};

mod array;
use self::array::{ArrayRangeAllocator, EmptyArrayAlloc};

enum DescriptorState<T> {
    Reserved(Reservation),
    Allocated(Allocation<T>),
}

impl<T> DescriptorState<T> {
    fn new(is_oneway: bool, debug_id: usize, pid: Pid) -> Self {
        DescriptorState::Reserved(Reservation {
            debug_id,
            is_oneway,
            pid,
        })
    }

    fn pid(&self) -> Pid {
        match self {
            DescriptorState::Reserved(inner) => inner.pid,
            DescriptorState::Allocated(inner) => inner.reservation.pid,
        }
    }

    fn is_oneway(&self) -> bool {
        match self {
            DescriptorState::Reserved(inner) => inner.is_oneway,
            DescriptorState::Allocated(inner) => inner.reservation.is_oneway,
        }
    }
}

#[derive(Clone)]
struct Reservation {
    debug_id: usize,
    is_oneway: bool,
    pid: Pid,
}

impl Reservation {
    fn allocate<T>(self, data: Option<T>) -> Allocation<T> {
        Allocation {
            data,
            reservation: self,
        }
    }
}

struct Allocation<T> {
    reservation: Reservation,
    data: Option<T>,
}

impl<T> Allocation<T> {
    fn deallocate(self) -> (Reservation, Option<T>) {
        (self.reservation, self.data)
    }

    fn debug_id(&self) -> usize {
        self.reservation.debug_id
    }

    fn take(&mut self) -> Option<T> {
        self.data.take()
    }
}

/// The array implementation must switch to the tree if it wants to go beyond this number of
/// ranges.
const TREE_THRESHOLD: usize = 8;

/// Represents a range of pages that have just become completely free.
#[derive(Copy, Clone)]
pub(crate) struct FreedRange {
    pub(crate) start_page_idx: usize,
    pub(crate) end_page_idx: usize,
}

impl FreedRange {
    fn interior_pages(offset: usize, size: usize) -> FreedRange {
        FreedRange {
            // Divide round up
            start_page_idx: offset.div_ceil(PAGE_SIZE),
            // Divide round down
            end_page_idx: (offset + size) / PAGE_SIZE,
        }
    }
}

struct Range<T> {
    offset: usize,
    size: usize,
    state: DescriptorState<T>,
}

impl<T> Range<T> {
    fn endpoint(&self) -> usize {
        self.offset + self.size
    }
}

pub(crate) struct RangeAllocator<T> {
    inner: Impl<T>,
}

enum Impl<T> {
    Empty(usize),
    Array(ArrayRangeAllocator<T>),
    Tree(TreeRangeAllocator<T>),
}

impl<T> RangeAllocator<T> {
    pub(crate) fn new(size: usize) -> Self {
        Self {
            inner: Impl::Empty(size),
        }
    }

    pub(crate) fn free_oneway_space(&self) -> usize {
        match &self.inner {
            Impl::Empty(size) => size / 2,
            Impl::Array(array) => array.free_oneway_space(),
            Impl::Tree(tree) => tree.free_oneway_space(),
        }
    }

    pub(crate) fn count_buffers(&self) -> usize {
        match &self.inner {
            Impl::Empty(_size) => 0,
            Impl::Array(array) => array.count_buffers(),
            Impl::Tree(tree) => tree.count_buffers(),
        }
    }

    pub(crate) fn debug_print(&self, m: &SeqFile) -> Result<()> {
        match &self.inner {
            Impl::Empty(_size) => Ok(()),
            Impl::Array(array) => array.debug_print(m),
            Impl::Tree(tree) => tree.debug_print(m),
        }
    }

    /// Try to reserve a new buffer, using the provided allocation if necessary.
    pub(crate) fn reserve_new(&mut self, mut args: ReserveNewArgs<T>) -> Result<ReserveNew<T>> {
        match &mut self.inner {
            Impl::Empty(size) => {
                let empty_array = match args.empty_array_alloc.take() {
                    Some(empty_array) => ArrayRangeAllocator::new(*size, empty_array),
                    None => {
                        return Ok(ReserveNew::NeedAlloc(ReserveNewNeedAlloc {
                            args,
                            need_empty_array_alloc: true,
                            need_new_tree_alloc: false,
                            need_tree_alloc: false,
                        }))
                    }
                };

                self.inner = Impl::Array(empty_array);
                self.reserve_new(args)
            }
            Impl::Array(array) if array.is_full() => {
                let allocs = match args.new_tree_alloc {
                    Some(ref mut allocs) => allocs,
                    None => {
                        return Ok(ReserveNew::NeedAlloc(ReserveNewNeedAlloc {
                            args,
                            need_empty_array_alloc: false,
                            need_new_tree_alloc: true,
                            need_tree_alloc: true,
                        }))
                    }
                };

                let new_tree =
                    TreeRangeAllocator::from_array(array.total_size(), &mut array.ranges, allocs);

                self.inner = Impl::Tree(new_tree);
                self.reserve_new(args)
            }
            Impl::Array(array) => {
                let offset =
                    array.reserve_new(args.debug_id, args.size, args.is_oneway, args.pid)?;
                Ok(ReserveNew::Success(ReserveNewSuccess {
                    offset,
                    oneway_spam_detected: false,
                    _empty_array_alloc: args.empty_array_alloc,
                    _new_tree_alloc: args.new_tree_alloc,
                    _tree_alloc: args.tree_alloc,
                }))
            }
            Impl::Tree(tree) => {
                let alloc = match args.tree_alloc {
                    Some(alloc) => alloc,
                    None => {
                        return Ok(ReserveNew::NeedAlloc(ReserveNewNeedAlloc {
                            args,
                            need_empty_array_alloc: false,
                            need_new_tree_alloc: false,
                            need_tree_alloc: true,
                        }));
                    }
                };
                let (offset, oneway_spam_detected) =
                    tree.reserve_new(args.debug_id, args.size, args.is_oneway, args.pid, alloc)?;
                Ok(ReserveNew::Success(ReserveNewSuccess {
                    offset,
                    oneway_spam_detected,
                    _empty_array_alloc: args.empty_array_alloc,
                    _new_tree_alloc: args.new_tree_alloc,
                    _tree_alloc: None,
                }))
            }
        }
    }

    /// Deletes the allocations at `offset`.
    pub(crate) fn reservation_abort(&mut self, offset: usize) -> Result<FreedRange> {
        match &mut self.inner {
            Impl::Empty(_size) => Err(EINVAL),
            Impl::Array(array) => array.reservation_abort(offset),
            Impl::Tree(tree) => {
                let freed_range = tree.reservation_abort(offset)?;
                if tree.is_empty() {
                    self.inner = Impl::Empty(tree.total_size());
                }
                Ok(freed_range)
            }
        }
    }

    /// Called when an allocation is no longer in use by the kernel.
    ///
    /// The value in `data` will be stored, if any. A mutable reference is used to avoid dropping
    /// the `T` when an error is returned.
    pub(crate) fn reservation_commit(&mut self, offset: usize, data: &mut Option<T>) -> Result {
        match &mut self.inner {
            Impl::Empty(_size) => Err(EINVAL),
            Impl::Array(array) => array.reservation_commit(offset, data),
            Impl::Tree(tree) => tree.reservation_commit(offset, data),
        }
    }

    /// Called when the kernel starts using an allocation.
    ///
    /// Returns the size of the existing entry and the data associated with it.
    pub(crate) fn reserve_existing(&mut self, offset: usize) -> Result<(usize, usize, Option<T>)> {
        match &mut self.inner {
            Impl::Empty(_size) => Err(EINVAL),
            Impl::Array(array) => array.reserve_existing(offset),
            Impl::Tree(tree) => tree.reserve_existing(offset),
        }
    }

    /// Call the provided callback at every allocated region.
    ///
    /// This destroys the range allocator. Used only during shutdown.
    pub(crate) fn take_for_each<F: Fn(usize, usize, usize, Option<T>)>(&mut self, callback: F) {
        match &mut self.inner {
            Impl::Empty(_size) => {}
            Impl::Array(array) => array.take_for_each(callback),
            Impl::Tree(tree) => tree.take_for_each(callback),
        }
    }
}

/// The arguments for `reserve_new`.
#[derive(Default)]
pub(crate) struct ReserveNewArgs<T> {
    pub(crate) size: usize,
    pub(crate) is_oneway: bool,
    pub(crate) debug_id: usize,
    pub(crate) pid: Pid,
    pub(crate) empty_array_alloc: Option<EmptyArrayAlloc<T>>,
    pub(crate) new_tree_alloc: Option<FromArrayAllocs<T>>,
    pub(crate) tree_alloc: Option<ReserveNewTreeAlloc<T>>,
}

/// The return type of `ReserveNew`.
pub(crate) enum ReserveNew<T> {
    Success(ReserveNewSuccess<T>),
    NeedAlloc(ReserveNewNeedAlloc<T>),
}

/// Returned by `reserve_new` when the reservation was successul.
pub(crate) struct ReserveNewSuccess<T> {
    pub(crate) offset: usize,
    pub(crate) oneway_spam_detected: bool,

    // If the user supplied an allocation that we did not end up using, then we return it here.
    // The caller will kfree it outside of the lock.
    _empty_array_alloc: Option<EmptyArrayAlloc<T>>,
    _new_tree_alloc: Option<FromArrayAllocs<T>>,
    _tree_alloc: Option<ReserveNewTreeAlloc<T>>,
}

/// Returned by `reserve_new` to request the caller to make an allocation before calling the method
/// again.
pub(crate) struct ReserveNewNeedAlloc<T> {
    args: ReserveNewArgs<T>,
    need_empty_array_alloc: bool,
    need_new_tree_alloc: bool,
    need_tree_alloc: bool,
}

impl<T> ReserveNewNeedAlloc<T> {
    /// Make the necessary allocations for another call to `reserve_new`.
    pub(crate) fn make_alloc(mut self) -> Result<ReserveNewArgs<T>> {
        if self.need_empty_array_alloc && self.args.empty_array_alloc.is_none() {
            self.args.empty_array_alloc = Some(EmptyArrayAlloc::try_new(TREE_THRESHOLD)?);
        }
        if self.need_new_tree_alloc && self.args.new_tree_alloc.is_none() {
            self.args.new_tree_alloc = Some(FromArrayAllocs::try_new(TREE_THRESHOLD)?);
        }
        if self.need_tree_alloc && self.args.tree_alloc.is_none() {
            self.args.tree_alloc = Some(ReserveNewTreeAlloc::try_new()?);
        }
        Ok(self.args)
    }
}