// Eryn Wells <eryn@erynwells.me>

//! Defines routines for computing sight of a piece. Sight is the set of squares
//! that a piece can see. In other words, it's the set of squares attacked or
//! controled by a piece.

use chessfriend_bitboard::BitBoard;
use chessfriend_board::Board;
use chessfriend_core::{Color, Direction, PlacedPiece, Shape, Square};

macro_rules! ray_in_direction {
    ($square:expr, $blockers:expr, $direction:ident, $occupied_squares:tt) => {{
        let ray = BitBoard::ray($square, Direction::$direction);
        if let Some(first_occupied_square) = BitBoard::$occupied_squares(&(ray & $blockers)).next()
        {
            let remainder = BitBoard::ray(first_occupied_square, Direction::$direction);
            let attack_ray = ray & !remainder;
            attack_ray
        } else {
            *ray
        }
    }};
}

/// Compute sight of a white pawn.
fn _white_pawn_sight(
    pawn: BitBoard,
    occupancy: BitBoard,
    blockers: BitBoard,
    en_passant_square: BitBoard,
) -> BitBoard {
    let possible_squares = !occupancy | blockers | en_passant_square;
    let pawn = pawn.shift_north_west_one() | pawn.shift_north_east_one();
    pawn & possible_squares
}

fn _black_pawn_sight(
    pawn: BitBoard,
    occupancy: BitBoard,
    blockers: BitBoard,
    en_passant_square: BitBoard,
) -> BitBoard {
    let possible_squares = !occupancy | blockers | en_passant_square;
    let pawn = pawn.shift_south_west_one() | pawn.shift_south_east_one();
    pawn & possible_squares
}

fn _knight_sight(knight_square: Square, blockers: BitBoard) -> BitBoard {
    BitBoard::knight_moves(knight_square) & !blockers
}

fn _bishop_sight(bishop_square: Square, occupancy: BitBoard) -> BitBoard {
    #[rustfmt::skip]
    let sight = ray_in_direction!(bishop_square, occupancy, NorthEast, occupied_squares_trailing)
        | ray_in_direction!(bishop_square, occupancy, NorthWest, occupied_squares_trailing)
        | ray_in_direction!(bishop_square, occupancy, SouthEast, occupied_squares)
        | ray_in_direction!(bishop_square, occupancy, SouthWest, occupied_squares);

    sight
}

fn _rook_sight(rook_square: Square, occupancy: BitBoard) -> BitBoard {
    #[rustfmt::skip]
    let sight = ray_in_direction!(rook_square, occupancy, North, occupied_squares_trailing)
        | ray_in_direction!(rook_square, occupancy, East, occupied_squares_trailing)
        | ray_in_direction!(rook_square, occupancy, South, occupied_squares)
        | ray_in_direction!(rook_square, occupancy, West, occupied_squares);

    sight
}

fn _queen_sight(queen_square: Square, occupancy: BitBoard) -> BitBoard {
    #[rustfmt::skip]
    let sight = ray_in_direction!(queen_square, occupancy, NorthWest, occupied_squares_trailing)
        | ray_in_direction!(queen_square, occupancy, North, occupied_squares_trailing)
        | ray_in_direction!(queen_square, occupancy, NorthEast, occupied_squares_trailing)
        | ray_in_direction!(queen_square, occupancy, East, occupied_squares_trailing)
        | ray_in_direction!(queen_square, occupancy, SouthEast, occupied_squares)
        | ray_in_direction!(queen_square, occupancy, South, occupied_squares)
        | ray_in_direction!(queen_square, occupancy, SouthWest, occupied_squares)
        | ray_in_direction!(queen_square, occupancy, West, occupied_squares);

    sight
}

fn _king_sight(king_square: Square, blockers: BitBoard) -> BitBoard {
    BitBoard::king_moves(king_square) & !blockers
}
pub(crate) trait BishopSightExt {
    fn bishop_sight(&self, occupancy: BitBoard) -> BitBoard;
}

pub(crate) trait KingSightExt {
    fn king_sight(&self, board: &Board) -> BitBoard;
}

pub(crate) trait KnightSightExt {
    fn knight_sight(&self, board: &Board) -> BitBoard;
}

pub(crate) trait PawnSightExt {
    fn pawn_sight(&self, board: &Board, en_passant_square: Option<Square>) -> BitBoard;
    fn white_pawn_sight(&self, board: &Board, en_passant_square: Option<Square>) -> BitBoard;
    fn black_pawn_sight(&self, board: &Board, en_passant_square: Option<Square>) -> BitBoard;
}
pub(crate) trait QueenSightExt {
    fn queen_sight(&self, occupancy: BitBoard) -> BitBoard;
}

pub(crate) trait RookSightExt {
    fn rook_sight(&self, occupancy: BitBoard) -> BitBoard;
}

pub(crate) trait SliderSightExt: BishopSightExt + QueenSightExt + RookSightExt {}

pub(crate) trait SightExt {
    fn sight(&self, board: &Board, en_passant_square: Option<Square>) -> BitBoard;
}

pub(crate) trait SliderRayToSquareExt {
    fn ray_to_square(&self, origin: Square, target: Square) -> Option<BitBoard>;
}

impl SightExt for PlacedPiece {
    fn sight(&self, board: &Board, en_passant_square: Option<Square>) -> BitBoard {
        match self.shape() {
            Shape::Pawn => match self.color() {
                Color::White => self.white_pawn_sight(board, en_passant_square),
                Color::Black => self.black_pawn_sight(board, en_passant_square),
            },
            Shape::Knight => self.knight_sight(board),
            Shape::Bishop => self.bishop_sight(board.all_pieces_bitboard()),
            Shape::Rook => self.rook_sight(board.all_pieces_bitboard()),
            Shape::Queen => self.queen_sight(board.all_pieces_bitboard()),
            Shape::King => self.king_sight(board),
        }
    }
}

impl KingSightExt for PlacedPiece {
    fn king_sight(&self, board: &Board) -> BitBoard {
        _king_sight(
            self.square(),
            board.all_pieces_of_color_bitboard(self.color()),
        )
    }
}

impl KnightSightExt for PlacedPiece {
    fn knight_sight(&self, board: &Board) -> BitBoard {
        _knight_sight(
            self.square(),
            board.all_pieces_of_color_bitboard(self.color()),
        )
    }
}

impl PawnSightExt for PlacedPiece {
    fn pawn_sight(&self, board: &Board, en_passant_square: Option<Square>) -> BitBoard {
        match self.color() {
            Color::White => self.white_pawn_sight(board, en_passant_square),
            Color::Black => self.black_pawn_sight(board, en_passant_square),
        }
    }

    fn white_pawn_sight(&self, board: &Board, en_passant_square: Option<Square>) -> BitBoard {
        let opponent = self.color().other();
        _white_pawn_sight(
            self.square().into(),
            board.all_pieces_bitboard(),
            board.all_pieces_of_color_bitboard(opponent),
            en_passant_square.into(),
        )
    }

    fn black_pawn_sight(&self, board: &Board, en_passant_square: Option<Square>) -> BitBoard {
        let opponent = self.color().other();
        _black_pawn_sight(
            self.square().into(),
            board.all_pieces_bitboard(),
            board.all_pieces_of_color_bitboard(opponent),
            en_passant_square.into(),
        )
    }
}

impl BishopSightExt for PlacedPiece {
    fn bishop_sight(&self, occupancy: BitBoard) -> BitBoard {
        _bishop_sight(self.square(), occupancy)
    }
}

impl RookSightExt for PlacedPiece {
    fn rook_sight(&self, occupancy: BitBoard) -> BitBoard {
        _rook_sight(self.square(), occupancy)
    }
}

impl QueenSightExt for PlacedPiece {
    fn queen_sight(&self, occupancy: BitBoard) -> BitBoard {
        _queen_sight(self.square(), occupancy)
    }
}

impl SliderSightExt for PlacedPiece {}

impl SliderRayToSquareExt for Shape {
    fn ray_to_square(&self, origin: Square, target: Square) -> Option<BitBoard> {
        macro_rules! ray {
            ($square:expr, $direction:ident) => {
                (
                    BitBoard::ray($square, Direction::$direction),
                    Direction::$direction,
                )
            };
        }

        let target_bitboard: BitBoard = target.into();

        let ray_and_direction = match self {
            Shape::Bishop => [
                ray!(origin, NorthEast),
                ray!(origin, SouthEast),
                ray!(origin, SouthWest),
                ray!(origin, NorthWest),
            ]
            .into_iter()
            .find(|(&ray, _)| !(target_bitboard & ray).is_empty()),
            Shape::Rook => [
                ray!(origin, North),
                ray!(origin, East),
                ray!(origin, South),
                ray!(origin, West),
            ]
            .into_iter()
            .find(|(&ray, _)| !(target_bitboard & ray).is_empty()),
            Shape::Queen => [
                ray!(origin, North),
                ray!(origin, NorthEast),
                ray!(origin, East),
                ray!(origin, SouthEast),
                ray!(origin, South),
                ray!(origin, SouthWest),
                ray!(origin, West),
                ray!(origin, NorthWest),
            ]
            .into_iter()
            .find(|(&ray, _)| !(target_bitboard & ray).is_empty()),
            _ => None,
        };

        if let Some((ray, direction)) = ray_and_direction {
            let remainder = BitBoard::ray(target, direction);
            return Some(ray & !remainder);
        }

        None
    }
}

impl BishopSightExt for Square {
    fn bishop_sight(&self, occupancy: BitBoard) -> BitBoard {
        _bishop_sight(*self, occupancy)
    }
}

impl QueenSightExt for Square {
    fn queen_sight(&self, occupancy: BitBoard) -> BitBoard {
        _queen_sight(*self, occupancy)
    }
}

impl RookSightExt for Square {
    fn rook_sight(&self, occupancy: BitBoard) -> BitBoard {
        _rook_sight(*self, occupancy)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chessfriend_bitboard::bitboard;
    use chessfriend_core::{piece, Square};

    macro_rules! sight_test {
        ($test_name:ident, $position:expr, $piece:expr, $bitboard:expr) => {
            #[test]
            fn $test_name() {
                let pos = $position;
                let piece = $piece;
                let sight = pos.sight_of_piece(&piece);

                assert_eq!(sight, $bitboard);
            }
        };
        ($test_name:ident, $piece:expr, $bitboard:expr) => {
            sight_test! {$test_name, $crate::Position::empty(), $piece, $bitboard}
        };
    }

    #[test]
    fn pawns_and_knights_cannot_make_rays() {
        assert_eq!(Shape::Pawn.ray_to_square(Square::F7, Square::E8), None);
        assert_eq!(Shape::Knight.ray_to_square(Square::F6, Square::E8), None);
    }

    mod pawn {
        use crate::test_position;
        use chessfriend_bitboard::{bitboard, BitBoard};
        use chessfriend_core::{piece, Square};
        use chessfriend_moves::EnPassant;

        sight_test!(e4_pawn, piece!(White Pawn on E4), bitboard!(D5, F5));

        sight_test!(
            e4_pawn_one_blocker,
            test_position![
                White Bishop on D5,
                White Pawn on E4,
            ],
            piece!(White Pawn on E4),
            bitboard!(F5)
        );

        #[test]
        fn e4_pawn_two_blocker() {
            let pos = test_position!(
                White Bishop on D5,
                White Queen on F5,
                White Pawn on E4,
            );

            let piece = piece!(White Pawn on E4);
            let sight = pos.sight_of_piece(&piece);

            assert_eq!(sight, BitBoard::empty());
        }

        #[test]
        fn e4_pawn_capturable() {
            let pos = test_position!(
                Black Bishop on D5,
                White Queen on F5,
                White Pawn on E4,
            );

            let piece = piece!(White Pawn on E4);
            let sight = pos.sight_of_piece(&piece);

            assert_eq!(sight, bitboard!(D5));
        }

        #[test]
        fn e5_en_passant() {
            let mut pos = test_position!(
                White Pawn on E5,
                Black Pawn on D5,
            );
            pos.test_set_en_passant(EnPassant::from_target_square(Square::D6).unwrap());
            let piece = piece!(White Pawn on E5);
            let sight = pos.sight_of_piece(&piece);

            assert_eq!(sight, bitboard!(D6, F6));
        }
    }

    #[macro_use]
    mod knight {
        use chessfriend_bitboard::bitboard;
        use chessfriend_core::piece;

        sight_test!(
            f6_knight,
            piece!(Black Knight on F6),
            bitboard!(H7, G8, E8, D7, D5, E4, G4, H5)
        );
    }

    mod bishop {
        use super::*;

        sight_test!(
            c2_bishop,
            piece!(Black Bishop on C2),
            bitboard!(D1, B3, A4, B1, D3, E4, F5, G6, H7)
        );

        #[test]
        fn ray_to_square() {
            let generated_ray = Shape::Bishop.ray_to_square(Square::C5, Square::E7);
            let expected_ray = bitboard![D6, E7];
            assert_eq!(generated_ray, Some(expected_ray));
        }
    }

    mod rook {
        use super::*;
        use crate::test_position;

        sight_test!(
            g3_rook,
            piece!(White Rook on G3),
            bitboard!(G1, G2, G4, G5, G6, G7, G8, A3, B3, C3, D3, E3, F3, H3)
        );

        sight_test!(
            e4_rook_with_e1_white_king_e7_black_king,
            test_position![
                White Rook on E4,
                White King on E1,
                Black King on E7,
            ],
            piece!(White Rook on E4),
            bitboard!(A4, B4, C4, D4, F4, G4, H4, E2, E3, E5, E6, E7, E1)
        );

        #[test]
        fn ray_to_square() {
            let generated_ray = Shape::Rook.ray_to_square(Square::C2, Square::C6);
            let expected_ray = bitboard![C3, C4, C5, C6];
            assert_eq!(generated_ray, Some(expected_ray));

            let generated_ray = Shape::Rook.ray_to_square(Square::D2, Square::H2);
            let expected_ray = bitboard![E2, F2, G2, H2];
            assert_eq!(generated_ray, Some(expected_ray));

            let generated_ray = Shape::Rook.ray_to_square(Square::G6, Square::B6);
            let expected_ray = bitboard![B6, C6, D6, E6, F6];
            assert_eq!(generated_ray, Some(expected_ray));

            let generated_ray = Shape::Rook.ray_to_square(Square::A6, Square::A3);
            let expected_ray = bitboard![A5, A4, A3];
            assert_eq!(generated_ray, Some(expected_ray));
        }
    }

    mod king {
        use chessfriend_bitboard::bitboard;
        use chessfriend_core::piece;

        sight_test!(
            e1_king,
            piece!(White King on E1),
            bitboard![D1, D2, E2, F2, F1]
        );
    }
}