use crate::bitset::{Bitset, SubregionStatus, UnexpectedBitsetState};
use contracts::requires;
use core::ptr::NonNull;

/// A single region of the allocator. See the comment on the `crate::allocators::buddy` module for
/// more information.
///
/// This type is valid when zero-initialized.
#[derive(Debug)]
pub struct Tree<const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize> {
    /// The base address of the tree.
    pub base_ptr: Option<NonNull<[u8; PAGE_SIZE]>>,

    /// The log2 of the number of pages in the region represented by the tree.
    pub size_class: usize,

    /// The offset in the bitset to the bits responsible for this tree's pages.
    pub bitset_offset: usize,
}

impl<const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize> Tree<PAGE_SIZE, PAGE_SIZE_BITS> {
    /// Returns the base address of the tree.
    #[requires(self.base_ptr.is_some())]
    pub fn base_addr(&self) -> usize {
        self.base_ptr.unwrap().as_ptr() as usize
    }

    /// Reads a bit from the bitset.
    #[requires(size_class <= self.size_class)]
    #[requires(offset_bytes < (PAGE_SIZE << (self.size_class + 1)))]
    #[requires(offset_bytes.trailing_zeros() as usize >= PAGE_SIZE_BITS + size_class)]
    pub fn bitset_get(
        &self,
        bitset: &Bitset,
        size_class: usize,
        offset_bytes: usize,
    ) -> SubregionStatus {
        bitset.get(self.bitset_index(size_class, offset_bytes))
    }

    /// Returns the index of a bit in the bitset that corresponds to the given size class and
    /// offset.
    #[requires(size_class <= self.size_class)]
    #[requires(offset_bytes < (PAGE_SIZE << (self.size_class + 1)))]
    #[requires(offset_bytes.trailing_zeros() as usize >= PAGE_SIZE_BITS + size_class)]
    fn bitset_index(&self, size_class: usize, offset_bytes: usize) -> usize {
        // We store the largest size classes first in the bitset. Count how many we are away from
        // the largest.
        let skipped_size_classes = self.size_class - size_class;
        let bits_skipped_for_size_class = (1 << skipped_size_classes) - 1;

        // Next, find our index in the size class.
        let bits_skipped_for_index = offset_bytes >> (PAGE_SIZE_BITS + size_class);

        // The sum of those two with our offset is simply our index.
        self.bitset_offset + bits_skipped_for_size_class + bits_skipped_for_index
    }

    /// Changes a bit in the bitset from `InFreeList` to `NotInFreeList`.
    #[requires(size_class <= self.size_class)]
    #[requires(offset_bytes < (PAGE_SIZE << (self.size_class + 1)))]
    #[requires(offset_bytes.trailing_zeros() as usize >= PAGE_SIZE_BITS + size_class)]
    pub fn bitset_mark_as_absent(
        &self,
        bitset: &mut Bitset,
        size_class: usize,
        offset_bytes: usize,
    ) -> Result<(), UnexpectedBitsetState> {
        bitset.replace(
            self.bitset_index(size_class, offset_bytes),
            SubregionStatus::InFreeList,
            SubregionStatus::NotInFreeList,
        )
    }

    /// Changes a bit in the bitset from `NotInFreeList` to `InFreeList`.
    #[requires(size_class <= self.size_class)]
    #[requires(offset_bytes < (PAGE_SIZE << (self.size_class + 1)))]
    #[requires(offset_bytes.trailing_zeros() as usize >= PAGE_SIZE_BITS + size_class)]
    pub fn bitset_mark_as_present(
        &self,
        bitset: &mut Bitset,
        size_class: usize,
        offset_bytes: usize,
    ) -> Result<(), UnexpectedBitsetState> {
        bitset.replace(
            self.bitset_index(size_class, offset_bytes),
            SubregionStatus::NotInFreeList,
            SubregionStatus::InFreeList,
        )
    }

    /// Returns whether the tree contains an address.
    #[requires(self.base_ptr.is_some())]
    pub fn contains(&self, addr: usize) -> bool {
        let tree_addr_lo = self.base_addr();
        let tree_addr_hi = tree_addr_lo + (PAGE_SIZE << self.size_class);
        (tree_addr_lo..tree_addr_hi).contains(&addr)
    }
}