// Eryn Wells <eryn@erynwells.me>

use crate::sight::SliderRayToSquareExt;
use chessfriend_bitboard::BitBoard;
use chessfriend_core::Shape;

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CheckingPieces {
    bitboards: [BitBoard; 5],
}

impl CheckingPieces {
    pub(crate) fn new(
        pawn: BitBoard,
        knight: BitBoard,
        bishop: BitBoard,
        rook: BitBoard,
        queen: BitBoard,
    ) -> CheckingPieces {
        CheckingPieces {
            bitboards: [pawn, knight, bishop, rook, queen],
        }
    }

    /// The number of checking pieces.
    pub fn count(&self) -> u32 {
        self.bitboards.iter().map(BitBoard::population_count).sum()
    }

    /// A BitBoard representing the set of pieces that must be captured to resolve check.
    pub fn capture_mask(&self) -> BitBoard {
        if self.count() == 0 {
            BitBoard::FULL
        } else {
            self.bitboards
                .iter()
                .fold(BitBoard::EMPTY, std::ops::BitOr::bitor)
        }
    }

    /// A BitBoard representing the set of squares to which a player can move a piece to block a
    /// checking piece.
    pub fn push_mask(&self, king: &BitBoard) -> BitBoard {
        let target = king.first_occupied_square().unwrap();

        macro_rules! push_mask_for_shape {
            ($push_mask:expr, $shape:ident, $king:expr) => {{
                let checking_pieces = self.bitboard_for_shape(Shape::$shape);
                if !checking_pieces.is_empty() {
                    if let Some(checking_ray) = checking_pieces
                        .occupied_squares()
                        .flat_map(|sq| Shape::$shape.ray_to_square(sq, target).into_iter())
                        .find(|bb| !(bb & $king).is_empty())
                    {
                        $push_mask |= checking_ray & !$king
                    }
                }
            }};
        }

        let mut push_mask = BitBoard::EMPTY;

        push_mask_for_shape!(push_mask, Bishop, king);
        push_mask_for_shape!(push_mask, Rook, king);
        push_mask_for_shape!(push_mask, Queen, king);

        push_mask
    }

    fn bitboard_for_shape(&self, shape: Shape) -> &BitBoard {
        &self.bitboards[shape as usize]
    }
}