// Eryn Wells <eryn@erynwells.me>

//! Declares the KingMoveGenerator type. This struct is responsible for
//! generating the possible moves for the king in the given position.

use super::{move_generator_declaration, MoveGeneratorInternal, MoveSet};
use chessfriend_bitboard::BitBoard;
use chessfriend_board::{castle::Castle, Board};
use chessfriend_core::{PlacedPiece, Shape};

move_generator_declaration!(KingMoveGenerator, struct);
move_generator_declaration!(KingMoveGenerator, new);
move_generator_declaration!(KingMoveGenerator, getters);

impl MoveGeneratorInternal for KingMoveGenerator {
    fn shape() -> Shape {
        Shape::King
    }

    fn move_set_for_piece(
        board: &Board,
        placed_piece: &PlacedPiece,
        _capture_mask: BitBoard,
        _push_mask: BitBoard,
    ) -> MoveSet {
        let piece = placed_piece.piece();
        let color = piece.color();
        let square = placed_piece.square();

        let safe_squares = BitBoard::FULL;
        let all_king_moves = BitBoard::king_moves(square);

        let empty_squares = board.empty_squares();
        let safe_empty_squares = empty_squares & safe_squares;

        let opposing_pieces = board.bitboard_for_color(color.other());
        let opposing_pieces_on_safe_squares = opposing_pieces & safe_squares;

        let quiet_moves = all_king_moves & safe_empty_squares;
        let capture_moves = all_king_moves & opposing_pieces_on_safe_squares;

        let mut move_set = MoveSet::new(*placed_piece)
            .quiet_moves(quiet_moves)
            .capture_moves(capture_moves);

        if board.player_can_castle(color, Castle::KingSide) {
            move_set.kingside_castle();
        }
        if board.player_can_castle(color, Castle::QueenSide) {
            move_set.queenside_castle();
        }

        move_set
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{assert_move_list, test_position, testing::*};
    use chessfriend_bitboard::bitboard;
    use chessfriend_board::castle::Castle;
    use chessfriend_core::{piece, Color, Square};
    use chessfriend_moves::{Builder as MoveBuilder, Move};
    use std::collections::HashSet;

    #[test]
    fn one_king() -> TestResult {
        let pos = test_position![White King on E4];

        let generator =
            KingMoveGenerator::new(&pos.board, Color::White, BitBoard::FULL, BitBoard::FULL);

        assert_eq!(
            generator._test_bitboard(),
            bitboard![E5 F5 F4 F3 E3 D3 D4 D5]
        );

        let builder = MoveBuilder::push(&piece!(White King on E4));
        let expected_moves: HashSet<Move> = HashSet::from_iter([
            builder.clone().to(Square::D5).build()?,
            builder.clone().to(Square::E5).build()?,
            builder.clone().to(Square::F5).build()?,
            builder.clone().to(Square::F4).build()?,
            builder.clone().to(Square::F3).build()?,
            builder.clone().to(Square::E3).build()?,
            builder.clone().to(Square::D3).build()?,
            builder.clone().to(Square::D4).build()?,
        ]);

        let generated_moves: HashSet<Move> = generator.iter().collect();

        assert_move_list!(generated_moves, expected_moves, pos);

        Ok(())
    }

    #[test]
    fn one_king_corner() -> TestResult {
        let pos = test_position![White King on A1];

        let generator =
            KingMoveGenerator::new(&pos.board, Color::White, BitBoard::FULL, BitBoard::FULL);

        let generated_bitboard = generator._test_bitboard();
        let expected_bitboard = bitboard![A2 B2 B1];
        assert_eq!(
            generator._test_bitboard(),
            bitboard![A2 B2 B1],
            "Generated:\n{generated_bitboard}\nExpected:\n{expected_bitboard}"
        );

        let builder = MoveBuilder::push(&piece!(White King on A1));
        let expected_moves = [
            builder.clone().to(Square::A2).build()?,
            builder.clone().to(Square::B1).build()?,
            builder.clone().to(Square::B2).build()?,
        ];

        let mut generated_moves: HashSet<_> = generator.iter().collect();

        for ex_move in expected_moves {
            assert!(
                generated_moves.remove(&ex_move),
                "{:#?} was not generated",
                &ex_move
            );
        }

        assert!(
            generated_moves.is_empty(),
            "Moves unexpectedly present: {:#?}",
            generated_moves
        );

        Ok(())
    }

    #[test]
    fn black_king_in_check_by_rook() {
        let pos = test_position!(Black, [
            White King on E1,
            White Rook on E4,
            Black King on E7,
        ]);

        assert!(pos.is_king_in_check());

        let generator =
            KingMoveGenerator::new(&pos.board, Color::Black, BitBoard::FULL, BitBoard::FULL);
        let generated_moves = generator._test_bitboard();

        let expected_moves = bitboard![F8 F7 F6 D6 D7 D8];

        assert_eq!(generated_moves, expected_moves);
    }

    #[test]
    fn white_king_unobstructed_castles() -> TestResult {
        let pos = test_position!(
            White King on E1,
            White Rook on A1,
            White Rook on H1,
        );

        assert!(pos.player_can_castle(Color::White, Castle::KingSide));
        assert!(pos.player_can_castle(Color::White, Castle::QueenSide));

        let generator =
            KingMoveGenerator::new(&pos.board, Color::White, BitBoard::FULL, BitBoard::FULL);
        let generated_moves: HashSet<Move> = generator.iter().collect();

        assert!(generated_moves
            .contains(&MoveBuilder::castling(Color::White, Castle::KingSide).build()?));
        assert!(generated_moves
            .contains(&MoveBuilder::castling(Color::White, Castle::QueenSide).build()?));

        Ok(())
    }

    #[test]
    fn white_king_obstructed_queenside_castle() -> TestResult {
        let pos = test_position!(
            White King on E1,
            White Knight on B1,
            White Rook on A1,
            White Rook on H1,
        );

        assert!(pos.player_can_castle(Color::White, Castle::KingSide));
        assert!(!pos.player_can_castle(Color::White, Castle::QueenSide));

        let generator =
            KingMoveGenerator::new(&pos.board, Color::White, BitBoard::FULL, BitBoard::FULL);
        let generated_moves: HashSet<Move> = generator.iter().collect();

        assert!(generated_moves
            .contains(&MoveBuilder::castling(Color::White, Castle::KingSide).build()?));
        assert!(!generated_moves
            .contains(&MoveBuilder::castling(Color::White, Castle::QueenSide).build()?));

        Ok(())
    }

    #[test]
    fn white_king_obstructed_kingside_castle() -> TestResult {
        let pos = test_position!(
            White King on E1,
            White Rook on A1,
            White Knight on G1,
            White Rook on H1,
        );

        assert!(!pos.player_can_castle(Color::White, Castle::KingSide));
        assert!(pos.player_can_castle(Color::White, Castle::QueenSide));

        let generator =
            KingMoveGenerator::new(&pos.board, Color::White, BitBoard::FULL, BitBoard::FULL);
        let generated_moves: HashSet<Move> = generator.iter().collect();

        assert!(!generated_moves
            .contains(&MoveBuilder::castling(Color::White, Castle::KingSide).build()?));
        assert!(generated_moves
            .contains(&MoveBuilder::castling(Color::White, Castle::QueenSide).build()?));

        Ok(())
    }
}