// Eryn Wells <eryn@erynwells.me>

//! Implements a move generator for pawns.

use super::GeneratedMove;
use crate::{Move, PromotionShape};
use chessfriend_bitboard::{bit_scanner::TrailingBitScanner, BitBoard};
use chessfriend_board::Board;
use chessfriend_core::{Color, Direction, Rank, Square};
use std::{iter::FusedIterator, slice};

pub struct PawnMoveGenerator {
    color: Color,
    single_pushes: BitBoard,
    double_pushes: BitBoard,
    left_captures: BitBoard,
    right_captures: BitBoard,
    en_passant: BitBoard,
    target_iterator: TrailingBitScanner,
    promotion_iterator: Option<PromotionIterator>,
    move_type: MoveType,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum MoveType {
    SinglePushes,
    DoublePushes,
    LeftCaptures,
    RightCaptures,
}

struct PromotionIterator {
    origin: Square,
    target: Square,
    iterator: slice::Iter<'static, PromotionShape>,
}

impl PawnMoveGenerator {
    #[must_use]
    pub fn new(board: &Board, color: Option<Color>) -> Self {
        let color = board.unwrap_color(color);

        let pawns = board.pawns(color);
        let occupied = board.occupancy();
        let empty = !occupied;
        let enemies = board.enemies(color);

        // En passant captures present a particular challenge. Include the
        // target e.p. square when computing captures (i.e. treat it like an
        // enemy piece is on that square) but do not include it when generating
        // capture moves. If it is included, a regular capture move will be
        // generated where a special e.p. move should be created instead.
        //
        // So, include it in the enemies set when computing captures, then
        // remove it from the left and right captures bitboards before passing
        // them into the move generator. Additionally, include the target e.p.
        // square in the e.p. bitboard iff the capture bitboards include it.

        let en_passant: BitBoard = board.en_passant_target().into();

        let (single_pushes, double_pushes) = Self::pushes(pawns, color, empty);
        let (left_captures, right_captures) = Self::captures(pawns, color, enemies | en_passant);

        Self {
            color,
            single_pushes,
            double_pushes,
            left_captures,
            right_captures,
            en_passant,
            target_iterator: single_pushes.occupied_squares_trailing(),
            promotion_iterator: None,
            move_type: MoveType::SinglePushes,
        }
    }

    fn pushes(pawns: BitBoard, color: Color, empty: BitBoard) -> (BitBoard, BitBoard) {
        match color {
            Color::White => {
                const THIRD_RANK: BitBoard = BitBoard::rank(Rank::THREE);

                let single_pushes = pawns.shift_north_one() & empty;
                let double_pushes = (single_pushes & THIRD_RANK).shift_north_one() & empty;

                (single_pushes, double_pushes)
            }
            Color::Black => {
                const SIXTH_RANK: BitBoard = BitBoard::rank(Rank::SIX);

                let single_pushes = pawns.shift_south_one() & empty;
                let double_pushes = (single_pushes & SIXTH_RANK).shift_south_one() & empty;

                (single_pushes, double_pushes)
            }
        }
    }

    fn captures(pawns: BitBoard, color: Color, enemies: BitBoard) -> (BitBoard, BitBoard) {
        match color {
            Color::White => {
                let left_captures = pawns.shift_north_west_one() & enemies;
                let right_captures = pawns.shift_north_east_one() & enemies;
                (left_captures, right_captures)
            }
            Color::Black => {
                let left_captures = pawns.shift_south_east_one() & enemies;
                let right_captures = pawns.shift_south_west_one() & enemies;
                (left_captures, right_captures)
            }
        }
    }

    fn calculate_origin_square(&self, target: Square) -> Option<Square> {
        match self.move_type {
            MoveType::SinglePushes => match self.color {
                Color::White => target.neighbor(Direction::South),
                Color::Black => target.neighbor(Direction::North),
            },
            MoveType::DoublePushes => match self.color {
                Color::White => target
                    .neighbor(Direction::South)?
                    .neighbor(Direction::South),
                Color::Black => target
                    .neighbor(Direction::North)?
                    .neighbor(Direction::North),
            },
            MoveType::LeftCaptures => match self.color {
                Color::White => target.neighbor(Direction::SouthEast),
                Color::Black => target.neighbor(Direction::NorthWest),
            },
            MoveType::RightCaptures => match self.color {
                Color::White => target.neighbor(Direction::SouthWest),
                Color::Black => target.neighbor(Direction::NorthEast),
            },
        }
    }

    fn next_move_type(&mut self) -> Option<MoveType> {
        let next_move_type = self.move_type.next();

        if let Some(next_move_type) = next_move_type {
            let next_bitboard = match next_move_type {
                MoveType::SinglePushes => self.single_pushes,
                MoveType::DoublePushes => self.double_pushes,
                MoveType::LeftCaptures => self.left_captures,
                MoveType::RightCaptures => self.right_captures,
            };

            self.target_iterator = next_bitboard.occupied_squares_trailing();
            self.move_type = next_move_type;
        }

        next_move_type
    }

    fn next_promotion_move(&mut self) -> Option<GeneratedMove> {
        if let Some(promotion_iterator) = self.promotion_iterator.as_mut() {
            if let Some(shape) = promotion_iterator.iterator.next() {
                let origin = promotion_iterator.origin;
                let target = promotion_iterator.target;

                return if matches!(
                    self.move_type,
                    MoveType::LeftCaptures | MoveType::RightCaptures
                ) {
                    Some(Move::capture_promotion(origin, target, *shape).into())
                } else {
                    Some(Move::promotion(origin, target, *shape).into())
                };
            }
        }

        None
    }
}

impl std::iter::Iterator for PawnMoveGenerator {
    type Item = GeneratedMove;

    fn next(&mut self) -> Option<Self::Item> {
        let next_promotion_move = self.next_promotion_move();
        if next_promotion_move.is_some() {
            return next_promotion_move;
        }

        self.promotion_iterator = None;

        if let Some(target) = self.target_iterator.next() {
            let origin = self
                .calculate_origin_square(target)
                .expect("bogus origin square");

            if target.rank().is_promotable_rank() {
                self.promotion_iterator = Some(PromotionIterator {
                    origin,
                    target,
                    iterator: PromotionShape::ALL.iter(),
                });

                return self.next();
            }

            match self.move_type {
                MoveType::SinglePushes => Some(GeneratedMove {
                    ply: Move::quiet(origin, target),
                }),
                MoveType::DoublePushes => Some(GeneratedMove {
                    ply: Move::double_push(origin, target),
                }),
                MoveType::LeftCaptures | MoveType::RightCaptures => {
                    let target_bitboard: BitBoard = target.into();
                    Some(GeneratedMove {
                        ply: if (target_bitboard & self.en_passant).is_populated() {
                            Move::en_passant_capture(origin, target)
                        } else {
                            Move::capture(origin, target)
                        },
                    })
                }
            }
        } else if self.next_move_type().is_some() {
            self.next()
        } else {
            None
        }
    }
}

impl FusedIterator for PawnMoveGenerator {}

impl MoveType {
    fn next(self) -> Option<Self> {
        match self {
            MoveType::SinglePushes => Some(MoveType::DoublePushes),
            MoveType::DoublePushes => Some(MoveType::LeftCaptures),
            MoveType::LeftCaptures => Some(MoveType::RightCaptures),
            MoveType::RightCaptures => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{assert_move_list, assert_move_list_contains, assert_move_list_does_not_contain, ply, Move};
    use chessfriend_board::{fen::FromFenStr, test_board};
    use chessfriend_core::{Color, Square};
    use std::collections::HashSet;

    #[test]
    fn black_b7_pushes_and_double_pushes() {
        let board = test_board!(Black Pawn on B7);
        let generated_moves: HashSet<GeneratedMove> =
            PawnMoveGenerator::new(&board, Some(Color::Black)).collect();
        assert_eq!(
            generated_moves,
            [
                GeneratedMove {
                    ply: Move::quiet(Square::B7, Square::B6)
                },
                GeneratedMove {
                    ply: Move::double_push(Square::B7, Square::B5)
                }
            ]
            .into()
        );
    }

    #[test]
    fn black_e2_pushes_and_double_pushes() {
        let board = test_board!(White Pawn on E2);
        let generated_moves: HashSet<GeneratedMove> =
            PawnMoveGenerator::new(&board, Some(Color::White)).collect();
        assert_eq!(
            generated_moves,
            [
                GeneratedMove {
                    ply: Move::quiet(Square::E2, Square::E3)
                },
                GeneratedMove {
                    ply: Move::double_push(Square::E2, Square::E4)
                }
            ]
            .into()
        );
    }

    #[test]
    fn black_b7_pushes_with_b5_blocker() {
        let board = test_board!(Black Pawn on B7, White Bishop on B5);
        let generated_moves: HashSet<GeneratedMove> =
            PawnMoveGenerator::new(&board, Some(Color::Black)).collect();
        assert_eq!(
            generated_moves,
            [GeneratedMove {
                ply: Move::quiet(Square::B7, Square::B6)
            }]
            .into()
        );
    }

    #[test]
    fn black_b7_pushes_with_b6_blocker() {
        let board = test_board!(Black Pawn on B7, Black Bishop on B6);
        let generated_moves: HashSet<GeneratedMove> =
            PawnMoveGenerator::new(&board, Some(Color::Black)).collect();
        assert!(generated_moves.is_empty());
    }

    #[test]
    fn white_e2_moves_with_e4_blocker() {
        let board = test_board!(White Pawn on E2, White Bishop on E4);
        let generated_moves: HashSet<GeneratedMove> =
            PawnMoveGenerator::new(&board, Some(Color::White)).collect();
        assert_eq!(
            generated_moves,
            [GeneratedMove {
                ply: Move::quiet(Square::E2, Square::E3)
            }]
            .into()
        );
    }

    #[test]
    fn white_e2_moves_with_e3_blocker() {
        let board = test_board!(White Pawn on E2, White Bishop on E3);
        let generated_moves: HashSet<GeneratedMove> =
            PawnMoveGenerator::new(&board, Some(Color::White)).collect();
        assert!(generated_moves.is_empty());
    }

    #[test]
    fn white_f6_left_captures() {
        let white_captures_board = test_board!(White Pawn on F6, Black Queen on E7);
        let generated_moves: HashSet<GeneratedMove> =
            PawnMoveGenerator::new(&white_captures_board, Some(Color::White)).collect();
        assert_eq!(
            generated_moves,
            [
                GeneratedMove {
                    ply: Move::capture(Square::F6, Square::E7)
                },
                GeneratedMove {
                    ply: Move::quiet(Square::F6, Square::F7)
                }
            ]
            .into()
        );
    }

    #[test]
    fn black_d5_left_captures() {
        let black_captures_board = test_board!(Black Pawn on D5, White Queen on E4);
        let generated_moves = PawnMoveGenerator::new(&black_captures_board, Some(Color::Black));
        assert_move_list!(generated_moves, [ply!(D5 x E4), ply!(D5 - D4),]);
    }

    #[test]
    fn white_f6_right_captures() {
        let white_captures_board = test_board!(White Pawn on F6, Black Queen on G7);
        let generated_moves: HashSet<GeneratedMove> =
            PawnMoveGenerator::new(&white_captures_board, Some(Color::White)).collect();
        assert_eq!(
            generated_moves,
            [
                GeneratedMove {
                    ply: Move::capture(Square::F6, Square::G7)
                },
                GeneratedMove {
                    ply: Move::quiet(Square::F6, Square::F7)
                }
            ]
            .into()
        );
    }

    #[test]
    fn black_d5_right_captures() {
        let black_captures_board = test_board!(Black Pawn on D5, White Queen on C4);
        let generated_moves: HashSet<GeneratedMove> =
            PawnMoveGenerator::new(&black_captures_board, Some(Color::Black)).collect();
        assert_eq!(
            generated_moves,
            [
                GeneratedMove {
                    ply: Move::capture(Square::D5, Square::C4)
                },
                GeneratedMove {
                    ply: Move::quiet(Square::D5, Square::D4)
                }
            ]
            .into()
        );
    }

    #[test]
    fn white_g7_push_promotions() {
        let board = test_board!(White Pawn on G7);
        let generated_moves: HashSet<GeneratedMove> =
            PawnMoveGenerator::new(&board, Some(Color::White)).collect();
        assert_eq!(
            generated_moves,
            [
                Move::promotion(Square::G7, Square::G8, PromotionShape::Knight).into(),
                Move::promotion(Square::G7, Square::G8, PromotionShape::Bishop).into(),
                Move::promotion(Square::G7, Square::G8, PromotionShape::Rook).into(),
                Move::promotion(Square::G7, Square::G8, PromotionShape::Queen).into(),
            ]
            .into()
        );
    }

    #[test]
    fn white_d7_push_and_capture_promotions() {
        let board = test_board!(
            White Pawn on D7,
            Black Queen on E8
        );
        let generated_moves: HashSet<GeneratedMove> =
            PawnMoveGenerator::new(&board, Some(Color::White)).collect();
        assert_eq!(
            generated_moves,
            [
                Move::promotion(Square::D7, Square::D8, PromotionShape::Knight).into(),
                Move::promotion(Square::D7, Square::D8, PromotionShape::Bishop).into(),
                Move::promotion(Square::D7, Square::D8, PromotionShape::Rook).into(),
                Move::promotion(Square::D7, Square::D8, PromotionShape::Queen).into(),
                Move::capture_promotion(Square::D7, Square::E8, PromotionShape::Knight).into(),
                Move::capture_promotion(Square::D7, Square::E8, PromotionShape::Bishop).into(),
                Move::capture_promotion(Square::D7, Square::E8, PromotionShape::Rook).into(),
                Move::capture_promotion(Square::D7, Square::E8, PromotionShape::Queen).into(),
            ]
            .into()
        );
    }

    #[test]
    fn black_a2_push_promotions() {
        let board = test_board!(Black Pawn on A2);
        let generated_moves: HashSet<GeneratedMove> =
            PawnMoveGenerator::new(&board, Some(Color::Black)).collect();
        assert_eq!(
            generated_moves,
            [
                Move::promotion(Square::A2, Square::A1, PromotionShape::Knight).into(),
                Move::promotion(Square::A2, Square::A1, PromotionShape::Bishop).into(),
                Move::promotion(Square::A2, Square::A1, PromotionShape::Rook).into(),
                Move::promotion(Square::A2, Square::A1, PromotionShape::Queen).into(),
            ]
            .into()
        );
    }

    #[test]
    fn black_h2_push_and_capture_promotions() {
        let board = test_board!(
            Black Pawn on H2,
            White Queen on G1,
        );
        let generated_moves: HashSet<GeneratedMove> =
            PawnMoveGenerator::new(&board, Some(Color::Black)).collect();
        assert_eq!(
            generated_moves,
            [
                Move::promotion(Square::H2, Square::H1, PromotionShape::Knight).into(),
                Move::promotion(Square::H2, Square::H1, PromotionShape::Bishop).into(),
                Move::promotion(Square::H2, Square::H1, PromotionShape::Rook).into(),
                Move::promotion(Square::H2, Square::H1, PromotionShape::Queen).into(),
                Move::capture_promotion(Square::H2, Square::G1, PromotionShape::Knight).into(),
                Move::capture_promotion(Square::H2, Square::G1, PromotionShape::Bishop).into(),
                Move::capture_promotion(Square::H2, Square::G1, PromotionShape::Rook).into(),
                Move::capture_promotion(Square::H2, Square::G1, PromotionShape::Queen).into(),
            ]
            .into()
        );
    }

    #[test]
    fn black_e4_captures_d4_en_passant() {
        let board = test_board!(Black, [
            White Pawn on D4,
            Black Pawn on E4
        ], D3);

        let generated_moves = PawnMoveGenerator::new(&board, None);

        assert_move_list!(generated_moves, [ply!(E4 - E3), ply!(E4 x D3 e.p.),]);
    }

    #[test]
    fn white_e5_captures_f5_en_passant() {
        let board = test_board!(White, [
            White Pawn on E5,
            Black Pawn on F5 
        ], F6);

        let generated_moves = PawnMoveGenerator::new(&board, None);

        assert_move_list!(generated_moves, [ply!(E5 - E6), ply!(E5 x F6 e.p.)]);
    }

    #[test]
    fn white_no_en_passant_if_no_pawn() {
        let board = test_board!(White, [
            White Pawn on A3,
            Black Pawn on F5,
        ], F6);

        let generated_moves: HashSet<_> = PawnMoveGenerator::new(&board, None).collect();

        assert_move_list_does_not_contain!(
            generated_moves,
            [ply!(E5 x F6 e.p.), ply!(G5 x F6 e.p.)]
        );
    }

    #[test]
    fn black_no_en_passant_if_no_pawn() {
        let board = test_board!(Black, [
            White Pawn on A4,
            Black Pawn on D4,
        ], A3);

        let generated_moves: HashSet<_> = PawnMoveGenerator::new(&board, None).collect();

        assert_move_list_does_not_contain!(generated_moves, [ply!(B4 x A3 e.p.)]);
    }

    #[test]
    fn black_en_passant_check() {
        let board = Board::from_fen_str("8/8/8/2k5/2pP4/8/B7/4K3 b - d3 0 3").expect("invalid fen");
        println!("{}", board.display());

        let generated_moves: HashSet<_> = PawnMoveGenerator::new(&board, None).collect();
        assert_move_list_contains!(generated_moves, [ply!(C4 x D3 e.p.)]);
    }
}