// Eryn Wells <eryn@erynwells.me>

use super::Position;
use crate::piece::{Color, Piece, PlacedPiece, Shape};
use chessfriend_bitboard::BitBoard;
use chessfriend_core::Square;

pub struct Pieces<'a> {
    color: Color,
    position: &'a Position,

    current_shape: Option<Shape>,

    shape_iterator: Box<dyn Iterator<Item = &'static Shape>>,
    square_iterator: Option<Box<dyn Iterator<Item = Square>>>,
}

impl<'a> Pieces<'a> {
    pub(crate) fn new(position: &Position, color: Color) -> Pieces {
        Pieces {
            color,
            position,
            current_shape: None,
            shape_iterator: Box::new(Shape::iter()),
            square_iterator: None,
        }
    }
}

impl<'a> Iterator for Pieces<'a> {
    type Item = PlacedPiece;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(square_iterator) = &mut self.square_iterator {
            if let (Some(square), Some(shape)) = (square_iterator.next(), self.current_shape) {
                return Some(PlacedPiece::new(Piece::new(self.color, shape), square));
            }
        }

        let mut current_shape: Option<Shape> = None;
        let mut next_nonempty_bitboard: Option<&BitBoard> = None;

        while let Some(shape) = self.shape_iterator.next() {
            let piece = Piece::new(self.color, *shape);

            let bitboard = self.position.bitboard_for_piece(piece);
            if bitboard.is_empty() {
                continue;
            }

            next_nonempty_bitboard = Some(bitboard);
            current_shape = Some(*shape);

            break;
        }

        if let (Some(bitboard), Some(shape)) = (next_nonempty_bitboard, current_shape) {
            let mut square_iterator = bitboard.occupied_squares();

            let mut next_placed_piece: Option<PlacedPiece> = None;
            if let Some(square) = square_iterator.next() {
                next_placed_piece = Some(PlacedPiece::new(Piece::new(self.color, shape), square));
            }

            self.square_iterator = Some(Box::new(square_iterator));
            self.current_shape = Some(shape);

            return next_placed_piece;
        }

        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::piece::{Color, Piece, Shape};
    use crate::{Position, PositionBuilder};
    use std::collections::HashSet;

    #[test]
    fn empty() {
        let pos = Position::empty();
        let mut pieces = pos.pieces(Color::White);
        assert_eq!(pieces.next(), None);
    }

    #[test]
    fn one() {
        let pos = PositionBuilder::new()
            .place_piece(piece!(White Queen on E4))
            .build();
        println!("{:#?}", &pos);

        let mut pieces = pos.pieces(Color::White);
        assert_eq!(pieces.next(), Some(piece!(White Queen on E4)));
        assert_eq!(pieces.next(), Some(piece!(White King on E1)));
        assert_eq!(pieces.next(), None);
    }

    #[test]
    fn multiple_pieces() {
        let pos = PositionBuilder::new()
            .place_piece(piece!(White Queen on E4))
            .place_piece(piece!(White King on A1))
            .place_piece(piece!(White Pawn on B2))
            .place_piece(piece!(White Pawn on C2))
            .build();
        println!("{}", crate::position::DiagramFormatter::new(&pos));

        let expected_placed_pieces = HashSet::from([
            piece!(White Queen on E4),
            piece!(White King on A1),
            piece!(White Pawn on B2),
            piece!(White Pawn on C2),
        ]);

        let placed_pieces = HashSet::from_iter(pos.pieces(Color::White));

        assert_eq!(
            placed_pieces,
            expected_placed_pieces,
            "{:#?}",
            placed_pieces
                .symmetric_difference(&expected_placed_pieces)
                .into_iter()
                .map(|pp| format!("{}", pp))
                .collect::<Vec<String>>()
        );
    }
}