// SPDX-License-Identifier: MIT OR Apache-2.0
// Copyright (c) 2025 lacklustr@protonmail.com https://github.com/eadf

use std::marker::PhantomData;
use vob::Vob;

/// Trait for types that can be used as pool indices
pub(crate) trait PoolIndex: Copy + From<usize> + Into<usize> + Ord {}

#[derive(Debug, PartialEq, Clone, Copy)]
pub(crate) enum Allocation<T> {
    Reused(T),
    New(T),
}

mod sealed {
    use super::PoolIndex;
    use std::cmp::Reverse;
    use std::collections::BinaryHeap;

    /// Internal trait for handling the available indices container
    pub(crate) trait AvailableHandler<I: PoolIndex>: Clone {
        fn new() -> Self;
        fn pop(&mut self) -> Option<I>;
        fn push(&mut self, index: I);
        fn len(&self) -> usize;
        fn clear(&mut self);
        fn sort_available(&mut self);
        fn pop_smallest_available(&mut self) -> Option<I>;
    }

    /// Policy trait that determines pool behavior
    pub(crate) trait IndexPolicy<I: PoolIndex> {
        type Handler: AvailableHandler<I>;
    }

    // LIFO Handler (Vec-based)
    #[derive(Clone)]
    pub struct LifoHandler<I: PoolIndex> {
        indices: Vec<I>,
    }

    impl<I: PoolIndex> AvailableHandler<I> for LifoHandler<I> {
        #[inline(always)]
        fn new() -> Self {
            Self {
                indices: Vec::new(),
            }
        }

        #[inline(always)]
        fn pop(&mut self) -> Option<I> {
            self.indices.pop()
        }

        #[inline(always)]
        fn push(&mut self, index: I) {
            self.indices.push(index);
        }

        #[inline(always)]
        fn len(&self) -> usize {
            self.indices.len()
        }

        #[inline(always)]
        fn clear(&mut self) {
            self.indices.clear();
        }

        #[inline(always)]
        fn sort_available(&mut self) {
            self.indices.sort_unstable_by(|a, b| b.cmp(a));
        }

        #[inline(always)]
        fn pop_smallest_available(&mut self) -> Option<I> {
            self.indices.pop()
        }
    }

    // SIFO Handler (BinaryHeap-based)
    #[derive(Clone)]
    pub struct SifoHandler<I: PoolIndex> {
        indices: BinaryHeap<Reverse<I>>,
    }

    impl<I: PoolIndex> AvailableHandler<I> for SifoHandler<I> {
        #[inline(always)]
        fn new() -> Self {
            Self {
                indices: BinaryHeap::new(),
            }
        }

        #[inline(always)]
        fn pop(&mut self) -> Option<I> {
            self.indices.pop().map(|r| r.0)
        }

        #[inline(always)]
        fn push(&mut self, index: I) {
            self.indices.push(Reverse(index));
        }

        #[inline(always)]
        fn len(&self) -> usize {
            self.indices.len()
        }

        #[inline(always)]
        fn clear(&mut self) {
            self.indices.clear();
        }

        #[inline(always)]
        fn sort_available(&mut self) {}

        #[inline(always)]
        fn pop_smallest_available(&mut self) -> Option<I> {
            self.indices.pop().map(|r| r.0)
        }
    }
}

use sealed::{AvailableHandler, IndexPolicy};

#[derive(Clone)]
/// Generic index pool with configurable policy
pub(crate) struct IndexPool<I: PoolIndex, P: IndexPolicy<I>, const ENABLE_UNSAFE: bool> {
    used: Vob,
    available: P::Handler,
    _phantom: PhantomData<I>,
}

impl<I: PoolIndex, P: IndexPolicy<I>, const ENABLE_UNSAFE: bool> IndexPool<I, P, ENABLE_UNSAFE> {
    /// Initiates the pool with densely packed `in_use` number of elements.
    pub(crate) fn new(already_in_use: u32) -> Self {
        Self {
            used: Vob::from_elem(true, already_in_use as usize),
            available: P::Handler::new(),
            _phantom: PhantomData,
        }
    }

    /// Deconstruct the container by taking the `used` field.
    /// Note that the collection is not working as intended after this operation
    pub(crate) fn take_used(&mut self) -> Vob {
        std::mem::take(&mut self.used)
    }

    /// Returns an index from the pool or creates a new one if the pool is empty
    pub(crate) fn pop(&mut self) -> Allocation<I> {
        if let Some(index) = self.available.pop() {
            if ENABLE_UNSAFE {
                unsafe {
                    let _ = self.used.set_unchecked(index.into(), true);
                }
            } else {
                let _ = self.used.set(index.into(), true);
            }
            Allocation::Reused(index)
        } else {
            let index = I::from(self.used.len());
            self.used.push(true);
            Allocation::New(index)
        }
    }

    /// Release an index back to the pool
    pub(crate) fn push(&mut self, index: I) {
        debug_assert!(index.into() < self.used.len(), "Invalid index");
        if self.used.get(index.into()).unwrap_or(false) {
            let _ = self.used.set(index.into(), false);
            self.available.push(index);
        } else {
            debug_assert!(false, "Index already released");
        }
    }

    /// The number of indices available for use, but not in active use (i.e. pool size)
    #[inline(always)]
    pub(crate) fn available_count(&self) -> u32 {
        self.available.len() as u32
    }

    /// The number of active indices + available in the pool
    #[inline(always)]
    pub(crate) fn total_count(&self) -> u32 {
        self.used.len() as u32
    }

    /// The number of active indices
    #[inline(always)]
    pub(crate) fn active_count(&self) -> u32 {
        debug_assert!(self.total_count() >= self.available_count());
        self.total_count() - self.available_count()
    }
    #[allow(dead_code)]
    #[inline(always)]
    pub(crate) fn is_used(&self, index: I) -> bool {
        unsafe { self.used.get_unchecked(index.into()) }
    }

    /// Resets the pool so that the first 'new_size' elements are marked as used,
    /// and the available list is cleared.
    pub(crate) fn truncate(&mut self, new_size: usize) {
        self.used.resize(new_size, true);
        self.used.set_all(true);
        self.available.clear();
    }

    /// Defragment the pool by moving items from high indices into low index holes.
    ///
    /// This compacts all active indices to be contiguous from 0..active_count,
    /// eliminating any holes in the middle. The callback is invoked for each move
    /// to allow the caller to update any references to the moved indices.
    ///
    /// After defragmentation, the pool is truncated to active_count size.
    ///
    /// # Arguments
    /// * `remap` - Callback invoked as `remap(from_index, to_index)` for each move
    pub(crate) fn defragment<F: FnMut(I, I)>(&mut self, mut remap: F) {
        let old_size = self.active_count();

        // Sort available indices so we process holes from lowest to highest
        self.available.sort_available();

        let mut move_candidate = self.used.len().saturating_sub(1);

        while let Some(hole_idx) = self.available.pop_smallest_available() {
            let hole_usize = hole_idx.into();

            // Find the last used index (skip any holes at the end)
            while move_candidate > hole_usize && !unsafe { self.used.get_unchecked(move_candidate) }
            {
                move_candidate -= 1;
            }

            // If hole is beyond active range, we're done
            if hole_usize >= move_candidate {
                break;
            }

            // Notify caller to remap references from move_candidate to hole_idx
            remap(I::from(move_candidate), hole_idx);

            move_candidate -= 1;
        }

        // Reset pool: all items are now densely packed
        self.truncate(old_size as usize);
    }
}

// ===== Policy Implementations =====

#[derive(Clone)]
/// Last In First Out policy - uses Vec for O(1) pop/push
pub(crate) struct Lifo;

#[derive(Clone)]
/// Smallest Index First Out policy - uses BinaryHeap for smallest index first
pub(crate) struct Sifo;

// Policy trait implementations
impl<I: PoolIndex> IndexPolicy<I> for Lifo {
    type Handler = sealed::LifoHandler<I>;
}

impl<I: PoolIndex> IndexPolicy<I> for Sifo {
    type Handler = sealed::SifoHandler<I>;
}

// ===== Type Aliases for Convenience =====

/// A fast 'Last In First Out' index pool
pub(crate) type LifoPool<I, const ENABLE_UNSAFE: bool> = IndexPool<I, Lifo, ENABLE_UNSAFE>;

/// A 'Smallest Index First Out' index pool
pub(crate) type SifoPool<I, const ENABLE_UNSAFE: bool> = IndexPool<I, Sifo, ENABLE_UNSAFE>;
