kasino 0.1.0

use core::marker::PhantomData;

use crate::{
    Collection,
    Signature,
    components::PushPopCollection,
    storage::StorageBackend,
    strategy::{Hooked, Strategy},
};

pub(crate) const DEFAULT_QUEUE_CAP: usize = 32;

#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Hash)]
pub(crate) struct BanditCore<Q, S, B, C, const SUB_CAP: usize = DEFAULT_QUEUE_CAP> {
    strategy: S,
    sub_collections: B,
    collection_state: C,
    _p: PhantomData<Q>,
}

impl<Q, S, B, C, const SUB_CAP: usize> BanditCore<Q, S, B, C, SUB_CAP>
where
    S: Default,
{
    pub(crate) fn new_with(queues: B, states: C) -> Self {
        Self {
            strategy: S::default(),
            sub_collections: queues,
            collection_state: states,
            _p: PhantomData,
        }
    }
}

impl<Q, S, B, C, const SUB_CAP: usize> BanditCore<Q, S, B, C, SUB_CAP>
where
    B: StorageBackend<Q>,
{
    /// returns the number of sub collections
    pub(crate) fn arm_count(&self) -> usize {
        self.sub_collections.len()
    }
}

impl<Q, S, B, C, const SUB_CAP: usize> BanditCore<Q, S, B, C, SUB_CAP>
where
    S: Strategy<Q>,
    Q: Collection,
{
    pub(crate) fn buy_in(&self) -> BanditHandle<'_, Q, S, B, C, SUB_CAP> {
        BanditHandle {
            parent: self,
            gambler: self.strategy.create_gambler(),
        }
    }
}

impl<Q, S, B, C, const SUB_CAP: usize> BanditCore<Q, S, B, C, SUB_CAP>
where
    B: StorageBackend<Q>,
{
    pub(crate) fn into_arms(self) -> impl Iterator<Item = B::Item> {
        self.sub_collections.into_iter()
    }
}

impl<Q, S, B, C, const SUB_CAP: usize> BanditCore<Q, S, B, C, SUB_CAP>
where
    Q: IntoIterator,
    B: IntoIterator<Item = Q>,
{
    pub(crate) fn into_items(self) -> impl Iterator<Item = Q::Item> {
        self.sub_collections
            .into_iter()
            .flat_map(|collection| collection.into_iter())
    }
}

/// An owned handle into the core bandit.
///
/// This handle provides access to the functionality of the wrapped [`Collection`].
#[must_use]
#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Hash)]
pub struct BanditHandle<
    'a,
    Q: Collection,
    S: Strategy<Q>,
    B,
    C,
    const SUB_CAP: usize = DEFAULT_QUEUE_CAP,
> {
    parent: &'a BanditCore<Q, S, B, C, SUB_CAP>,
    gambler: S::Gambler,
}

impl<'a, Q, S, B, C, const SUB_CAP: usize> BanditHandle<'a, Q, S, B, C, SUB_CAP>
where
    Q: Collection,
    S: Strategy<Q>,
    B: StorageBackend<Q>,
    C: StorageBackend<<S::Gambler as Hooked>::Stake>,
{
    /// Fork this handle into a new one
    #[inline]
    pub fn fork(&mut self) -> Self {
        Self {
            parent: self.parent,
            gambler: S::fork_gambler(&self.parent.strategy, &mut self.gambler),
        }
    }

    /// Make a call to [`Collection::offer`] to an arms as chosen by this handles gambler.
    #[inline]
    pub fn offer<'b, 'c>(
        &'c mut self,
        item: <Q::OfferSignature as Signature>::Input<'b>,
    ) -> Result<
        <Q::OfferSignature as Signature>::Output<'b, 'c>,
        <Q::OfferSignature as Signature>::Error<'b, 'c>,
    > {
        let i = self
            .parent
            .strategy
            .choose_offer_arm(&self.parent.collection_state, &mut self.gambler);
        match self.parent.sub_collections[i].offer(item) {
            Ok(r) => {
                self.gambler.on_offer_succ(&self.parent.collection_state[i]);
                Ok(r)
            }
            Err(e) => {
                self.gambler.on_offer_fail(&self.parent.collection_state[i]);
                Err(e)
            }
        }
    }

    /// Make a call to [`Collection::poll`] to an arm as chosen by this handles gambler.
    ///
    /// If the call fails, [`Strategy::collect`] may be called to ensure consistency across all arms.
    #[inline]
    pub fn poll<'b, 'c>(
        &'c mut self,
        input: <Q::PollSignature as Signature>::Input<'b>,
    ) -> Result<
        <Q::PollSignature as Signature>::Output<'b, 'c>,
        <Q::PollSignature as Signature>::Error<'b, 'c>,
    > {
        Self::poll_internal(self.parent, &mut self.gambler, input).0
    }

    /// Makes a call to [`Self::poll`] and returns the stake associated with the arm we pulled.
    ///
    /// On failure returns the info associated with the arm originally pulled.
    #[expect(clippy::type_complexity)]
    #[inline]
    pub fn poll_with_info<'b, 'c>(
        &'c mut self,
        input: <Q::PollSignature as Signature>::Input<'b>,
    ) -> (
        Result<
            <Q::PollSignature as Signature>::Output<'b, 'c>,
            <Q::PollSignature as Signature>::Error<'b, 'c>,
        >,
        <S::Gambler as Hooked>::Stake,
    )
    where
        <S::Gambler as Hooked>::Stake: Clone,
    {
        let (res, idx) = Self::poll_internal(self.parent, &mut self.gambler, input);
        (res, self.parent.collection_state[idx].clone())
    }

    /// Makes a call to [`Self::poll`] and returns the index associated with the arm we pulled.
    #[expect(clippy::type_complexity)]
    pub(crate) fn poll_internal<'b, 'c>(
        parent: &'c BanditCore<Q, S, B, C, SUB_CAP>,
        gambler: &mut S::Gambler,
        input: <Q::PollSignature as Signature>::Input<'b>,
    ) -> (
        Result<
            <Q::PollSignature as Signature>::Output<'b, 'c>,
            <Q::PollSignature as Signature>::Error<'b, 'c>,
        >,
        usize,
    ) {
        let i = parent
            .strategy
            .choose_poll_arm(&parent.collection_state, gambler);
        match parent.sub_collections[i].poll(input) {
            Ok(r) => {
                gambler.on_poll_succ(&parent.collection_state[i]);
                (Ok(r), i)
            }
            Err(e) => {
                gambler.on_poll_fail(&parent.collection_state[i]);
                let r = parent.strategy.collect(
                    &parent.collection_state,
                    &parent.sub_collections,
                    input,
                );
                if let Some((r, state)) = r {
                    gambler.on_poll_succ(&parent.collection_state[state]);
                    (Ok(r), state)
                } else {
                    (Err(e), i)
                }
            }
        }
    }

    /// Returns an iterator over all stakes in all arms
    #[inline]
    pub fn state(&self) -> impl Iterator<Item = &<S::Gambler as Hooked>::Stake> {
        self.parent.collection_state.iter()
    }

    /// the total len of all arms
    #[inline]
    pub fn len(&self) -> usize {
        self.parent.sub_collections.iter().map(|q| q.len()).sum()
    }

    /// the total capacity of all arms
    #[inline]
    pub fn capacity(&self) -> usize {
        self.parent
            .sub_collections
            .iter()
            .map(|q| q.capacity())
            .sum()
    }

    /// are all arms empty?
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

impl<'a, Q, S, B, C, const SUB_CAP: usize> BanditHandle<'a, Q, S, B, C, SUB_CAP>
where
    B: StorageBackend<Q>,
    S: Strategy<Q>,
    Q: Collection,
{
    /// returns the number of arms
    #[inline]
    pub fn arm_count(&self) -> usize {
        self.parent.arm_count()
    }
}

impl<'a, Q, S, B, C, const SUB_CAP: usize> BanditHandle<'a, Q, S, B, C, SUB_CAP>
where
    Q: PushPopCollection,
    S: Strategy<Q>,
    B: StorageBackend<Q>,
    C: StorageBackend<<S::Gambler as Hooked>::Stake>,
{
    /// Pushes an item to the collection.
    ///
    /// Returns the item on an erorr.
    ///
    /// This method is a convenience wrapper around [`Self::offer`].
    #[inline]
    pub fn push(&mut self, item: Q::Item) -> Result<(), Q::Item> {
        self.offer(item)
    }

    /// Attempts to pop an item from the collection.
    ///
    /// This method is a convenience wrapper around [`Self::poll`].
    #[inline]
    pub fn pop(&mut self) -> Option<Q::Item> {
        self.poll(()).ok()
    }
}