// 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 crate::{r#move::Castle, Position};
use chessfriend_bitboard::BitBoard;
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(
        position: &Position,
        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 = !position.king_danger(color);
        let all_king_moves = BitBoard::king_moves(square);

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

        let opposing_pieces = position.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 position.player_can_castle(color, Castle::KingSide) {
            move_set.kingside_castle();
        }
        if position.player_can_castle(color, Castle::QueenSide) {
            move_set.queenside_castle();
        }

        move_set
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{assert_move_list, position, test_position, Move, MoveBuilder, PositionBuilder};
    use chessfriend_bitboard::bitboard;
    use chessfriend_core::{piece, Color, Square};
    use std::collections::HashSet;

    #[test]
    fn one_king() {
        let pos = position![White King on E4];

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

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

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

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

        assert_move_list!(generated_moves, expected_moves, pos);
    }

    #[test]
    fn one_king_corner() {
        let pos = position![White King on A1];

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

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

        let expected_moves = [
            MoveBuilder::new(piece!(White King), Square::A1, Square::A2).build(),
            MoveBuilder::new(piece!(White King), Square::A1, Square::B1).build(),
            MoveBuilder::new(piece!(White King), Square::A1, Square::B2).build(),
        ];

        let mut generated_moves: HashSet<Move> = 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
        );
    }

    #[test]
    fn black_king_in_check_by_rook() {
        let pos = PositionBuilder::new()
            .place_piece(piece!(White King on E1))
            .place_piece(piece!(White Rook on E4))
            .place_piece(piece!(Black King on E7))
            .to_move(Color::Black)
            .build();

        assert!(pos.is_king_in_check());

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

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

        assert_eq!(generated_moves, expected_moves);
    }

    #[test]
    fn white_king_unobstructed_castles() {
        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, Color::White, BitBoard::FULL, BitBoard::FULL);
        let generated_moves: HashSet<Move> = generator.iter().collect();

        let king = piece!(White King);
        assert!(generated_moves.contains(
            &MoveBuilder::new(king, Square::E1, Square::G1)
                .castle(Castle::KingSide)
                .build()
        ));
        assert!(generated_moves.contains(
            &MoveBuilder::new(king, Square::E1, Square::C1)
                .castle(Castle::QueenSide)
                .build()
        ));
    }

    #[test]
    fn white_king_obstructed_queenside_castle() {
        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, Color::White, BitBoard::FULL, BitBoard::FULL);
        let generated_moves: HashSet<Move> = generator.iter().collect();

        let king = piece!(White King);
        assert!(generated_moves.contains(
            &MoveBuilder::new(king, Square::E1, Square::G1)
                .castle(Castle::KingSide)
                .build()
        ));
        assert!(!generated_moves.contains(
            &MoveBuilder::new(king, Square::E1, Square::C1)
                .castle(Castle::QueenSide)
                .build()
        ));
    }

    #[test]
    fn white_king_obstructed_kingside_castle() {
        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, Color::White, BitBoard::FULL, BitBoard::FULL);
        let generated_moves: HashSet<Move> = generator.iter().collect();

        let king = piece!(White King);
        assert!(!generated_moves.contains(
            &MoveBuilder::new(king, Square::E1, Square::G1)
                .castle(Castle::KingSide)
                .build()
        ));
        assert!(generated_moves.contains(
            &MoveBuilder::new(king, Square::E1, Square::C1)
                .castle(Castle::QueenSide)
                .build()
        ));
    }
}