// 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 crate::bitboard::BitBoard;
use crate::piece::{Color, Piece};
use crate::{Move, Position};

enum MoveList {
    Pushes = 0,
    Captures = 1,
}

pub(super) struct KingMoveGenerator<'a> {
    position: &'a Position,
    color: Color,

    did_populate_move_lists: bool,

    pushes: BitBoard,
    captures: BitBoard,

    move_lists: [Vec<Move>; 2],
}

impl<'a> KingMoveGenerator<'a> {
    pub(super) fn new(position: &Position, color: Color) -> KingMoveGenerator {
        let mut generator = KingMoveGenerator {
            position,
            color,
            did_populate_move_lists: false,
            pushes: BitBoard::empty(),
            captures: BitBoard::empty(),
            move_lists: [Vec::new(), Vec::new()],
        };

        generator.generate_moves();

        generator
    }

    pub(super) fn iter(&self) -> impl Iterator<Item = &Move> + '_ {
        self.move_lists.iter().flatten()
    }

    fn bitboard(&self) -> BitBoard {
        self.quiet_moves_bitboard() | self.capture_moves_bitboard()
    }

    fn quiet_moves_iter(&self) -> impl Iterator<Item = &Move> + '_ {
        self.pushes_move_list().iter()
    }

    fn quiet_moves_bitboard(&self) -> BitBoard {
        self.pushes
    }

    fn capture_moves_bitboard(&self) -> BitBoard {
        self.captures
    }

    fn capture_moves_iter(&self) -> impl Iterator<Item = &Move> + '_ {
        self.captures_move_list().iter()
    }

    fn generate_moves(&mut self) {
        self.generate_bitboards();
        self.populate_move_lists();
        self.did_populate_move_lists = true;
    }

    fn pushes_move_list(&self) -> &Vec<Move> {
        &self.move_lists[MoveList::Pushes as usize]
    }

    fn pushes_move_list_mut(&mut self) -> &mut Vec<Move> {
        &mut self.move_lists[MoveList::Pushes as usize]
    }

    fn captures_move_list(&self) -> &Vec<Move> {
        &self.move_lists[MoveList::Captures as usize]
    }

    fn captures_move_list_mut(&mut self) -> &mut Vec<Move> {
        &mut self.move_lists[MoveList::Captures as usize]
    }

    fn generate_bitboards(&mut self) {
        let king = self
            .position
            .bitboard_for_piece(Piece::king(self.color))
            .occupied_squares()
            .next()
            .unwrap();
        let all_moves = BitBoard::king_moves(king);

        let empty_squares = self.position.empty_squares();
        let opposing_pieces = self.position.bitboard_for_color(self.color.other());

        self.pushes = all_moves & empty_squares;
        self.captures = all_moves & opposing_pieces;

        // TODO: Handle checks. Prevent moving a king to a square attacked by a
        // piece of the opposite color.
    }

    fn populate_move_lists(&mut self) {
        let king = Piece::king(self.color);

        let from_sq = self
            .position
            .bitboard_for_piece(king)
            .occupied_squares()
            .next()
            .unwrap();

        for to_sq in self.pushes.occupied_squares() {
            self.pushes_move_list_mut()
                .push(Move::new(king, from_sq, to_sq));
        }

        for to_sq in self.captures.occupied_squares() {
            let captured_piece = self.position.piece_on_square(to_sq).unwrap();
            self.captures_move_list_mut()
                .push(Move::new(king, from_sq, to_sq).capturing(captured_piece));
        }

        #[allow(unused_variables)]
        if let Some(castle) = self.king_side_castle() {
            // TODO: Add castle to the pushes move list.
        }

        #[allow(unused_variables)]
        if let Some(castle) = self.queen_side_castle() {
            // TODO: Add castle to the pushes move list.
        }
    }

    fn king_side_castle(&self) -> Option<Move> {
        None
    }

    fn queen_side_castle(&self) -> Option<Move> {
        None
    }
}

impl<'pos> Iterator for KingMoveGenerator<'pos> {
    type Item = Move;

    fn next(&mut self) -> Option<Self::Item> {
        if !self.did_populate_move_lists {
            self.generate_moves();
        }

        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{Position, Square};
    use std::collections::HashSet;

    #[test]
    fn one_king() {
        let mut pos = Position::empty();
        pos.place_piece(&Piece::king(Color::White), &Square::e4())
            .expect("Failed to place king on e4");

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

        let bb = generator.bitboard();
        assert_eq!(
            bb,
            BitBoard::new(
                0b00000000_00000000_00000000_00111000_00101000_00111000_00000000_00000000
            )
        );

        let expected_moves = [
            Move::new(Piece::king(Color::White), Square::e4(), Square::d5()),
            Move::new(Piece::king(Color::White), Square::e4(), Square::e5()),
            Move::new(Piece::king(Color::White), Square::e4(), Square::f5()),
            Move::new(Piece::king(Color::White), Square::e4(), Square::f4()),
            Move::new(Piece::king(Color::White), Square::e4(), Square::f3()),
            Move::new(Piece::king(Color::White), Square::e4(), Square::e3()),
            Move::new(Piece::king(Color::White), Square::e4(), Square::d3()),
            Move::new(Piece::king(Color::White), Square::e4(), Square::d4()),
        ];

        let mut generated_moves: HashSet<Move> = generator.iter().cloned().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 one_king_corner() {
        let mut pos = Position::empty();
        pos.place_piece(&Piece::king(Color::White), &Square::a1())
            .expect("Failed to place king on a1");

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

        let bb = generator.bitboard();
        assert_eq!(
            bb,
            BitBoard::new(
                0b00000000_00000000_00000000_00000000_00000000_00000000_00000011_00000010
            )
        );

        let expected_moves = [
            Move::new(Piece::king(Color::White), Square::a1(), Square::a2()),
            Move::new(Piece::king(Color::White), Square::a1(), Square::b1()),
            Move::new(Piece::king(Color::White), Square::a1(), Square::b2()),
        ];

        let mut generated_moves: HashSet<Move> = generator.iter().cloned().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
        );
    }
}