twgame 0.11.0

use slotmap::SecondaryMap;

use crate::entities::SpawnOrder;
use crate::ids::PlayerUid;
use crate::state::Tees;

#[derive(Debug)]
/// Entities and Tees need to access all Tees in the game world. Tees never need to access the
/// entities. This struct allows getting tees in either player_id or spawn order.
/// This struct is about caching the iterators and is the only housekeeping about available Tees
/// in the team within the accessible game state. It needs to be updated on Tee join/leave.
/// For creating snapshots it is also necessary to have a TeeUid -> PlayerUid lookup.
pub struct TeeOrder {
    /// tees in Id order TODO: use smallvec or similar to optimize heap allocation for each (mostly 1-2 tee sized) team
    id: SecondaryMap<PlayerUid, ()>,
    /// tees in Spawn order
    spawn: Vec<(PlayerUid, SpawnOrder)>,
}

impl TeeOrder {
    pub(super) fn new() -> Self {
        Self {
            id: Default::default(),
            spawn: vec![],
        }
    }
    pub(super) fn reset(&mut self) {
        self.id.clear();
        self.spawn.clear();
    }
}

impl TeeOrder {
    pub(super) fn add_tee(&mut self, player_uid: PlayerUid, spawn_order: SpawnOrder) {
        // append to spawn order
        self.spawn.push((player_uid, spawn_order));
        self.id.insert(player_uid, ());
    }

    pub(super) fn add_tee_from_other_team(
        &mut self,
        player_uid: PlayerUid,
        spawn_order: SpawnOrder,
    ) {
        // do binary search of id order
        // unwrap_err, because tee mit same exact SpawnOrder shouldn't exist
        let pos = self
            .spawn
            .binary_search_by(|(_, other_order)| other_order.cmp(&spawn_order))
            .unwrap_err();
        self.spawn.insert(pos, (player_uid, spawn_order));
        self.id.insert(player_uid, ());
    }

    pub(super) fn remove_tee(&mut self, player_uid: PlayerUid) {
        self.id.remove(player_uid).unwrap();
        self.spawn.retain(|(pid, _)| *pid != player_uid);
    }

    pub(super) fn member_count(&self) -> usize {
        self.spawn.len()
    }
}

pub(crate) struct TeeOrderIter<T> {
    iter: T,
}

impl<T> Iterator for TeeOrderIter<T>
where
    T: Iterator<Item = PlayerUid>,
{
    type Item = PlayerUid;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next()
    }
}

impl TeeOrder {
    pub(crate) fn id_order(&self) -> TeeOrderIter<impl Iterator<Item = PlayerUid> + '_> {
        TeeOrderIter {
            iter: self.id.iter().map(|(tid, _)| tid),
        }
    }

    pub(crate) fn spawn_order(&self) -> TeeOrderIter<impl Iterator<Item = PlayerUid> + '_> {
        TeeOrderIter {
            iter: self.spawn.iter().map(|(tid, _)| *tid).rev(),
        }
    }
}

impl<T> TeeOrderIter<T> {
    // returns an iter over tees
    pub(crate) fn tees<U>(self, tees: &Tees<U>) -> TeesIter<'_, T, U> {
        TeesIter { tees, iter: self }
    }

    // returns an iter over tees
    pub(crate) fn tees_mut<U>(self, tees: &mut Tees<U>) -> TeesIterMut<'_, T, U> {
        TeesIterMut { tees, iter: self }
    }

    // returns an iter over tees
    pub(crate) fn tees_except<U>(
        self,
        tees: &Tees<U>,
        except: PlayerUid,
    ) -> TeesExceptIter<'_, T, U> {
        TeesExceptIter {
            except,
            tees,
            iter: self,
        }
    }

    // returns an iter over tees
    pub(crate) fn tees_except_mut<U>(
        self,
        tees: &mut Tees<U>,
        except: PlayerUid,
    ) -> TeesExceptIterMut<'_, T, U> {
        TeesExceptIterMut {
            except,
            tees,
            iter: self,
        }
    }
}

pub(crate) struct TeesIter<'a, T, U> {
    tees: &'a Tees<U>,
    iter: TeeOrderIter<T>,
}

impl<'a, T, U> Iterator for TeesIter<'a, T, U>
where
    T: Iterator<Item = PlayerUid>,
{
    type Item = (PlayerUid, &'a U);

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(tee_uid) = self.iter.next() {
            Some((tee_uid, self.tees.get_tee(tee_uid).unwrap()))
        } else {
            None
        }
    }
}

pub(crate) struct TeesIterMut<'a, T, U> {
    tees: &'a mut Tees<U>,
    iter: TeeOrderIter<T>,
}

impl<T, U> TeesIterMut<'_, T, U>
where
    T: Iterator<Item = PlayerUid>,
{
    pub fn next(&mut self) -> Option<(PlayerUid, &mut U)> {
        if let Some(tee_uid) = self.iter.next() {
            Some((tee_uid, self.tees.get_tee_mut(tee_uid).unwrap()))
        } else {
            None
        }
    }
}

pub(crate) struct TeesExceptIter<'a, T, U> {
    except: PlayerUid,
    tees: &'a Tees<U>,
    iter: TeeOrderIter<T>,
}

impl<'a, T, U> Iterator for TeesExceptIter<'a, T, U>
where
    T: Iterator<Item = PlayerUid>,
{
    type Item = (PlayerUid, &'a U);

    fn next(&mut self) -> Option<Self::Item> {
        for tee_uid in self.iter.by_ref() {
            if tee_uid == self.except {
                continue;
            }
            return Some((tee_uid, self.tees.get_tee(tee_uid).unwrap()));
        }
        None
    }
}

pub(crate) struct TeesExceptIterMut<'a, T, U> {
    except: PlayerUid,
    tees: &'a mut Tees<U>,
    iter: TeeOrderIter<T>,
}

impl<T, U> TeesExceptIterMut<'_, T, U>
where
    T: Iterator<Item = PlayerUid>,
{
    pub fn next(&mut self) -> Option<(PlayerUid, &mut U)> {
        for tee_uid in self.iter.by_ref() {
            if tee_uid == self.except {
                continue;
            }
            return Some((tee_uid, self.tees.get_tee_mut(tee_uid).unwrap()));
        }
        None
    }
}