// Eryn Wells <eryn@erynwells.me>

use chessfriend_core::{Piece, Square};
use std::iter::FromIterator;

#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub(crate) struct Mailbox([Option<Piece>; Square::NUM]);

impl Mailbox {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn get(&self, square: Square) -> Option<Piece> {
        self.0[square as usize]
    }

    pub fn set(&mut self, piece: Piece, square: Square) {
        self.0[square as usize] = Some(piece);
    }

    pub fn remove(&mut self, square: Square) {
        self.0[square as usize] = None;
    }

    pub fn iter(&self) -> impl Iterator<Item = (Square, Piece)> {
        Square::ALL
            .into_iter()
            .zip(self.0)
            .filter_map(|(square, piece)| piece.map(|piece| (square, piece)))
    }
}

impl Default for Mailbox {
    fn default() -> Self {
        Self([None; Square::NUM])
    }
}

impl From<[Option<Piece>; Square::NUM]> for Mailbox {
    fn from(value: [Option<Piece>; Square::NUM]) -> Self {
        Mailbox(value)
    }
}

impl FromIterator<(Square, Piece)> for Mailbox {
    fn from_iter<T: IntoIterator<Item = (Square, Piece)>>(iter: T) -> Self {
        iter.into_iter()
            .fold(Self::new(), |mut mailbox, (square, piece)| {
                mailbox.set(piece, square);
                mailbox
            })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chessfriend_core::piece;
    use std::collections::HashSet;

    #[test]
    fn iter_returns_all_pieces() {
        let mut mailbox = Mailbox::new();
        mailbox.set(piece!(White Queen), Square::C3);
        mailbox.set(piece!(White Rook), Square::H8);
        mailbox.set(piece!(Black Bishop), Square::E4);
        mailbox.set(piece!(Black Pawn), Square::F6);

        let pieces = mailbox.iter().collect::<HashSet<_>>();
        assert_eq!(
            pieces,
            HashSet::from([
                (Square::C3, piece!(White Queen)),
                (Square::H8, piece!(White Rook)),
                (Square::E4, piece!(Black Bishop)),
                (Square::F6, piece!(Black Pawn)),
            ])
        );
    }
}