// Eryn Wells <eryn@erynwells.me>

use super::BitScanner;
use crate::Square;
use std::fmt;
use std::ops::{BitAnd, BitOr, Not};

#[derive(Clone, Copy, Eq, Hash, PartialEq)]
pub(crate) struct BitBoard(u64);

impl BitBoard {
    pub fn empty() -> BitBoard {
        BitBoard(0)
    }

    pub fn from_bit_field(bits: u64) -> BitBoard {
        BitBoard(bits)
    }

    pub fn rank(rank: u8) -> BitBoard {
        BitBoard(0xFF).shift_north(rank)
    }

    pub fn is_empty(&self) -> bool {
        self.0 == 0
    }

    pub fn has_piece_at(&self, sq: &Square) -> bool {
        (self.0 & (1 << sq.index)) > 0
    }

    pub fn place_piece_at(&mut self, sq: &Square) {
        self.0 |= 1 << sq.index
    }

    fn remove_piece_at(&mut self, sq: &Square) {
        self.0 &= !(1 << sq.index)
    }
}

impl BitBoard {
    const NOT_A_FILE: u64 = 0xfefefefefefefefe;
    const NOT_H_FILE: u64 = 0x7f7f7f7f7f7f7f7f;

    #[inline]
    pub fn shift_north(&self, n: u8) -> BitBoard {
        BitBoard(self.0 << (8 * n))
    }

    #[inline]
    pub fn shift_north_one(&self) -> BitBoard {
        BitBoard(self.0 << 8)
    }

    #[inline]
    pub fn shift_north_east_one(&self) -> BitBoard {
        BitBoard(self.0 << 9 & BitBoard::NOT_A_FILE)
    }

    #[inline]
    pub fn shift_east_one(&self) -> BitBoard {
        BitBoard(self.0 << 1 & BitBoard::NOT_A_FILE)
    }

    #[inline]
    pub fn shift_south_east_one(&self) -> BitBoard {
        BitBoard(self.0 >> 7 & BitBoard::NOT_A_FILE)
    }

    #[inline]
    pub fn shift_south(&self, n: u8) -> BitBoard {
        BitBoard(self.0 >> (8 * n))
    }

    #[inline]
    pub fn shift_south_one(&self) -> BitBoard {
        BitBoard(self.0 >> 8)
    }

    #[inline]
    pub fn shift_south_west_one(&self) -> BitBoard {
        BitBoard(self.0 >> 9 & BitBoard::NOT_H_FILE)
    }

    #[inline]
    pub fn shift_west_one(&self) -> BitBoard {
        BitBoard(self.0 >> 1 & BitBoard::NOT_H_FILE)
    }

    #[inline]
    pub fn shift_north_west_one(&self) -> BitBoard {
        BitBoard(self.0 << 7 & BitBoard::NOT_H_FILE)
    }
}

impl BitBoard {
    pub(crate) fn occupied_squares(&self) -> impl Iterator<Item = Square> {
        BitScanner::new(self.0)
            .map(|x| u8::try_from(x))
            .take_while(|x| x.is_ok())
            .map(|x| Square::from_index_unsafe(x.unwrap()))
    }
}

impl fmt::Debug for BitBoard {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        let bits = self.0;
        write!(f, "BitBoard({bits:064b})")
    }
}

macro_rules! infix_op {
    ($trait_type:ident, $func_name:ident, $left_type:ty, $right_type:ty, $op:tt) => {
        impl $trait_type<$right_type> for $left_type {
            type Output = BitBoard;

            #[inline]
            fn $func_name(self, rhs: $right_type) -> Self::Output {
                BitBoard(self.0 $op rhs.0)
            }
        }
    };
}

infix_op!(BitAnd, bitand, BitBoard, BitBoard, &);
infix_op!(BitAnd, bitand, &BitBoard, BitBoard, &);
infix_op!(BitAnd, bitand, BitBoard, &BitBoard, &);

infix_op!(BitOr, bitor, BitBoard, BitBoard, |);
infix_op!(BitOr, bitor, &BitBoard, BitBoard, |);
infix_op!(BitOr, bitor, BitBoard, &BitBoard, |);

impl Not for BitBoard {
    type Output = BitBoard;

    #[inline]
    fn not(self) -> Self::Output {
        BitBoard(!self.0)
    }
}

impl Not for &BitBoard {
    type Output = BitBoard;

    #[inline]
    fn not(self) -> Self::Output {
        BitBoard(!self.0)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Square;

    #[test]
    fn rank() {
        assert_eq!(BitBoard::rank(0).0, 0xFF, "Rank 1");
        assert_eq!(BitBoard::rank(1).0, 0xFF00, "Rank 2");
        assert_eq!(BitBoard::rank(2).0, 0xFF0000, "Rank 3");
        assert_eq!(BitBoard::rank(3).0, 0xFF000000, "Rank 4");
        assert_eq!(BitBoard::rank(4).0, 0xFF00000000, "Rank 5");
        assert_eq!(BitBoard::rank(5).0, 0xFF0000000000, "Rank 6");
        assert_eq!(BitBoard::rank(6).0, 0xFF000000000000, "Rank 7");
        assert_eq!(BitBoard::rank(7).0, 0xFF00000000000000, "Rank 8");
    }

    #[test]
    fn is_empty() {
        assert!(BitBoard(0).is_empty());
        assert!(!BitBoard(0xFF).is_empty());
    }

    #[test]
    fn has_piece_at() {
        let bb = BitBoard(0b1001100);

        let c1 = Square::from_algebraic_str("c1").expect("Unable to get square for 'c1'");
        assert!(bb.has_piece_at(&c1));

        let b1 = Square::from_algebraic_str("b1").expect("Unable to get square for 'b1'");
        assert!(!bb.has_piece_at(&b1));
    }

    #[test]
    fn place_piece_at() {
        let sq = Square::from_index(34).expect("Invalid square index");

        let mut bb = BitBoard(0b1001100);
        bb.place_piece_at(&sq);
        assert!(bb.has_piece_at(&sq));
    }

    #[test]
    fn remove_piece_at() {
        let sq = Square::from_index(2).expect("Invalid square index");

        let mut bb = BitBoard(0b1001100);
        bb.remove_piece_at(&sq);
        assert!(!bb.has_piece_at(&sq));
    }

    #[test]
    fn cardinal_direction_shifts() {
        let bb = BitBoard(0x1000_0000); // e4

        assert_eq!(bb.shift_north_one(), BitBoard(0x0010_0000_0000), "North");
        assert_eq!(bb.shift_east_one(), BitBoard(0x2000_0000), "East");
        assert_eq!(bb.shift_south_one(), BitBoard(0x0010_0000), "South");
        assert_eq!(bb.shift_west_one(), BitBoard(0x0800_0000), "West");
    }

    #[test]
    fn intercardinal_direction_shifts() {
        let bb = BitBoard(0x1000_0000); // e4

        assert_eq!(
            bb.shift_north_east_one(),
            BitBoard(0x0020_0000_0000),
            "North East"
        );
        assert_eq!(
            bb.shift_south_east_one(),
            BitBoard(0x0020_0000),
            "South East"
        );
        assert_eq!(
            bb.shift_south_west_one(),
            BitBoard(0x0008_0000),
            "South West"
        );
        assert_eq!(
            bb.shift_north_west_one(),
            BitBoard(0x0008_0000_0000),
            "North West"
        );
    }

    #[test]
    fn shift_n() {
        assert_eq!(
            BitBoard(0x0008_0000_0000).shift_north(2),
            BitBoard(0x0008_0000_0000_0000)
        );

        assert_eq!(
            BitBoard(0x0008_0000_0000).shift_south(2),
            BitBoard(0x0008_0000)
        );
    }

    #[test]
    fn pieces() {
        let bb = BitBoard(0x1001_1010); // e4

        let mut occupied_squares = bb.occupied_squares();
        assert_eq!(occupied_squares.next(), Some(Square::from_index_unsafe(28)));
        assert_eq!(occupied_squares.next(), Some(Square::from_index_unsafe(16)));
        assert_eq!(occupied_squares.next(), Some(Square::from_index_unsafe(12)));
        assert_eq!(occupied_squares.next(), Some(Square::from_index_unsafe(4)));
        assert_eq!(occupied_squares.next(), None);
    }
}